Port interpolated bounce back

Author: shkodm

src/ArbConnectivity.hpp

@@ -4,6 +4,7 @@
 #include <memory>
 #include <numeric>
 #include <vector>
+#include "types.h"
 
 struct ArbLatticeConnectivity {
     using Index = long;
@@ -15,6 +16,7 @@ struct ArbLatticeConnectivity {
     std::unique_ptr<Index[]> nbrs;
     std::unique_ptr<ZoneIndex[]> zones_per_node;
     std::vector<ZoneIndex> zones;
+    std::unique_ptr<cut_t[]> cuts;
     double grid_size{};
 
     ArbLatticeConnectivity() = default;
@@ -26,7 +28,8 @@ struct ArbLatticeConnectivity {
           coords(std::make_unique<double[]>(3 * (chunk_end_ - chunk_begin_))),
           og_index(std::make_unique<Index[]>(chunk_end_ - chunk_begin_)),
           nbrs(std::make_unique<Index[]>((chunk_end_ - chunk_begin_) * Q)),
-          zones_per_node(std::make_unique<ZoneIndex[]>(chunk_end_ - chunk_begin_)) {
+          zones_per_node(std::make_unique<ZoneIndex[]>(chunk_end_ - chunk_begin_)),
+          cuts(std::make_unique<cut_t[]>(26 * (chunk_end_ - chunk_begin_))) {
         zones.reserve(getLocalSize());
     }
 
@@ -52,6 +55,8 @@ struct ArbLatticeConnectivity {
     double coord(size_t dim, size_t local_node_ind) const { return coords[local_node_ind + dim * getLocalSize()]; }
     Index& neighbor(size_t q, size_t local_node_ind) { return nbrs[local_node_ind + q * getLocalSize()]; }
     Index neighbor(size_t q, size_t local_node_ind) const { return nbrs[local_node_ind + q * getLocalSize()]; }
+    cut_t& cut_distance(size_t d, size_t local_node_ind) { return cuts[local_node_ind + d * getLocalSize()]; }
+    cut_t cut_distance(size_t d, size_t local_node_ind) const { return cuts[local_node_ind + d * getLocalSize()]; }
 };
 
 inline auto computeInitialNodeDist(size_t num_nodes_global, size_t comm_size) -> std::vector<long> {

src/ArbLattice.cpp

@@ -192,6 +192,11 @@ void ArbLattice::readFromCxn(const std::string& cxn_path) {
                 auto& zone = connect.zones.emplace_back();
                 file >> zone;
             }
+             
+            for (size_t d = 0; d != 26; ++d) {
+                auto& cut = connect.cut_distance(d, local_node_ind);
+                file >> cut;
+            }
 
             check_file_ok("Failed to read node data");
         }
@@ -293,6 +298,8 @@ void ArbLattice::allocDeviceMemory() {
     neighbors_device = cudaMakeUnique2D<unsigned>(sizes.neighbors_pitch, Q);
     sizes.coords_pitch = local_sz;
     coords_device = cudaMakeUnique2D<real_t>(sizes.coords_pitch, 3);
+    sizes.cuts_pitch = local_sz;
+    cut_distances_device = cudaMakeUnique2D<cut_t>(sizes.cuts_pitch, 26);
     sizes.snaps_pitch = local_sz + ghost_nodes.size() + 1;
     snaps_device = cudaMakeUnique2D<storage_t>(sizes.snaps_pitch, sizes.snaps * NF);
     node_types_device = cudaMakeUnique<flag_t>(local_sz);
@@ -358,6 +365,17 @@ std::vector<real_t> ArbLattice::computeCoords() const {
     return retval;
 }
 
+std::vector<cut_t> ArbLattice::computeCutDistances() const {
+    const auto local_sz = connect.getLocalSize();
+    std::vector<cut_t> retval(sizes.cuts_pitch * 26);
+    for (size_t d = 0; d != 26; ++d) {
+        size_t i = 0;
+        for (; i != local_sz; ++i) retval[local_permutation[i] + d * sizes.cuts_pitch] = connect.cut_distance(d, i);
+        for (; i != sizes.cuts_pitch; ++i) retval[i + d * sizes.cuts_pitch] = NO_CUT;  // padding
+    }
+    return retval;
+}
+
 unsigned int ArbLattice::lookupLocalGhostIndex(ArbLatticeConnectivity::Index gid) const {
     const unsigned local_sz = connect.getLocalSize();
     const auto it = std::lower_bound(ghost_nodes.begin(), ghost_nodes.end(), gid);
@@ -393,6 +411,8 @@ void ArbLattice::initDeviceData(pugi::xml_node arb_node, const std::map<std::str
     copyVecToDeviceAsync(neighbors_device.get(), nbrs, inStream);
     const auto coords = computeCoords();
     copyVecToDeviceAsync(coords_device.get(), coords, inStream);
+    const auto cuts = computeCutDistances();
+    copyVecToDeviceAsync(cut_distances_device.get(), cuts, inStream);
     CudaStreamSynchronize(inStream);
 }
 
@@ -404,6 +424,8 @@ void ArbLattice::initContainer() {
 #endif
     launcher.container.nbrs = neighbors_device.get();
     launcher.container.coords = coords_device.get();
+    launcher.container.Q = cut_distances_device.get();
+    launcher.container.cuts_pitch = sizes.cuts_pitch;
     launcher.container.node_types = node_types_device.get();
     launcher.container.nbrs_pitch = sizes.neighbors_pitch;
     launcher.container.coords_pitch = sizes.coords_pitch;
@@ -843,4 +865,4 @@ void ArbLattice::resetAverage(){
             CudaMemset(&getSnapPtr(Snap)[f.id*sizes.snaps_pitch], 0, sizes.snaps_pitch*sizeof(real_t));
         }
     }
-}
+}

src/ArbLattice.hpp

@@ -33,6 +33,7 @@ class ArbLattice : public LatticeBase {
         size_t neighbors_pitch;  /// B + I + padding
         size_t coords_pitch;     /// B + I + padding (should be the same as neighbors_pitch, but let's be extra safe since they come from separate pitched allocation calls)
         size_t snaps_pitch;      /// B + I + G + 1 + padding
+        size_t cuts_pitch;     /// B + I + padding (should be the same as neighbors_pitch, but same as above)
     };
 
     struct CommManager {
@@ -63,6 +64,7 @@ class ArbLattice : public LatticeBase {
     std::unordered_map<std::string, int> label_to_ind_map;  /// Label string to unique ID
     CudaUniquePtr<unsigned> neighbors_device;               /// Device allocation of the neighbor table: (B + I) x Q
     CudaUniquePtr<real_t> coords_device;                    /// Device allocation of node coordinates: (B + I) x 3
+    CudaUniquePtr<cut_t> cut_distances_device;              /// Device allocation of cut-distances: (B + I) x 26
     CudaUniquePtr<storage_t> snaps_device;                  /// Device allocation of snaps: (B + I + G + 1) x NF x num_snaps
     CudaUniquePtr<flag_t> node_types_device;                /// Device allocation of node type array: (B + I)
     std::vector<flag_t, pinned_allocator<flag_t> > node_types_host;               /// Host (pinned) allocation of node type array: (B + I)
@@ -145,6 +147,7 @@ class ArbLattice : public LatticeBase {
     void computeNodeTypesOnHost(pugi::xml_node arb_node, const std::map<std::string, int>& setting_zones, bool permute);           /// Compute the node types to be stored on the device, `permute` enables better code reuse
     std::vector<real_t> computeCoords() const;                                                                                /// Compute the coordinates 2D array to be stored on the device
     std::vector<unsigned> computeNeighbors() const;                                                                           /// Compute the neighbors 2D array to be stored on the device
+    std::vector<cut_t> computeCutDistances() const;                                                                                /// Compute the cut-distances 2D array to be stored on the device
     void initDeviceData(pugi::xml_node arb_node, const std::map<std::string, int>& setting_zones);                                 /// Initialize data residing in device memory
     void initCommManager();                                                                                                        /// Compute which fields need to be sent to/received from which neighbors
     void initContainer();                                                                                                          /// Initialize the data residing in launcher.container

src/ArbLatticeAccess.hpp.Rt

@@ -86,10 +86,9 @@ class ArbLatticeAccess {
     CudaDeviceFunction real_t getY() const { return getDim(1); }
     CudaDeviceFunction real_t getZ() const { return getDim(2); }
     CudaDeviceFunction flag_t getNodeType() const { return node_type; }
-    CudaDeviceFunction cut_t getQ(int) const { /// TODO
-        printf("Cuts not implemented for arbitrary lattice");
-        assert(false);
-        return NO_CUT;
+    CudaDeviceFunction cut_t getQ(int d) const {
+        if (container->Q == nullptr) return NO_CUT;
+        return container->Q[container->cuts_pitch*d + lid];
     }
 <?R
 for (f in rows(Fields)) { ?>

src/ArbLatticeContainer.hpp

@@ -10,14 +10,15 @@
 struct ArbLatticeContainer {
     const unsigned* nbrs;
     const real_t* coords;
+    cut_t* Q;  // cut-distances
     const storage_t* snap_in;
     storage_t* snap_out;
 #ifdef ADJOINT
     const storage_t* adj_snap_in;
     storage_t* adj_snap_out;
 #endif
     const flag_t* node_types;
-    unsigned nbrs_pitch, coords_pitch, snaps_pitch, num_border_nodes, num_interior_nodes;
+    unsigned nbrs_pitch, coords_pitch, cuts_pitch, snaps_pitch, num_border_nodes, num_interior_nodes;
 
     // Packing/unpacking on device
     storage_t* pack_buf;

src/toArb.cpp

@@ -120,6 +120,20 @@ static int writeArbLatticeNodes(const Geometry& geo,
                         if (zone_flag == zf) file << gz_ind << ' ';
                         ++gz_ind;
                     }
+
+                if (geo.Q != nullptr){
+                    size_t regsize = geo.region.sizeL();
+                    for (int d=0; d < 26; d++){
+                        size_t k = geo.region.offset(x, y, z);
+                        const cut_t q = geo.Q[regsize*d + k];
+                        file << q << ' ';
+                    }
+                } else {
+                    for (int i=0; i < 26; i++){
+                        const cut_t q = NO_CUT;
+                        file << q << ' ';
+                    }
+                }
                 file << '\n';
                 if (!file.good()) break;  // Fail early
             }