Various fixes

src/pcms/coupler/field_communicator.hpp

@@ -35,7 +35,11 @@ class FieldCommunicator
       serializer_(std::move(serializer))
   {
     PCMS_ALWAYS_ASSERT(&layout_comm.GetLayout() == &field.GetLayout());
-    comm_buffer_.resize(layout_comm.GetMsgSize());
+    // The exchange plan is per DOF holder; the field payload carries
+    // num_components values per holder, so the buffer is scaled accordingly.
+    comm_buffer_.resize(
+      layout_comm.GetMsgSize() *
+      static_cast<size_t>(field.GetLayout().GetNumComponents()));
     comm_ =
       layout_comm_.GetChannel().CreateComm<T>(name, layout_comm.GetMPIComm());
     layout_comm_.SetOutMessageLayout<T>(comm_);

src/pcms/coupler/field_exchange_planner.cpp

@@ -180,7 +180,7 @@ static ReversePartitionMap2 BuildReversePartitionMap(
   auto owned = layout.GetOwnedHost();
   auto class_dims = layout.GetDOFHolderClassificationDimensionsHost();
   auto class_ids = layout.GetDOFHolderClassificationIdsHost();
-  auto coords = layout.GetDOFHolderCoordinates().GetCoordinates();
+  auto coords = layout.GetDOFHolderCoordinates().GetValues();
   auto ent_offsets = layout.GetEntOffsets();
   int mesh_dim = static_cast<int>(coords.extent(1));
 

src/pcms/coupler/field_layout_communicator.h

@@ -41,7 +41,19 @@ class FieldLayoutCommunicator
   template <typename T>
   void SetOutMessageLayout(redev::BidirectionalComm<T>& comm)
   {
-    comm.SetOutMessageLayout(plan_.dest_ranks, plan_.offsets);
+    // plan_.offsets is a per-holder CSR layout. Each holder carries
+    // num_components field values, so the field message's offsets are scaled
+    // by num_components while the destination ranks are unchanged.
+    const int num_comp = layout_.GetNumComponents();
+    if (num_comp == 1) {
+      comm.SetOutMessageLayout(plan_.dest_ranks, plan_.offsets);
+      return;
+    }
+    redev::LOs scaled_offsets(plan_.offsets.size());
+    for (size_t i = 0; i < plan_.offsets.size(); ++i) {
+      scaled_offsets[i] = plan_.offsets[i] * num_comp;
+    }
+    comm.SetOutMessageLayout(plan_.dest_ranks, scaled_offsets);
   }
 
   void UpdateLayout();

src/pcms/coupler/field_serializer.h

@@ -20,27 +20,42 @@ class FieldSerializer
   {
     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.
     if (buffer.size() > 0) {
-      for (LO i = 0; i < static_cast<LO>(data.size()); ++i) {
-        if (owned[i])
-          buffer[permutation[i]] = data[i];
+      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) {
+        if (owned[i]) {
+          for (LO c = 0; c < num_comp; ++c) {
+            buffer[permutation[i] * num_comp + c] = data(i, c);
+          }
+        }
       }
     }
-    return static_cast<int>(data.size());
+    return static_cast<int>(buffer.size());
   }
 
   virtual void Deserialize(
     FieldData<T>& field, const FieldLayout& layout,
     Rank1View<const T, HostMemorySpace> buffer,
     Rank1View<const LO, HostMemorySpace> permutation) const
   {
-    Kokkos::View<T*, HostMemorySpace> sorted("sorted", permutation.size());
+    const LO num_dof = layout.GetNumOwnedDofHolder();
+    const LO num_comp = layout.GetNumComponents();
+    Kokkos::View<T*, HostMemorySpace> sorted("sorted", layout.OwnedSize());
     auto owned = layout.GetOwnedHost();
-    for (LO i = 0; i < static_cast<LO>(sorted.size()); ++i) {
-      if (owned[i])
-        sorted[i] = buffer[permutation[i]];
+    for (LO i = 0; i < num_dof; ++i) {
+      if (owned[i]) {
+        for (LO c = 0; c < num_comp; ++c) {
+          sorted[i * num_comp + c] = buffer[permutation[i] * num_comp + c];
+        }
+      }
     }
-    field.SetDOFHolderDataHost(make_const_array_view(sorted));
+    field.SetDOFHolderDataHost(
+      Rank2View<const T, HostMemorySpace>(sorted.data(), num_dof, num_comp));
   }
 
   virtual ~FieldSerializer() noexcept = default;

src/pcms/coupler/serializer/xgc.h

@@ -35,10 +35,16 @@ class XGCFieldSerializer : public FieldSerializer<T>
 
     auto data = xgc_field->GetDOFHolderDataHost();
     auto owned = layout.GetOwnedHost();
+    // Per-holder plan: owned[i]/permutation[i] index holders; a holder's
+    // num_components values form one contiguous block in the wire buffer.
     if (buffer.size() > 0) {
-      for (LO i = 0; i < static_cast<LO>(data.size()); ++i) {
+      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) {
         if (owned[i]) {
-          buffer[permutation[i]] = data[i];
+          for (LO c = 0; c < num_comp; ++c) {
+            buffer[permutation[i] * num_comp + c] = data(i, c);
+          }
         }
       }
     }
@@ -58,23 +64,30 @@ class XGCFieldSerializer : public FieldSerializer<T>
 
     auto current = xgc_field->GetDOFHolderDataHost();
     auto owned = layout.GetOwnedHost();
+    const LO num_dof = static_cast<LO>(current.extent(0));
+    const LO num_comp = static_cast<LO>(current.extent(1));
     std::vector<T> full_data(current.size());
-    for (size_t i = 0; i < current.size(); ++i) {
-      full_data[i] = current[i];
+    for (LO i = 0; i < num_dof; ++i) {
+      for (LO c = 0; c < num_comp; ++c) {
+        full_data[i * num_comp + c] = current(i, c);
+      }
     }
     if (rank_participates_) {
-      for (LO i = 0; i < static_cast<LO>(current.size()); ++i) {
+      for (LO i = 0; i < num_dof; ++i) {
         if (owned[i]) {
-          full_data[i] = buffer[permutation[i]];
+          for (LO c = 0; c < num_comp; ++c) {
+            full_data[i * num_comp + c] = buffer[permutation[i] * num_comp + c];
+          }
         }
       }
     }
 
     MPI_Bcast(full_data.data(), static_cast<int>(full_data.size()),
               pcms::GetMPIType(T{}), 0, plane_comm_);
 
-    xgc_field->SetDOFHolderDataHost(
-      Rank1View<const T, HostMemorySpace>(full_data.data(), full_data.size()));
+    xgc_field->SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace>(
+      full_data.data(), layout.GetNumOwnedDofHolder(),
+      layout.GetNumComponents()));
   }
 
 private:

src/pcms/discretization/discretization/omega_h.hpp

@@ -28,6 +28,9 @@ class OmegaHDiscretization : public Discretization
   Rank1View<const ClassificationId, DeviceMemorySpace>
   GetEntityClassificationIds(int entity_dim) const override;
 
+  // The Omega_h mesh this discretization is defined on.
+  Omega_h::Mesh& GetMesh() const noexcept { return mesh_; }
+
 private:
   static constexpr int max_entity_dim_ = Region;
 

src/pcms/field/coordinate_system.h

@@ -32,8 +32,8 @@ class CoordinateView
     return coordinate_system_;
   }
 
-  [[nodiscard]] Rank2View<const Real, MemorySpace, LayoutPolicy>
-  GetCoordinates() const noexcept
+  [[nodiscard]] Rank2View<const Real, MemorySpace, LayoutPolicy> GetValues()
+    const noexcept
   {
     return coordinates_;
   }
@@ -42,7 +42,7 @@ class CoordinateView
   // CoordinateTransformation as they are unsafe (i.e., you can break class
   // invariant) passkey pattern?
   [[nodiscard]] Rank2View<const Real, MemorySpace, LayoutPolicy>
-  GetCoordinates() noexcept
+  GetValues() noexcept
   {
     return coordinates_;
   }

src/pcms/field/data/mesh_fields.h

@@ -36,31 +36,35 @@ class MeshFieldsFieldData : public FieldData<T>
 
   const FieldMetadata& GetMetadata() const override { return metadata_; }
 
-  Rank1View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
+  Rank2View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
   {
     Kokkos::deep_copy(host_data_, device_data_);
-    return make_const_array_view(host_data_);
+    return Rank2View<const T, HostMemorySpace>(host_data_.data(),
+                                               layout_->GetNumOwnedDofHolder(),
+                                               layout_->GetNumComponents());
   }
 
-  void SetDOFHolderDataHost(Rank1View<const T, HostMemorySpace> values) override
+  void SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace> values) override
   {
     PCMS_ALWAYS_ASSERT(values.size() ==
                        static_cast<size_t>(layout_->OwnedSize()));
-    CopyHostRank1ViewToDeviceView(device_data_, values);
-    SyncBackend(make_const_array_view(device_data_));
+    CopyHostRank2ViewToDeviceView(device_data_, values);
+    SyncBackend(GetDOFHolderData());
   }
 
-  Rank1View<const T, DeviceMemorySpace> GetDOFHolderData() const override
+  Rank2View<const T, DeviceMemorySpace> GetDOFHolderData() const override
   {
-    return make_const_array_view(device_data_);
+    return Rank2View<const T, DeviceMemorySpace>(
+      device_data_.data(), layout_->GetNumOwnedDofHolder(),
+      layout_->GetNumComponents());
   }
 
-  void SetDOFHolderData(Rank1View<const T, DeviceMemorySpace> values) override
+  void SetDOFHolderData(Rank2View<const T, DeviceMemorySpace> values) override
   {
     PCMS_ALWAYS_ASSERT(values.size() ==
                        static_cast<size_t>(layout_->OwnedSize()));
-    CopyDeviceRank1ViewToDeviceView(device_data_, values);
-    SyncBackend(make_const_array_view(device_data_));
+    CopyDeviceRank2ViewToDeviceView(device_data_, values);
+    SyncBackend(GetDOFHolderData());
   }
 
   std::shared_ptr<MeshFieldBackend<T>> GetMeshFieldBackend() const
@@ -69,21 +73,25 @@ class MeshFieldsFieldData : public FieldData<T>
   }
 
 private:
