Partition Wrapper

.github/workflows/cmake-test.yml

@@ -6,7 +6,7 @@ on:
 
 jobs:
   test-build:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     timeout-minutes: 30
 
     defaults:
@@ -31,7 +31,7 @@ jobs:
       run: sudo apt-get update
 
     - name: Install mpi
-      run: sudo apt-get install -yq mpich libmpich-dev
+      run: sudo apt-get install -yq mpich
 
     - uses: actions/checkout@v4
 

.github/workflows/perlmutter/install.sh

@@ -9,7 +9,7 @@ module load cray-fftw
 # # kokkos
 # rm build-kokkos -rf
 # # rm kokkos -rf
-# git clone -b 4.2.00 https://github.com/kokkos/kokkos.git
+# git clone -b 4.5.00 https://github.com/kokkos/kokkos.git
 # cmake -S kokkos -B build-kokkos \
 #   -DCMAKE_INSTALL_PREFIX=build-kokkos/install \
 #   -DCMAKE_BUILD_TYPE="Release" \
@@ -100,7 +100,7 @@ cmake -S pcms -B build-pcms \
   -DCMAKE_C_COMPILER=cc \
   -DCMAKE_CXX_COMPILER=CC \
   -DCMAKE_BUILD_TYPE=Release \
-  -DPCMS_TIMEOUT=0 \
+  -DPCMS_TIMEOUT=100 \
   -Dperfstubs_DIR=$PWD/build-perfstubs \
   -Dredev_DIR=$PWD/build-redev/install/lib64/cmake/redev \
   -DOmega_h_DIR=$PWD/build-omega_h/install/lib64/cmake/Omega_h/ \

CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.19)
-project(pcms VERSION 0.1.1 LANGUAGES CXX)
+project(pcms VERSION 0.1.2 LANGUAGES CXX)
 
 include(GNUInstallDirs)
 include(CMakePackageConfigHelpers)

src/CMakeLists.txt

@@ -17,6 +17,7 @@ set(PCMS_HEADERS
         pcms/inclusive_scan.h
         pcms/profile.h
         pcms/print.h
+        pcms/partition.h
         )
 
 set(PCMS_SOURCES

src/pcms/dummy_field_adapter.h

@@ -8,7 +8,7 @@ class DummyFieldAdapter
   using value_type = int;
   [[nodiscard]] std::vector<GO> GetGids() const { return {}; }
   [[nodiscard]] ReversePartitionMap GetReversePartitionMap(
-    const redev::Partition& partition) const
+    const Partition& partition) const
   {
     return {};
   }

src/pcms/field_communicator.h

@@ -5,6 +5,8 @@
 #include <numeric>
 #include "pcms/inclusive_scan.h"
 #include "pcms/profile.h"
+#include "pcms/partition.h"
+
 namespace pcms
 {
 
@@ -203,7 +205,7 @@ struct FieldCommunicator
       auto gids = field_adapter_.GetGids();
       if (redev_.GetProcessType() == redev::ProcessType::Client) {
         const ReversePartitionMap reverse_partition =
-          field_adapter_.GetReversePartitionMap(redev_.GetPartition());
+          field_adapter_.GetReversePartitionMap(Partition{redev_.GetPartition()});
         auto out_message = ConstructOutMessage(reverse_partition);
         comm_.SetOutMessageLayout(out_message.dest, out_message.offset);
         gid_comm_.SetOutMessageLayout(out_message.dest, out_message.offset);

src/pcms/omega_h_field.h

@@ -17,6 +17,7 @@
 #include "pcms/transfer_field.h"
 #include "pcms/memory_spaces.h"
 #include "pcms/profile.h"
+#include "pcms/partition.h"
 
 // FIXME add executtion spaces (don't use kokkos exe spaces directly)
 
@@ -85,29 +86,6 @@ Omega_h::Read<T> filter_array(Omega_h::Read<T> array,
     });
   return filtered_field;
 }
-struct GetRankOmegaH
-{
-  GetRankOmegaH(int i, Omega_h::I8 dim, Omega_h::ClassId id, std::array<pcms::Real,3> & coord)
-    : i_(i), id_(id), dim_(dim), coord_(coord)
-  {
-    PCMS_FUNCTION_TIMER;
-  }
-  auto operator()(const redev::ClassPtn& ptn) const
-  {
-    PCMS_FUNCTION_TIMER;
-    const auto ent = redev::ClassPtn::ModelEnt({dim_, id_});
-    return ptn.GetRank(ent);
-  }
-  auto operator()(const redev::RCBPtn& ptn)
-  {
-    PCMS_FUNCTION_TIMER;
-    return ptn.GetRank(coord_);
-  }
-  int i_;
-  Omega_h::ClassId id_;
-  Omega_h::I8 dim_;
-  std::array<pcms::Real,3> coord_;
-};
 } // namespace detail
 
 template <typename T,
@@ -557,7 +535,7 @@ class OmegaHFieldAdapter
   }
   // REQUIRED
   [[nodiscard]] ReversePartitionMap GetReversePartitionMap(
-    const redev::Partition& partition) const
+    const Partition& partition) const
   {
     PCMS_FUNCTION_TIMER;
     auto classIds_h = Omega_h::HostRead<Omega_h::ClassId>(field_.GetClassIDs());
@@ -575,8 +553,7 @@ class OmegaHFieldAdapter
       coord[0] = coords[i * dim];
       coord[1] = coords[i * dim + 1];
       coord[2] = (dim == 3) ? coords[i * dim + 2] : 0.0;
-      auto dr = std::visit(detail::GetRankOmegaH{i, classDims_h[i], classIds_h[i], coord},
-                           partition);
+      auto dr = partition.GetDr(classIds_h[i], classDims_h[i], coord);
       reverse_partition[dr].emplace_back(local_index++);
     }
     return reverse_partition;

src/pcms/partition.h

@@ -0,0 +1,47 @@
+#ifndef PCMS_PARTITION_H
+#define PCMS_PARTITION_H
+#include "pcms/common.h"
+
+
+namespace pcms
+{
+  struct GetRank
+  {
+    GetRank(const LO id, const LO dim, const std::array<Real,3>& coord)
+      : id_(id), dim_(dim), coord_(coord)
+    {
+      PCMS_FUNCTION_TIMER;
+    }
+    auto operator()(const redev::ClassPtn& ptn) const
+    {
+      PCMS_FUNCTION_TIMER;
+      const auto ent = redev::ClassPtn::ModelEnt({dim_, id_});
+      return ptn.GetRank(ent);
+    }
+    auto operator()(const redev::RCBPtn& ptn)
+    {
+      PCMS_FUNCTION_TIMER;
+      return ptn.GetRank(coord_);
+    }
+    LO id_;
+    LO dim_;
+    std::array<Real,3> coord_;
+  };
+
+  struct Partition
+  {
+    Partition(const redev::Partition& partition) : partition_(partition)
+    {
+      PCMS_FUNCTION_TIMER;
+    }
+
+    auto GetDr(const LO id, const LO dim, const std::array<Real,3> coord = {}) const
+    {
+      return std::visit(GetRank{id, dim, coord}, partition_);
+    }
+
+    redev::Partition partition_;
+  };
+} // namespace pcms
+
+#endif //PCMS_PARTITION_H

src/pcms/xgc_field_adapter.h

@@ -13,30 +13,6 @@
 
 namespace pcms
 {
-namespace detail
-{
-// Needed since NVHPC doesn't work with overloaded
-struct GetRank
-{
-  using GeomType = DimID;
-  GetRank(const GeomType& geom) : geom_(geom) {}
-  auto operator()(const redev::ClassPtn& ptn) const
-  {
-    PCMS_FUNCTION_TIMER;
-    const auto ent = redev::ClassPtn::ModelEnt({geom_.dim, geom_.id});
-    return ptn.GetRank(ent);
-  }
-  auto operator()(const redev::RCBPtn& /*unused*/) const
-  {
-    PCMS_FUNCTION_TIMER;
-    std::cerr << "RCB partition not handled yet\n";
-    std::terminate();
-    return 0;
-  }
-  const GeomType& geom_;
-};
-} // namespace detail
-
 template <typename T, typename CoordinateElementType = Real>
 class XGCFieldAdapter
 {
@@ -136,7 +112,7 @@ class XGCFieldAdapter
 
   // REQUIRED
   [[nodiscard]] ReversePartitionMap GetReversePartitionMap(
-    const redev::Partition& partition) const
+    const Partition& partition) const
   {
     PCMS_FUNCTION_TIMER;
     if (RankParticipatesCouplingCommunication()) {
@@ -148,7 +124,7 @@ class XGCFieldAdapter
         // if the geometry is in specified overlap region
         if (in_overlap_(geom.first.dim, geom.first.id)) {
 
-          auto dr = std::visit(detail::GetRank{geom.first}, partition);
+          auto dr = partition.GetDr(geom.first.id, geom.first.dim);
           auto [it, inserted] = reverse_partition.try_emplace(dr);
           // the map gives the local iteration order of the global ids
           auto map = mask_.GetMap();

test/test_ohClassPtn_appRibPtn.cpp

@@ -49,8 +49,6 @@ int main(int argc, char** argv)
                    static_cast<redev::ProcessType>(isRdv));
 
   const std::string name = "meshVtxIds";
-  const int rdvRanks = 2;
-  const int appRanks = 2;
 
   adios2::Params params{{"Streaming", "On"}, {"OpenTimeoutSecs", "60"}};
   auto channel =