Return Array CSR PIDS (#113)

SCOREC · Feb 2, 2024 · 7b55b1b · 7b55b1b
1 parent 35c49f0
commit 7b55b1b
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 9 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.0.0)
 
-project(pumipic VERSION 2.0.3 LANGUAGES CXX)
+project(pumipic VERSION 2.1.0 LANGUAGES CXX)
 
 include(cmake/bob.cmake)
 

diff --git a/particle_structs/src/particle_structure.hpp b/particle_structs/src/particle_structure.hpp
@@ -95,6 +95,9 @@ namespace pumipic {
                          MTVs new_particle_info = NULL) = 0;
     virtual void printMetrics() const = 0;
     virtual void printFormat(const char* prefix = "") const = 0;
+
+    template <typename ViewT>
+    void getPIDs(ViewT& pids, ViewT& offsets);
   protected:
     //String to identify the particle structure
     std::string name;

diff --git a/particle_structs/src/ps_for.hpp b/particle_structs/src/ps_for.hpp
@@ -53,4 +53,34 @@ namespace pumipic {
     throw 1;
     return NULL;
   }
+
+  /** This function initializes and populates the pids and offsets arrays
+   * @param[out] pids Returns a new array of PIDs sorted by elements
+   * @param[out] offsets Returns a new array of where
+   *                       each index is an element
+   *                       each value is the starting index in the pids array for that element
+  */
+  template <typename DataTypes, typename MemSpace>
+  template <typename ViewT>
+  void ParticleStructure<DataTypes, MemSpace>::getPIDs(ViewT& pids, ViewT& offsets) {
+    offsets = ViewT("offsets", num_elems+1);
+    ViewT ppe("ppe", num_elems+1);
+    auto setPPE = PS_LAMBDA(const lid_t& e, const lid_t& p, const bool& mask) {
+      if (mask) {
+        Kokkos::atomic_increment(&ppe(e));
+      }
+    };
+    parallel_for(this, setPPE, "setPPE");
+    exclusive_scan(ppe, offsets, execution_space());
+
+    pids = ViewT("pids", getLastValue(offsets));
+    ViewT currIndex("currIndex", num_elems);
+    auto setPIDs = PS_LAMBDA(const lid_t& e, const lid_t& p, const bool& mask) {
+      if (mask) {
+        auto index = Kokkos::atomic_fetch_add(&currIndex(e), 1);
+        pids(offsets(e)+index) = p;
+      }
+    };
+    parallel_for(this, setPIDs, "setPIDs");
+  }
 }
diff --git a/particle_structs/test/test_structure.cpp b/particle_structs/test/test_structure.cpp
@@ -35,6 +35,7 @@ int migrateSendToOne(const char* name, PS* structure);
 int testMetrics(const char* name, PS* structure);
 int testCopy(const char* name, PS* structure);
 int testSegmentComp(const char* name, PS* structure);
+int testPIDs(const char* name, PS* structure);
 
 //Edge Case tests
 int migrateToEmptyAndRefill(const char* name, PS* structure);
@@ -85,6 +86,7 @@ int main(int argc, char* argv[]) {
       fails += testCounts(name.c_str(), structure, num_elems, num_ptcls);
       fails += testParticleExistence(name.c_str(), structure, num_ptcls);
       fails += setValues(name.c_str(), structure);
+      fails += testPIDs(name.c_str(), structure);
       fails += pseudoPush(name.c_str(), structure);
       fails += testMetrics(name.c_str(), structure);
       fails += testRebuild(name.c_str(), structure);
@@ -163,14 +165,6 @@ PS* buildNextStructure(int num, lid_t num_elems, lid_t num_ptcls, kkLidView ppe,
                                           element_gids, particle_elements, particle_info);
     }
     else if (num == 5) {
-      //DPS
-      error_message = "DPS Host";
-      name = "dps host";
-      Kokkos::TeamPolicy<Kokkos::DefaultHostExecutionSpace> policy = pumipic::TeamPolicyAuto(num_elems,32);
-      return new ps::DPS<Types, Kokkos::HostSpace>(policy, num_elems, num_ptcls, ppe,
-                                          element_gids, particle_elements, particle_info);
-    }
-    else if (num == 6) {
       //DPS
       error_message = "DPS 2";
       name = "dps 2";
@@ -357,6 +351,30 @@ int testSegmentComp(const char* name, PS* structure) {
   return fails;
 }
 
+int testPIDs(const char* name, PS* structure) {
+  printf("testPIDs %s, rank %d\n", name, comm_rank);
+  kkLidView pids;
+  kkLidView offsets;
+  int fails = 0;
+  structure->getPIDs(pids, offsets);
+
+  kkLidView failures("failures", 1);
+  kkLidView unsortedElems("unsortedElems", structure->capacity());
+  auto setUnsortedElems = PS_LAMBDA(const lid_t& e, const lid_t& p, const bool& mask) {
+    if (mask) unsortedElems(p) = e;
+  };
+  pumipic::parallel_for(structure, setUnsortedElems, "setUnsortedElems");
+
+  Kokkos::parallel_for(pids.size(), KOKKOS_LAMBDA(const lid_t& i) {
+    lid_t pid = pids(i);
+    lid_t oldElem = unsortedElems(pid);
+    if (i < offsets(oldElem) || i >= offsets(oldElem+1))
+      Kokkos::atomic_add(&(failures[0]), 1);
+  });
+  fails += pumipic::getLastValue(failures);
+  return fails;
+}
+
 #include "test_constructor.cpp"
 #include "test_rebuild.cpp"
 #include "test_migrate.cpp"