-  void SyncBackend(Rank1View<const T, DeviceMemorySpace> flat)
+  void SyncBackend(Rank2View<const T, DeviceMemorySpace> data)
   {
     auto nodes_per_dim = layout_->GetNodesPerDim();
     auto num_components = layout_->GetNumComponents();
     auto& mesh = layout_->GetMesh();
-    size_t offset = 0;
+    // data is [dof_holder][component], contiguous node-major, so each mesh
+    // dimension owns a contiguous block of rows; SetData consumes a flat
+    // node-major span over that block.
+    size_t row_offset = 0;
     for (int i = 0; i <= mesh.dim(); ++i) {
       if (nodes_per_dim[i]) {
-        size_t len = static_cast<size_t>(mesh.nents(i)) *
-                     static_cast<size_t>(nodes_per_dim[i]) *
-                     static_cast<size_t>(num_components);
+        size_t num_rows = static_cast<size_t>(mesh.nents(i)) *
+                          static_cast<size_t>(nodes_per_dim[i]);
+        size_t len = num_rows * static_cast<size_t>(num_components);
         Rank1View<const T, DeviceMemorySpace> subspan{
-          flat.data_handle() + offset, len};
+          data.data_handle() + row_offset * static_cast<size_t>(num_components),
+          len};
         mesh_field_->SetData(subspan, nodes_per_dim[i], num_components, i);
-        offset += len;
+        row_offset += num_rows;
       }
     }
   }

