Merge pull request #542 from shkodm/arb_cut

Implement cuts for the arbitrary lattice

Author: llaniewski

src/ArbConnectivity.hpp

@@ -4,6 +4,7 @@
 #include <memory>
 #include <numeric>
 #include <vector>
+#include "types.h"
 
 struct ArbLatticeConnectivity {
     using Index = long;
@@ -15,18 +16,22 @@ 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{};
+    bool has_cuts{};
 
     ArbLatticeConnectivity() = default;
-    ArbLatticeConnectivity(size_t chunk_begin_, size_t chunk_end_, size_t num_nodes_global_, size_t Q_)
+    ArbLatticeConnectivity(size_t chunk_begin_, size_t chunk_end_, size_t num_nodes_global_, size_t Q_, bool has_cuts_ = false)
         : chunk_begin(chunk_begin_),
           chunk_end(chunk_end_),
           num_nodes_global(num_nodes_global_),
           Q(Q_),
           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(has_cuts_ ? std::make_unique<cut_t[]>(26 * (chunk_end_ - chunk_begin_)) : nullptr),
+          has_cuts(has_cuts_) {
         zones.reserve(getLocalSize());
     }
 
@@ -52,6 +57,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

@@ -160,14 +160,35 @@ void ArbLattice::readFromCxn(const std::string& cxn_path) {
     });
     for (size_t i = 0; i != labels.size(); ++i) label_to_ind_map.emplace(labels[i], i);
 
-    // Nodes header
-    process_section("NODES", [&](size_t num_nodes_global) {
+    // Optional CUTS section - check at the next word to see if it's CUTS or NODES
+    bool has_cuts = false;
+    {
+        file >> word;
+        check_file_ok("Failed to read section header: expected CUTS or NODES");
+        if (word == "CUTS") {
+            size_t n_cuts{};
+            file >> n_cuts;
+            check_file_ok("Failed to read CUTS size");
+            if (n_cuts != 26) throw std::logic_error(wrap_err_msg("Expected CUTS 26, got CUTS " + std::to_string(n_cuts)));
+            has_cuts = true;
+            file >> word;
+            check_file_ok("Failed to read section header: NODES");
+        }
+        check_expected_word("NODES", word);
+    }
+
+    // Nodes
+    {
+        size_t num_nodes_global{};
+        file >> num_nodes_global;
+        check_file_ok("Failed to read section size: NODES");
+
         // Compute the current rank's offset and number of nodes to read
         const auto chunk_offsets = computeInitialNodeDist(num_nodes_global, static_cast<size_t>(comm_size));
         const auto chunk_begin = static_cast<size_t>(chunk_offsets[comm_rank]), chunk_end = static_cast<size_t>(chunk_offsets[comm_rank + 1]);
         const auto num_nodes_local = chunk_end - chunk_begin;
 
-        connect = ArbLatticeConnectivity(chunk_begin, chunk_end, num_nodes_global, Q);
+        connect = ArbLatticeConnectivity(chunk_begin, chunk_end, num_nodes_global, Q, has_cuts);
         connect.grid_size = grid_size;
 
         // Skip chunk_begin + 1 (header) newlines
@@ -193,9 +214,16 @@ void ArbLattice::readFromCxn(const std::string& cxn_path) {
                 file >> zone;
             }
 
+            if (has_cuts) {
+                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");
         }
-    });
+    }
 }
 
 void ArbLattice::partition() {
@@ -293,6 +321,12 @@ 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;
+    if (connect.has_cuts) {
+        cut_distances_device = cudaMakeUnique2D<cut_t>(sizes.cuts_pitch, 26);
+    } else {
+        cut_distances_device.reset();
+    }
     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 +392,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 +438,10 @@ 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);
+    if (connect.has_cuts) {
+        const auto cuts = computeCutDistances();
+        copyVecToDeviceAsync(cut_distances_device.get(), cuts, inStream);
+    }
     CudaStreamSynchronize(inStream);
 }
 
@@ -404,6 +453,8 @@ void ArbLattice::initContainer() {
 #endif
     launcher.container.nbrs = neighbors_device.get();
     launcher.container.coords = coords_device.get();
+    launcher.container.Q = connect.has_cuts ? cut_distances_device.get() : nullptr;
+    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 +894,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

@@ -57,13 +57,14 @@ static auto makeArbLatticeIndexMap(const lbRegion& region, const std::vector<boo
     return retval;
 }
 
-static int writeArbLatticeHeader(std::fstream& file, size_t n_nodes, double grid_size, const Model& model, const std::map<std::string, int>& zone_map) {
+static int writeArbLatticeHeader(std::fstream& file, size_t n_nodes, double grid_size, const Model& model, const std::map<std::string, int>& zone_map, bool has_cuts) {
     file << "OFFSET_DIRECTIONS " << Model_m::offset_directions.size() << '\n';
     for (const auto [x, y, z] : Model_m::offset_directions) file << x << ' ' << y << ' ' << z << '\n';
     file << "GRID_SIZE " << grid_size << '\n';
     file << "NODE_LABELS " << model.nodetypeflags.size() + zone_map.size() << '\n';
     for (const auto& ntf : model.nodetypeflags) file << ntf.name << '\n';
     for (const auto& [name, zf] : zone_map) file << "_Z_" << name << '\n';
+    if (has_cuts) file << "CUTS 26\n";
     file << "NODES " << n_nodes << '\n';
     return file.good() ? EXIT_SUCCESS : EXIT_FAILURE;
 }
@@ -74,7 +75,8 @@ static int writeArbLatticeNodes(const Geometry& geo,
                                 const std::unordered_map<long, long>& lin_to_arb_index_map,
                                 const std::vector<bool>& bulk_bmp,
                                 std::fstream& file,
-                                double spacing) {
+                                double spacing,
+                                bool has_cuts) {
     const long nx = geo.totalregion.nx, ny = geo.totalregion.ny, nz = geo.totalregion.nz;
     const auto get_nbr_id = [&](long my_pos, long nbr_pos) -> long {
         if (my_pos != nbr_pos && bulk_bmp[my_pos] && bulk_bmp[nbr_pos]) return -1;  // ignore edges between bulk nodes
@@ -120,6 +122,15 @@ static int writeArbLatticeNodes(const Geometry& geo,
                         if (zone_flag == zf) file << gz_ind << ' ';
                         ++gz_ind;
                     }
+
+                if (has_cuts) {
+                    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 << ' ';
+                    }
+                }
                 file << '\n';
                 if (!file.good()) break;  // Fail early
             }
@@ -135,13 +146,14 @@ static int writeArbLattice(const Geometry& geo,
                            const std::vector<bool>& bulk_bmp,
                            const std::string& filename,
                            double spacing) {
+    const bool has_cuts = geo.Q != nullptr;
     std::fstream file(filename, std::ios_base::out);
     if (!file.good()) {
         ERROR("Failed to open .cxn file for writing");
         return EXIT_FAILURE;
     }
-    if (writeArbLatticeHeader(file, lin_to_arb_index_map.size(), spacing, model, zone_map)) return EXIT_FAILURE;
-    return writeArbLatticeNodes(geo, model, zone_map, lin_to_arb_index_map, bulk_bmp, file, spacing);
+    if (writeArbLatticeHeader(file, lin_to_arb_index_map.size(), spacing, model, zone_map, has_cuts)) return EXIT_FAILURE;
+    return writeArbLatticeNodes(geo, model, zone_map, lin_to_arb_index_map, bulk_bmp, file, spacing, has_cuts);
 }
 
 static int writeArbXml(const Solver& solver, const Geometry& geo, const Model& model, const std::string& cxn_path) {