Make cut array reading / writing conditions

Author: shkodm

src/ArbConnectivity.hpp

@@ -18,9 +18,10 @@ struct ArbLatticeConnectivity {
     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_),
@@ -29,7 +30,8 @@ struct ArbLatticeConnectivity {
           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_)),
-          cuts(std::make_unique<cut_t[]>(26 * (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());
     }
 

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
@@ -192,15 +213,17 @@ 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;
+
+            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() {
@@ -299,7 +322,11 @@ void ArbLattice::allocDeviceMemory() {
     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);
+    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);
@@ -411,8 +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);
-    const auto cuts = computeCutDistances();
-    copyVecToDeviceAsync(cut_distances_device.get(), cuts, inStream);
+    if (connect.has_cuts) {
+        const auto cuts = computeCutDistances();
+        copyVecToDeviceAsync(cut_distances_device.get(), cuts, inStream);
+    }
     CudaStreamSynchronize(inStream);
 }
 
@@ -424,7 +453,7 @@ 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.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;

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
@@ -121,18 +123,13 @@ static int writeArbLatticeNodes(const Geometry& geo,
                         ++gz_ind;
                     }
 
-                if (geo.Q != nullptr){
+                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 << ' ';
                     }
-                } else {
-                    for (int i=0; i < 26; i++){
-                        const cut_t q = NO_CUT;
-                        file << q << ' ';
-                    }
                 }
                 file << '\n';
                 if (!file.good()) break;  // Fail early
@@ -149,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) {