src/pcms/field/data/point_cloud.cpp

@@ -10,10 +10,8 @@ namespace pcms
 PointCloud::PointCloud(std::shared_ptr<const PointCloudLayout> layout)
   : layout_(std::move(layout)),
     metadata_{},
-    device_data_("",
-                 layout_->GetDOFHolderCoordinates().GetCoordinates().extent(0)),
-    data_host_("",
-               layout_->GetDOFHolderCoordinates().GetCoordinates().extent(0))
+    device_data_("", layout_->GetDOFHolderCoordinates().GetValues().extent(0)),
+    data_host_("", layout_->GetDOFHolderCoordinates().GetValues().extent(0))
 {
 }
 
@@ -22,30 +20,34 @@ const FieldMetadata& PointCloud::GetMetadata() const
   return metadata_;
 }
 
-Rank1View<const Real, HostMemorySpace> PointCloud::GetDOFHolderDataHost() const
+Rank2View<const Real, HostMemorySpace> PointCloud::GetDOFHolderDataHost() const
 {
   Kokkos::deep_copy(data_host_, device_data_);
-  return make_const_array_view(data_host_);
+  const auto nc = layout_->GetNumComponents();
+  return Rank2View<const Real, HostMemorySpace>(
+    data_host_.data(), static_cast<LO>(data_host_.size()) / nc, nc);
 }
 
 void PointCloud::SetDOFHolderDataHost(
-  Rank1View<const Real, HostMemorySpace> data)
+  Rank2View<const Real, HostMemorySpace> data)
 {
   PCMS_FUNCTION_TIMER;
   PCMS_ALWAYS_ASSERT(data.size() == device_data_.size());
-  CopyHostRank1ViewToDeviceView(device_data_, data);
+  CopyHostRank2ViewToDeviceView(device_data_, data);
 }
 
