SCOREC · Sichao25 · Sep 11, 2026
diff --git a/src/pcms/coupler/field_exchange_planner.cpp b/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) {

diff --git a/src/pcms/coupler/field_layout_communicator.cpp b/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_);

diff --git a/src/pcms/coupler/field_serializer.h b/src/pcms/coupler/field_serializer.h
@@ -18,18 +18,12 @@ class FieldSerializer
                         Rank1View<T, HostMemorySpace> buffer,
                         Rank1View<const LO, HostMemorySpace> permutation) const
   {
-    auto data = field.GetDOFHolderDataHost();
-    auto owned = layout.GetOwnedHost();
-    // The exchange plan is per DOF holder: owned[i] and permutation[i] are
-    // indexed by holder. All num_components components of a holder share its
-    // location, so they occupy one contiguous block permutation[i]*num_comp in
-    // the wire buffer.
+    // Only owned (rank-exclusive) DOF holders are serialized.
+    auto data = field.GetOwnedDOFHolderDataHost();
     if (buffer.size() > 0) {
       const LO num_dof = static_cast<LO>(data.extent(0));
       const LO num_comp = static_cast<LO>(data.extent(1));
       for (LO i = 0; i < num_dof; ++i) {
-        // A negative permutation entry marks a holder outside the exchange
-        // (non-owned, or owned but outside the overlap region); it has no slot.
         if (permutation[i] >= 0) {
           for (LO c = 0; c < num_comp; ++c) {
             buffer[permutation[i] * num_comp + c] = data(i, c);
@@ -45,21 +39,21 @@ class FieldSerializer
     Rank1View<const T, HostMemorySpace> buffer,
     Rank1View<const LO, HostMemorySpace> permutation) const
   {
-    const LO num_dof = layout.GetNumOwnedDofHolder();
+    const LO num_owned = layout.GetNumOwnedDofHolder();
+    const LO num_local = layout.GetNumLocalDofHolder();
     const LO num_comp = layout.GetNumComponents();
-    Kokkos::View<T*, HostMemorySpace> sorted("sorted", layout.OwnedSize());
-    for (LO i = 0; i < num_dof; ++i) {
-      // A negative permutation entry marks a holder outside the exchange (owned
-      // but outside the overlap region); no data was received for it, so its
-      // zero-initialized `sorted` slot is left as-is.
-      if (permutation[i] >= 0) {
+    const auto owned_to_local = layout.GetOwnedToLocalHost();
+    Kokkos::View<T*, HostMemorySpace> sorted("sorted", layout.LocalSize());
+    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.SetDOFHolderDataHost(
-      Rank2View<const T, HostMemorySpace>(sorted.data(), num_dof, num_comp));
+      Rank2View<const T, HostMemorySpace>(sorted.data(), num_local, num_comp));
   }
 
   virtual ~FieldSerializer() noexcept = default;

diff --git a/src/pcms/coupler/overlap_mask.h b/src/pcms/coupler/overlap_mask.h
@@ -46,16 +46,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_);

diff --git a/src/pcms/coupler/serializer/xgc.h b/src/pcms/coupler/serializer/xgc.h
@@ -90,7 +90,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()));
   }
 

diff --git a/src/pcms/field/data/mesh_fields.h b/src/pcms/field/data/mesh_fields.h
@@ -7,9 +7,11 @@
 #include "pcms/field/field_metadata.h"
 #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
 {
@@ -24,9 +26,9 @@ class MeshFieldsFieldData : public FieldData<T>
       metadata_(metadata),
       mesh_field_(MakeMeshFieldBackend<T>(*layout_)),
       host_data_("meshfields_field_data",
-                 static_cast<size_t>(layout_->OwnedSize())),
+                 static_cast<size_t>(layout_->LocalSize())),
       device_data_("meshfields_field_data_device",
-                   static_cast<size_t>(layout_->OwnedSize()))
+                   static_cast<size_t>(layout_->LocalSize()))
   {
     if (!mesh_field_) {
       throw pcms_error(
@@ -40,14 +42,25 @@ class MeshFieldsFieldData : public FieldData<T>
   {
     Kokkos::deep_copy(host_data_, device_data_);
     return Rank2View<const T, HostMemorySpace>(host_data_.data(),
-                                               layout_->GetNumOwnedDofHolder(),
+                                               layout_->GetNumLocalDofHolder(),
                                                layout_->GetNumComponents());
   }
 
+  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
   {
     PCMS_ALWAYS_ASSERT(values.size() ==
-                       static_cast<size_t>(layout_->OwnedSize()));
+                       static_cast<size_t>(layout_->LocalSize()));
     CopyHostRank2ViewToDeviceView(device_data_, values);
     SyncBackend(GetDOFHolderData());
   }
@@ -58,18 +71,43 @@ class MeshFieldsFieldData : public FieldData<T>
     // memory is enabled. This may cause issues in multi component cases. See
     // issue #342
     return Rank2View<const T, DeviceMemorySpace>(
-      device_data_.data(), layout_->GetNumOwnedDofHolder(),
+      device_data_.data(), layout_->GetNumLocalDofHolder(),
       layout_->GetNumComponents());
   }
 
   void SetDOFHolderData(Rank2View<const T, DeviceMemorySpace> values) override
   {
     PCMS_ALWAYS_ASSERT(values.size() ==
-                       static_cast<size_t>(layout_->OwnedSize()));
+                       static_cast<size_t>(layout_->LocalSize()));
     CopyDeviceRank2ViewToDeviceView(device_data_, values);
     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 flat_len = num_rows * nc;
+      const LO flat_off = static_cast<LO>(row_offset * static_cast<size_t>(nc));
+
+      auto block = Kokkos::subview(
+        device_data_, Kokkos::make_pair(flat_off, flat_off + flat_len));
+      SynchronizeOmegaHBlock<T>(mesh, dim, nc, block);
+
+      row_offset += static_cast<size_t>(num_rows);
+    }
+
+    SyncBackend(GetDOFHolderData());
+  }
+
   std::shared_ptr<MeshFieldBackend<T>> GetMeshFieldBackend() const
   {
     return mesh_field_;
@@ -103,6 +141,8 @@ class MeshFieldsFieldData : public FieldData<T>
   FieldMetadata metadata_;
   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_;
 };