API cleanup to use 2D arrays for Set/GetDOFs

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_serializer.h

@@ -20,10 +20,18 @@ class FieldSerializer
   {
     auto data = field.GetDOFHolderDataHost();
     auto owned = layout.GetOwnedHost();
+    // owned/permutation/buffer index the flat node-major wire order
+    // (dof * num_components + component); read the field values through the
+    // 2D [dof][comp] view.
     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) {
+        for (LO c = 0; c < num_comp; ++c) {
+          const LO flat = i * num_comp + c;
+          if (owned[flat])
+            buffer[permutation[flat]] = data(i, c);
+        }
       }
     }
     return static_cast<int>(data.size());
@@ -40,7 +48,8 @@ class FieldSerializer
       if (owned[i])
         sorted[i] = buffer[permutation[i]];
     }
-    field.SetDOFHolderDataHost(make_const_array_view(sorted));
+    field.SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace>(
+      sorted.data(), layout.GetNumOwnedDofHolder(), layout.GetNumComponents()));
   }
 
   virtual ~FieldSerializer() noexcept = default;

src/pcms/coupler/serializer/xgc.h

@@ -36,9 +36,14 @@ class XGCFieldSerializer : public FieldSerializer<T>
     auto data = xgc_field->GetDOFHolderDataHost();
     auto owned = layout.GetOwnedHost();
     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) {
+        for (LO c = 0; c < num_comp; ++c) {
+          const LO flat = i * num_comp + c;
+          if (owned[flat]) {
+            buffer[permutation[flat]] = data(i, c);
+          }
         }
       }
     }
@@ -58,9 +63,13 @@ 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) {
@@ -73,8 +82,9 @@ class XGCFieldSerializer : public FieldSerializer<T>
     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/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:

src/pcms/field/data/xgc.h

@@ -40,37 +40,46 @@ class XGCFieldData : public FieldData<T>
 
   const FieldMetadata& GetMetadata() const override { return metadata_; }
 
-  Rank1View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
+  Rank2View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
   {
-    return Rank1View<const T, HostMemorySpace>(data_.data_handle(),
-                                               data_.size());
+    const auto nc = layout_->GetNumComponents();
+    return Rank2View<const T, HostMemorySpace>(
+      data_.data_handle(), static_cast<LO>(data_.size()) / nc, nc);
   }
 
-  void SetDOFHolderDataHost(Rank1View<const T, HostMemorySpace> values) override
+  void SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace> values) override
   {
     if (values.size() != data_.size()) {
       throw pcms_error("XGCFieldData::SetDOFHolderDataHost: size mismatch");
     }
-    for (size_t i = 0; i < values.size(); ++i) {
-      data_(i) = values[i];
+    const auto num_dof = values.extent(0);
+    const auto num_comp = values.extent(1);
+    for (size_t i = 0; i < num_dof; ++i) {
+      for (size_t c = 0; c < num_comp; ++c) {
+        data_(i * num_comp + c) = values(i, c);
+      }
     }
   }
 
-  Rank1View<const T, DeviceMemorySpace> GetDOFHolderData() const override
+  Rank2View<const T, DeviceMemorySpace> GetDOFHolderData() const override
   {
     Kokkos::View<T*, HostMemorySpace> host_view("GetDOFHolderData_host_view",
                                                 data_.size());
     for (size_t i = 0; i < data_.size(); ++i) {
       host_view(i) = data_(i);
     }
     Kokkos::deep_copy(device_data_, host_view);
-    return make_const_array_view(device_data_);
+    const auto nc = layout_->GetNumComponents();
+    return Rank2View<const T, DeviceMemorySpace>(
+      device_data_.data(), static_cast<LO>(device_data_.extent(0)) / nc, nc);
   }
 
-  void SetDOFHolderData(Rank1View<const T, DeviceMemorySpace> values) override
+  void SetDOFHolderData(Rank2View<const T, DeviceMemorySpace> values) override
   {
-    CopyDeviceRank1ViewToDeviceView(device_data_, values);
-    CopyRank1ViewToHost(data_, values);
+    // Fill the node-major device scratch from the (possibly non-node-major)
+    // 2D view, then mirror it back into the canonical host storage.
+    CopyDeviceRank2ViewToDeviceView(device_data_, values);
+    CopyRank1ViewToHost(data_, make_const_array_view(device_data_));
   }
 
 private:

src/pcms/field/eqdsk_field.h

@@ -44,8 +44,8 @@ inline EQDSKField MakeEQDSKField(const EQDSKData& eqdsk_data)
   auto space = SplineFunctionSpace::FromUniformGrid(
     eqdsk_data.grid, CoordinateSystem::Cartesian);
   auto field = space.CreateField<Real>();
-  field.GetData().SetDOFHolderData(Rank1View<const Real, DeviceMemorySpace>(
-    eqdsk_data.PSIZR.data(), eqdsk_data.PSIZR.extent(0)));
+  field.GetData().SetDOFHolderData(Rank2View<const Real, DeviceMemorySpace>(
+    eqdsk_data.PSIZR.data(), eqdsk_data.PSIZR.extent(0), 1));
   return {std::move(space), std::move(field)};
 }
 

src/pcms/field/evaluator/mesh_fields.h

@@ -128,7 +128,7 @@ class MeshFieldsEvaluatorFactory : public FieldEvaluatorFactory<T>
         "MeshFieldsEvaluatorFactory: NearestBoundary is not supported");
     }
 
-    auto coordinates = coords.GetCoordinates();
+    auto coordinates = coords.GetValues();
     Kokkos::View<Real* [2], DeviceMemorySpace> coords_search(
       "coords_search", coordinates.extent(0));
     Kokkos::parallel_for(

src/pcms/field/evaluator/mls_point_cloud.h

@@ -54,7 +54,7 @@ class MLSPointEvaluator : public PointEvaluator<Real, LayoutPolicy>
     Kokkos::parallel_for(
       "CopyDeviceDataToOmegaHWrite",
       Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, n_sources),
-      KOKKOS_LAMBDA(int i) { src_w[i] = device_data(i); });
+      KOKKOS_LAMBDA(int i) { src_w[i] = device_data(i, 0); });
     Omega_h::Reals source_values(src_w);
 
     auto result = mls_interpolation(