-Rank1View<const Real, DeviceMemorySpace> PointCloud::GetDOFHolderData() const
+Rank2View<const Real, DeviceMemorySpace> PointCloud::GetDOFHolderData() const
 {
-  return make_const_array_view(device_data_);
+  const auto nc = layout_->GetNumComponents();
+  return Rank2View<const Real, DeviceMemorySpace>(
+    device_data_.data(), static_cast<LO>(device_data_.size()) / nc, nc);
 }
 
-void PointCloud::SetDOFHolderData(Rank1View<const Real, DeviceMemorySpace> data)
+void PointCloud::SetDOFHolderData(Rank2View<const Real, DeviceMemorySpace> data)
 {
   PCMS_FUNCTION_TIMER;
   PCMS_ALWAYS_ASSERT(data.size() == device_data_.size());
-  CopyDeviceRank1ViewToDeviceView(device_data_, data);
+  CopyDeviceRank2ViewToDeviceView(device_data_, data);
 }
 
 } // namespace pcms

src/pcms/field/data/point_cloud.h

@@ -16,12 +16,12 @@ class PointCloud : public FieldData<Real>
 
   const FieldMetadata& GetMetadata() const override;
 
-  Rank1View<const Real, HostMemorySpace> GetDOFHolderDataHost() const override;
+  Rank2View<const Real, HostMemorySpace> GetDOFHolderDataHost() const override;
   void SetDOFHolderDataHost(
-    Rank1View<const Real, HostMemorySpace> data) override;
+    Rank2View<const Real, HostMemorySpace> data) override;
 
-  Rank1View<const Real, DeviceMemorySpace> GetDOFHolderData() const override;
-  void SetDOFHolderData(Rank1View<const Real, DeviceMemorySpace> data) override;
+  Rank2View<const Real, DeviceMemorySpace> GetDOFHolderData() const override;
+  void SetDOFHolderData(Rank2View<const Real, DeviceMemorySpace> data) override;
 
 private:
   std::shared_ptr<const PointCloudLayout> layout_;

src/pcms/field/data/simple.h

@@ -37,29 +37,33 @@ class SimpleFieldData : public FieldData<T>
 
   const FieldMetadata& GetMetadata() const override { return metadata_; }
 
-  Rank1View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
+  Rank2View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
   {
     Kokkos::deep_copy(host_data_, device_data_);
-    return make_const_array_view(host_data_);
+    return Rank2View<const T, HostMemorySpace>(host_data_.data(),
+                                               layout_->GetNumOwnedDofHolder(),
+                                               layout_->GetNumComponents());
   }
 
-  void SetDOFHolderDataHost(Rank1View<const T, HostMemorySpace> values) override
+  void SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace> values) override
   {
     PCMS_ALWAYS_ASSERT(values.size() ==
                        static_cast<size_t>(layout_->OwnedSize()));
-    CopyHostRank1ViewToDeviceView(device_data_, values);
+    CopyHostRank2ViewToDeviceView(device_data_, values);
   }
 
-  Rank1View<const T, DeviceMemorySpace> GetDOFHolderData() const override
+  Rank2View<const T, DeviceMemorySpace> GetDOFHolderData() const override
   {
-    return make_const_array_view(device_data_);
+    return Rank2View<const T, DeviceMemorySpace>(
+      device_data_.data(), layout_->GetNumOwnedDofHolder(),
+      layout_->GetNumComponents());
   }
 
-  void SetDOFHolderData(Rank1View<const T, DeviceMemorySpace> values) override
+  void SetDOFHolderData(Rank2View<const T, DeviceMemorySpace> values) override
   {
     PCMS_ALWAYS_ASSERT(values.size() ==
                        static_cast<size_t>(layout_->OwnedSize()));
-    CopyDeviceRank1ViewToDeviceView(device_data_, values);
+    CopyDeviceRank2ViewToDeviceView(device_data_, values);
   }
 
 private: