Support distributed fields

src/pcms/coupler/field_communicator.hpp

@@ -65,6 +65,11 @@ class FieldCommunicator
     auto data = comm_.Recv(redev::Mode::Synchronous);
     serializer_->Deserialize(field_, make_const_array_view(data),
                              layout_comm_.GetPermutationArray());
+    // Deserialize writes owned DOF holders only, so ghost copies still hold
+    // their previous values. Anything that evaluates the field reads ghosts
+    // (an owned element can have vertices another rank owns), so bring them up
+    // to date. A no-op for non-distributed fields.
+    field_.SynchronizeGhosts();
   }
 
 private:

src/pcms/coupler/field_exchange_planner.cpp

@@ -29,20 +29,6 @@ struct OutMsg
   redev::LOs offset;
 };
 
-// Returns the mesh entity dimension for DOF local_index based on the
-// entity offsets array.  ent_offsets[d]..ent_offsets[d+1] is the range of
-// DOF indices belonging to mesh entity dimension d.
-static int GetMeshEntityDim(LO local_index, const EntOffsetsArray& ent_offsets)
-{
-  for (int d = 0; d < ent_offsets_len - 1; ++d) {
-    if (local_index >= static_cast<LO>(ent_offsets[d]) &&
-        local_index < static_cast<LO>(ent_offsets[d + 1])) {
-      return d;
-    }
-  }
-  return ent_offsets_len - 1;
-}
-
 static size_t GetMessageBlockIndex(
   LO permutation_entry, Rank1View<const redev::LO, HostMemorySpace> offsets)
 {
@@ -189,9 +175,10 @@ static ReversePartitionMap2 BuildReversePartitionMap(
 {
   PCMS_FUNCTION_TIMER;
   auto owned = layout.GetOwnedHost();
+  auto owned_to_local = layout.GetOwnedToLocalHost();
   auto class_dims = layout.GetDOFHolderClassificationDimensionsHost();
   auto class_ids = layout.GetDOFHolderClassificationIdsHost();
-  auto coords = layout.GetDOFHolderCoordinates().GetValues();
+  auto coords = layout.GetOwnedDOFHolderCoordinates().GetValues();
   auto ent_offsets = layout.GetEntOffsets();
   int mesh_dim = static_cast<int>(coords.extent(1));
 
@@ -209,19 +196,25 @@ static ReversePartitionMap2 BuildReversePartitionMap(
     Kokkos::create_mirror_view_and_copy(HostMemorySpace(), coords_device);
 
   ReversePartitionMap2 reverse_partition;
-  LO n = static_cast<LO>(owned.extent(0));
+  const LO n_owned = static_cast<LO>(coords.extent(0));
   std::array<Real, 3> coord{};
   auto overlap_mask_view = overlap_mask.GetMask(layout);
 
-  for (LO local_index = 0; local_index < n; ++local_index) {
+  for (LO owned_index = 0; owned_index < n_owned; ++owned_index) {
+    // Classification dims/ids, entity offsets, and the ownership mask are still
+    // local-indexed, so map the owned index back to its local index.
+    const LO local_index =
+      owned_to_local.size() == 0 ? owned_index : owned_to_local(owned_index);
+
+    // Skip holders not owned by this rank. This is required for xgc
     if (!owned(local_index))
       continue;
 
-    if (!overlap_mask_view[local_index])
+    if (!overlap_mask_view[owned_index])
       continue;
 
     for (int d = 0; d < mesh_dim; ++d)
-      coord[d] = coords_host(local_index, d);
+      coord[d] = coords_host(owned_index, d);
     for (int d = mesh_dim; d < 3; ++d)
       coord[d] = 0.0;
 
@@ -230,7 +223,7 @@ static ReversePartitionMap2 BuildReversePartitionMap(
     LO class_id = class_ids[local_index];
 
     auto dr = std::visit(GetRank{class_id, class_dim, coord}, partition);
-    reverse_partition[dr].indices.emplace_back(local_index);
+    reverse_partition[dr].indices.emplace_back(owned_index);
 
     for (size_t e = static_cast<size_t>(mesh_ent_dim) + 1; e < ent_offsets_len;
          ++e) {
@@ -248,7 +241,7 @@ ExchangePlan GenericFieldExchangePlanner::BuildExchangePlan(
 {
   PCMS_FUNCTION_TIMER;
   PCMS_ALWAYS_ASSERT(overlap_mask != nullptr);
-  auto gids = layout.GetGidsHost();
+  auto gids = layout.GetOwnedGidsHost();
 
   const ReversePartitionMap2 reverse_partition =
     BuildReversePartitionMap(layout, partition, *overlap_mask);
@@ -272,8 +265,8 @@ ExchangePlan GenericFieldExchangePlanner::BuildReceivePlan(
   int rank, int nproc, const redev::InMessageLayout& in_message_layout) const
 {
   PCMS_FUNCTION_TIMER;
-  auto gids = layout.GetGidsHost();
-  auto ent_offsets = layout.GetEntOffsets();
+  auto gids = layout.GetOwnedGidsHost();
+  auto ent_offsets = layout.GetOwnedEntOffsets();
 
   ExchangePlan plan;
   auto out_msg = ConstructOutMessage(rank, nproc, in_message_layout);
@@ -295,28 +288,28 @@ void GenericFieldExchangePlanner::FillGidMessage(
   PCMS_ALWAYS_ASSERT(static_cast<size_t>(gid_message.size()) ==
                      plan.msg_size + header_size);
 
-  auto gids = layout.GetGidsHost();
-  auto owned = layout.GetOwnedHost();
+  auto gids = layout.GetOwnedGidsHost();
+  auto owned_to_local = layout.GetOwnedToLocalHost();
   auto ent_offsets = layout.GetEntOffsets();
   auto offsets = Rank1View<const redev::LO, HostMemorySpace>(
     plan.offsets.data(), plan.offsets.size());
 
   std::vector<EntOffsetsArray> per_rank_offsets(plan.dest_ranks.size());
 
-  for (LO local_index = 0; local_index < static_cast<LO>(gids.size());
-       ++local_index) {
-    LO perm_index = plan.permutation[local_index];
+  for (LO owned_index = 0; owned_index < static_cast<LO>(gids.size());
+       ++owned_index) {
+    LO perm_index = plan.permutation[owned_index];
     // Owned holders outside the overlap region carry the sentinel and have no
     // slot in the message.
     if (perm_index < 0)
       continue;
-    // A holder with a valid permutation slot must be owned.
-    PCMS_ALWAYS_ASSERT(owned[local_index]);
     auto block_index = GetMessageBlockIndex(perm_index, offsets);
     const auto gid_index =
       perm_index + static_cast<LO>((block_index + 1) * ent_offsets_len);
-    gid_message(gid_index) = gids(local_index);
+    gid_message(gid_index) = gids(owned_index);
 
+    const LO local_index =
+      owned_to_local.size() == 0 ? owned_index : owned_to_local(owned_index);
     int mesh_ent_dim = GetMeshEntityDim(local_index, ent_offsets);
     for (size_t e = static_cast<size_t>(mesh_ent_dim) + 1; e < ent_offsets_len;
          ++e) {

src/pcms/coupler/field_layout_communicator.cpp

@@ -25,7 +25,7 @@ FieldLayoutCommunicator::FieldLayoutCommunicator(
     planner_(std::move(planner)),
     overlap_mask_(overlap_mask ? std::make_unique<OverlapMask>(*overlap_mask)
                                : std::make_unique<OverlapMask>(
-                                   layout.GetGidsHost().size())),
+                                   layout.GetNumOwnedDofHolder())),
     own_mpi_comm_(own_mpi_comm)
 {
   gid_comm_ = channel.CreateComm<GO>(name_ + "_gids", mpi_comm_);

src/pcms/coupler/field_serializer.h

@@ -18,7 +18,8 @@ class FieldSerializer
                         Rank1View<T, HostMemorySpace> buffer,
                         Rank1View<const LO, HostMemorySpace> permutation) const
   {
-    const auto data = field.GetDOFHolderDataHost().GetValues();
+    // Only owned (rank-exclusive) DOF holders are serialized.
+    const auto data = field.GetOwnedDOFHolderDataHost().GetValues();
     if (buffer.size() > 0) {
       const LO num_dof = static_cast<LO>(data.extent(0));
       const LO num_comp = static_cast<LO>(data.extent(1));
@@ -38,29 +39,30 @@ class FieldSerializer
     Rank1View<const LO, HostMemorySpace> permutation) const
   {
     const auto& layout = field.GetLayout();
-    const LO num_dof = layout.GetNumOwnedDofHolder();
+    const LO num_owned = layout.GetNumOwnedDofHolder();
+    const LO num_local = layout.GetNumLocalDofHolder();
     const LO num_comp = layout.GetNumComponents();
-    // We initialize with local data so that
-    // data that doesn't get received from deserialized
-    // doesn't get zeroed out. Long term FIXME
-    // is to make Get/SetData operate over only part of the
-    // arrays
+    const auto owned_to_local = layout.GetOwnedToLocalHost();
+    // We initialize with local data so that data that doesn't get received
+    // from deserialize doesn't get zeroed out. Long term FIXME is to make
+    // Get/SetData operate over only part of the arrays.
     const auto current = field.GetDOFHolderDataHost().GetValues();
-    Kokkos::View<T*, HostMemorySpace> sorted("sorted", layout.OwnedSize());
-    for (LO i = 0; i < num_dof; ++i) {
+    Kokkos::View<T*, HostMemorySpace> sorted("sorted", layout.LocalSize());
+    for (LO i = 0; i < num_local; ++i) {
       for (LO c = 0; c < num_comp; ++c) {
         sorted[i * num_comp + c] = current(i, c);
       }
     }
-    for (LO i = 0; i < num_dof; ++i) {
-      if (permutation[i] >= 0) {
+    for (LO o = 0; o < num_owned; ++o) {
+      if (permutation[o] >= 0) {
+        const LO local = owned_to_local.size() == 0 ? o : owned_to_local(o);
         for (LO c = 0; c < num_comp; ++c) {
-          sorted[i * num_comp + c] = buffer[permutation[i] * num_comp + c];
+          sorted[local * num_comp + c] = buffer[permutation[o] * num_comp + c];
         }
       }
     }
     field.SetDOFHolderDataUncheckedHost(
-      Rank2View<const T, HostMemorySpace>(sorted.data(), num_dof, num_comp));
+      Rank2View<const T, HostMemorySpace>(sorted.data(), num_local, num_comp));
   }
 
   virtual ~FieldSerializer() noexcept = default;

src/pcms/coupler/overlap_mask.h

@@ -56,16 +56,28 @@ struct OverlapMask
     }
   }
 
-  // Get the mask, evaluating the function if needed
+  // Get the mask, evaluating the function if needed. The mask is owned-indexed
+  // (one entry per owned DOF holder), matching the owned-indexed exchange plan.
   Rank1View<const bool, HostMemorySpace> GetMask(
     const FieldLayout& layout) const
   {
     if (in_overlap_func_) {
       auto class_dims = layout.GetDOFHolderClassificationDimensionsHost();
       auto class_ids = layout.GetDOFHolderClassificationIdsHost();
-      for (size_t i = 0; i < is_overlap_.extent(0); ++i) {
-        is_overlap_[i] =
-          static_cast<bool>(in_overlap_func_(class_dims[i], class_ids[i]));
+      auto owned_to_local = layout.GetOwnedToLocalHost();
+      const size_t n = is_overlap_.extent(0);
+      if (owned_to_local.size() == 0) {
+        // Non-distributed layout: owned == local.
+        for (size_t i = 0; i < n; ++i) {
+          is_overlap_[i] =
+            static_cast<bool>(in_overlap_func_(class_dims[i], class_ids[i]));
+        }
+      } else {
+        for (size_t o = 0; o < n; ++o) {
+          const LO local = owned_to_local(o);
+          is_overlap_[o] = static_cast<bool>(
+            in_overlap_func_(class_dims[local], class_ids[local]));
+        }
       }
     }
     return make_const_array_view(is_overlap_);

src/pcms/coupler/serializer/xgc.h

@@ -87,7 +87,7 @@ class XGCFieldSerializer : public FieldSerializer<T>
               pcms::GetMPIType(T{}), 0, plane_comm_);
 
     xgc_field->SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace>(
-      full_data.data(), layout.GetNumOwnedDofHolder(),
+      full_data.data(), layout.GetNumLocalDofHolder(),
       layout.GetNumComponents()));
   }
 

src/pcms/field/CMakeLists.txt

@@ -53,6 +53,7 @@ if (PCMS_ENABLE_OMEGA_H)
             layout/omega_h_entity.h
             layout/omega_h_lagrange.h
             evaluator/omega_h_lagrange.h
+            evaluator/distributed_point_evaluator.h
             layout/uniform_grid.h
             uniform_grid_binary_field.h
             evaluator/uniform_grid.h

src/pcms/field/data/mesh_fields.h

@@ -7,9 +7,11 @@
 #include "pcms/field/value_view.hpp"
 #include "pcms/utility/assert.h"
 #include "pcms/utility/arrays.h"
+#include "pcms/utility/omega_h_array_utils.h"
 
 #include <Kokkos_Core.hpp>
 #include <memory>
+#include <type_traits>
 
 namespace pcms
 {
@@ -38,6 +40,32 @@ struct FlattenFunctor
     }
   }
 };
+
+// Functor that scatters a flat dof-major buffer back into component-major 2D
+// data (LayoutLeft). Inverse of FlattenFunctor; indexes the full destination
+// view directly (rather than via Kokkos::subview) because a partial row-range
+// subview of a LayoutLeft view is not representable as LayoutLeft/LayoutRight
+// (Kokkos falls back to LayoutStride), which Rank2View does not detect.
+template <typename T>
+struct UnflattenFunctor
+{
+  Kokkos::View<const T*, DeviceMemorySpace> flat_;
+  Rank2View<T, DeviceMemorySpace> data_;
+  LO row_off_;
+  size_t nc_;
+  UnflattenFunctor(Kokkos::View<const T*, DeviceMemorySpace> f,
+                   Rank2View<T, DeviceMemorySpace> d, LO ro, size_t nc)
+    : flat_(f), data_(d), row_off_(ro), nc_(nc)
+  {
+  }
+  KOKKOS_INLINE_FUNCTION void operator()(LO local) const
+  {
+    LO global_dof = row_off_ + local;
+    for (size_t c = 0; c < nc_; ++c) {
+      data_(global_dof, c) = flat_(local * nc_ + c);
+    }
+  }
+};
 } // namespace
 
 template <typename T>
@@ -50,10 +78,10 @@ class MeshFieldsFieldData : public FieldData<T>
       basis_(std::move(basis)),
       mesh_field_(MakeMeshFieldBackend<T>(*layout_)),
       host_data_("meshfields_field_data",
-                 static_cast<size_t>(layout_->GetNumOwnedDofHolder()),
+                 static_cast<size_t>(layout_->GetNumLocalDofHolder()),
                  static_cast<size_t>(layout_->GetNumComponents())),
       device_data_("meshfields_field_data_device",
-                   static_cast<size_t>(layout_->GetNumOwnedDofHolder()),
+                   static_cast<size_t>(layout_->GetNumLocalDofHolder()),
                    static_cast<size_t>(layout_->GetNumComponents()))
   {
     if (!mesh_field_) {
@@ -74,6 +102,17 @@ class MeshFieldsFieldData : public FieldData<T>
     return MakeConstRank2View(host_data_);
   }
 
+  Rank2View<const T, HostMemorySpace> GetOwnedDOFHolderDataHost() const override
+  {
+    return GatherOwnedHostData(*layout_, device_data_, host_data_,
+                               owned_host_data_);
+  }
+
+  Rank2View<const T, DeviceMemorySpace> GetOwnedDOFHolderData() const override
+  {
+    return GatherOwnedDeviceData(*layout_, device_data_, owned_device_data_);
+  }
+
   void SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace> values) override
   {
     CopyHostRank2ViewToDeviceView(device_data_, values);
@@ -91,6 +130,50 @@ class MeshFieldsFieldData : public FieldData<T>
     SyncBackend(GetDOFHolderData());
   }
 
+  void SynchronizeGhosts() override
+  {
+    const int nc = layout_->GetNumComponents();
+    auto& mesh = layout_->GetMesh();
+    const auto nodes_per_dim = layout_->GetNodesPerDim();
+
+    size_t row_offset = 0;
+    for (int dim = 0; dim <= mesh.dim(); ++dim) {
+      if (!nodes_per_dim[dim]) {
+        continue;
+      }
+      const LO num_rows = static_cast<LO>(mesh.nents(dim)) * nodes_per_dim[dim];
+      const LO row_off = static_cast<LO>(row_offset);
+
+      // device_data_ may be component-major on device; flatten this block to
+      // dof-major order for Omega_h's sync_array, then scatter the
+      // synchronized values back. We index the full device_data_ view
+      // directly (not a Kokkos::subview of it) because a partial row-range
+      // subview of a LayoutLeft view is not representable as
+      // LayoutLeft/LayoutRight (Kokkos falls back to LayoutStride), which
+      // MakeConstRank2View/MakeRank2View do not detect.
+      Kokkos::View<T*, DeviceMemorySpace> flat("ghost_sync_flat",
+                                               static_cast<size_t>(num_rows) *
+                                                 static_cast<size_t>(nc));
+      Kokkos::parallel_for(
+        "SynchronizeGhostsFlatten",
+        Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, num_rows),
+        FlattenFunctor<T>(MakeConstRank2View(device_data_), flat, row_off,
+                          static_cast<size_t>(nc)));
+
+      SynchronizeOmegaHBlock<T>(mesh, dim, nc, flat);
+
+      Kokkos::parallel_for(
+        "SynchronizeGhostsUnflatten",
+        Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, num_rows),
+        UnflattenFunctor<T>(flat, MakeRank2View(device_data_), row_off,
+                            static_cast<size_t>(nc)));
+
+      row_offset += static_cast<size_t>(num_rows);
+    }
+
+    SyncBackend(GetDOFHolderData());
+  }
+
   std::shared_ptr<MeshFieldBackend<T>> GetMeshFieldBackend() const
   {
     return mesh_field_;
@@ -131,6 +214,8 @@ class MeshFieldsFieldData : public FieldData<T>
   ValueBasis basis_;
   std::shared_ptr<MeshFieldBackend<T>> mesh_field_;
   mutable Kokkos::View<T**, HostMemorySpace> host_data_;
+  mutable Kokkos::View<T**, HostMemorySpace> owned_host_data_;
+  mutable Kokkos::View<T**, DeviceMemorySpace> owned_device_data_;
   Kokkos::View<T**, DeviceMemorySpace> device_data_;
 };