Skip to content

Support multi component in meshFields backend #371

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Sichao25 wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

Support multi component in meshFields backend #371

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d0e32ad
support meshfields multi-component backend
Sichao25 Aug 27, 2026
28ae157
debug meshfields layout mismatch
Sichao25 Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 45 additions & 20 deletions src/pcms/field/data/mesh_fields.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,32 @@
namespace pcms
{

namespace
{
// Functor that flattens component-major 2D data (LayoutLeft) into dof-major
// order.
template <typename T>
struct FlattenFunctor
{
Rank2View<const T, DeviceMemorySpace> data_;
Kokkos::View<T*, DeviceMemorySpace> flat_;
LO row_off_;
size_t nc_;
FlattenFunctor(Rank2View<const T, DeviceMemorySpace> d,
Kokkos::View<T*, DeviceMemorySpace> f, LO ro, size_t nc)
: data_(d), flat_(f), 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) {
flat_(local * nc_ + c) = data_(global_dof, c);
}
}
};
} // namespace

template <typename T>
class MeshFieldsFieldData : public FieldData<T>
{
Expand All @@ -24,9 +50,11 @@ 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_->GetNumOwnedDofHolder()),
static_cast<size_t>(layout_->GetNumComponents())),
device_data_("meshfields_field_data_device",
static_cast<size_t>(layout_->OwnedSize()))
static_cast<size_t>(layout_->GetNumOwnedDofHolder()),
static_cast<size_t>(layout_->GetNumComponents()))
{
if (!mesh_field_) {
throw pcms_error(
Expand All @@ -38,10 +66,8 @@ class MeshFieldsFieldData : public FieldData<T>

Rank2View<const T, HostMemorySpace> GetDOFHolderDataHost() const override
{
Kokkos::deep_copy(host_data_, device_data_);
return Rank2View<const T, HostMemorySpace>(host_data_.data(),
layout_->GetNumOwnedDofHolder(),
layout_->GetNumComponents());
DeepCopyMismatchLayouts(host_data_, device_data_);
return MakeConstRank2View(host_data_);
}

void SetDOFHolderDataHost(Rank2View<const T, HostMemorySpace> values) override
Expand All @@ -54,12 +80,7 @@ class MeshFieldsFieldData : public FieldData<T>

Rank2View<const T, DeviceMemorySpace> GetDOFHolderData() const override
{
// The Rank2View will wrap the dof-major data with layout left when device
// memory is enabled. This may cause issues in multi component cases. See
// issue #342
return Rank2View<const T, DeviceMemorySpace>(
device_data_.data(), layout_->GetNumOwnedDofHolder(),
layout_->GetNumComponents());
return MakeConstRank2View(device_data_);
}

void SetDOFHolderData(Rank2View<const T, DeviceMemorySpace> values) override
Expand All @@ -81,18 +102,22 @@ class MeshFieldsFieldData : public FieldData<T>
auto nodes_per_dim = layout_->GetNodesPerDim();
auto num_components = layout_->GetNumComponents();
auto& mesh = layout_->GetMesh();
// 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.
// device_data_ is rank-2 LayoutLeft (component-major), but SetData
// expects a flat dof-major span.
size_t row_offset = 0;
for (int i = 0; i <= mesh.dim(); ++i) {
if (nodes_per_dim[i]) {
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{
data.data_handle() + row_offset * static_cast<size_t>(num_components),
len};
Kokkos::View<T*, DeviceMemorySpace> flat("sync_flat", len);
Kokkos::parallel_for(
"SyncBackendReorder",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(
0, static_cast<LO>(num_rows)),
FlattenFunctor<T>(data, flat, static_cast<LO>(row_offset),
num_components));
Rank1View<const T, DeviceMemorySpace> subspan(flat.data(), len);
mesh_field_->SetData(subspan, nodes_per_dim[i], num_components, i);
row_offset += num_rows;
}
Expand All @@ -102,8 +127,8 @@ class MeshFieldsFieldData : public FieldData<T>
std::shared_ptr<const MeshFieldsAdapterLayout> layout_;
FieldMetadata metadata_;
std::shared_ptr<MeshFieldBackend<T>> mesh_field_;
mutable Kokkos::View<T*, HostMemorySpace> host_data_;
Kokkos::View<T*, DeviceMemorySpace> device_data_;
mutable Kokkos::View<T**, HostMemorySpace> host_data_;
Kokkos::View<T**, DeviceMemorySpace> device_data_;
};

} // namespace pcms
Expand Down
25 changes: 9 additions & 16 deletions src/pcms/field/evaluator/mesh_fields.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,36 +42,33 @@ class MeshFieldsPointEvaluator : public PointEvaluator<T, LayoutPolicy>
hint_.coordinates_.extent(0) + hint_.num_missing_);
PCMS_ALWAYS_ASSERT(values.extent(1) ==
static_cast<size_t>(layout_->GetNumComponents()));
// ensure that only scalar fields are supported
PCMS_ALWAYS_ASSERT(layout_->GetNumComponents() == 1);
auto const* mesh_field_data =
dynamic_cast<const MeshFieldsFieldData<T>*>(&field.GetData());
if (!mesh_field_data) {
throw pcms_error(
"MeshFieldsPointEvaluator::Evaluate: incompatible FieldData type");
}

// Use device views directly from hint (no copy needed)
auto eval_results = mesh_field_data->GetMeshFieldBackend()->evaluate(
hint_.coordinates_d_, hint_.offsets_d_);

// Scatter results directly on device (no host copy)
int ncomp = layout_->GetNumComponents();
int num_eval = static_cast<int>(eval_results.extent(0));
Kokkos::parallel_for(
"CopyEvalResultsToValues",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(
0, eval_results.extent(0)),
KOKKOS_CLASS_LAMBDA(LO i) {
values(hint_.indices_d_(i), 0) = eval_results(i, 0);
Kokkos::MDRangePolicy<Kokkos::Rank<2>>({0, 0}, {num_eval, ncomp}),
KOKKOS_CLASS_LAMBDA(LO i, int c) {
values(hint_.indices_d_(i), c) = eval_results(i, c);
});

if (hint_.num_missing_ > 0 && hint_.mode_ == OutOfBoundsMode::FILL) {
T fill_val = static_cast<T>(fill_value_);
int num_missing = static_cast<int>(hint_.num_missing_);
Kokkos::parallel_for(
"FillMissingValues",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(
0, hint_.num_missing_),
KOKKOS_CLASS_LAMBDA(LO i) {
values(hint_.missing_indices_d_(i), 0) = fill_val;
Kokkos::MDRangePolicy<Kokkos::Rank<2>>({0, 0}, {num_missing, ncomp}),
KOKKOS_CLASS_LAMBDA(LO i, int c) {
values(hint_.missing_indices_d_(i), c) = fill_val;
});
}
}
Expand All @@ -98,10 +95,6 @@ class MeshFieldsEvaluatorFactory : public FieldEvaluatorFactory<T>
if (mesh_.dim() == 3) {
throw pcms_error("MeshFieldsEvaluatorFactory does not support 3D meshes");
}
if (layout_->GetNumComponents() != 1) {
throw pcms_error(
"MeshFieldsEvaluatorFactory only supports single-component fields");
}
}

const FieldLayout& GetLayout() const override { return *layout_; }
Expand Down
61 changes: 49 additions & 12 deletions src/pcms/field/evaluator/mesh_fields_backend.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@ class MeshFieldBackend
{
public:
virtual ~MeshFieldBackend() = default;
virtual Kokkos::View<T* [1]> evaluate(Kokkos::View<T**> localCoords,
Kokkos::View<LO*> offsets) const = 0;
virtual Kokkos::View<T**> evaluate(Kokkos::View<T**> localCoords,
Kokkos::View<LO*> offsets) const = 0;
virtual void SetData(Rank1View<const T, DeviceMemorySpace> data,
size_t num_nodes, size_t num_components, int dim) = 0;
virtual void GetData(Rank1View<T, DeviceMemorySpace> data, size_t num_nodes,
Expand All @@ -38,21 +38,23 @@ class MeshFieldBackend
// ---------------------------------------------------------------------------
// Concrete backend implementation
// ---------------------------------------------------------------------------
template <typename T, int Dim, int Order>
template <typename T, int Dim, int Order, int NumComponents>
class MeshFieldBackendImpl : public MeshFieldBackend<T>
{
public:
MeshFieldBackendImpl(Omega_h::Mesh& mesh)
: mesh_(mesh),
mesh_field_(mesh),
shape_field_(mesh_field_.template CreateLagrangeField<T, Order, 1>())
shape_field_(
mesh_field_.template CreateLagrangeField<T, Order, NumComponents>())
{
}

Kokkos::View<T* [1]> evaluate(Kokkos::View<T**> localCoords,
Kokkos::View<LO*> offsets) const override
Kokkos::View<T**> evaluate(Kokkos::View<T**> localCoords,
Kokkos::View<LO*> offsets) const override
{
auto self = const_cast<MeshFieldBackendImpl<T, Dim, Order>*>(this);
auto self =
const_cast<MeshFieldBackendImpl<T, Dim, Order, NumComponents>*>(this);
return self->mesh_field_.triangleLocalPointEval(localCoords, offsets,
shape_field_.field);
}
Expand Down Expand Up @@ -92,13 +94,47 @@ class MeshFieldBackendImpl : public MeshFieldBackend<T>
private:
Omega_h::Mesh& mesh_;
MeshField::OmegahMeshField<DefaultExecutionSpace, Dim> mesh_field_;
using FWC = decltype(mesh_field_.template CreateLagrangeField<T, Order, 1>());
using FWC =
decltype(mesh_field_
.template CreateLagrangeField<T, Order, NumComponents>());
FWC shape_field_; // FWC = FieldWithController; keeps ctrlr alive
};

// ---------------------------------------------------------------------------
// Factory function: create a MeshFieldBackend from a layout
// ---------------------------------------------------------------------------

// Helper to dispatch on num_components at runtime for a fixed (Dim, Order).
template <typename T, int Dim, int Order>
std::shared_ptr<MeshFieldBackend<T>> MakeBackendForComponents(
Omega_h::Mesh& mesh, int num_components)
{
switch (num_components) {
case 1:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 1>>(mesh);
case 2:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 2>>(mesh);
case 3:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 3>>(mesh);
case 4:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 4>>(mesh);
case 5:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 5>>(mesh);
case 6:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 6>>(mesh);
case 7:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 7>>(mesh);
case 8:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 8>>(mesh);
case 9:
return std::make_shared<MeshFieldBackendImpl<T, Dim, Order, 9>>(mesh);
default:
throw pcms_error("MeshFieldBackend: num_components " +
std::to_string(num_components) +
" exceeds maximum supported (9).");
}
}

template <typename T>
std::shared_ptr<MeshFieldBackend<T>> MakeMeshFieldBackend(
const MeshFieldsAdapterLayout& layout)
Expand All @@ -115,18 +151,19 @@ std::shared_ptr<MeshFieldBackend<T>> MakeMeshFieldBackend(
throw pcms_error("MeshFieldBackend does not support 3D meshes");
}
auto nodes_per_dim = layout.GetNodesPerDim();
int num_components = layout.GetNumComponents();
if (nodes_per_dim[0] == 1 && nodes_per_dim[1] == 0 && nodes_per_dim[2] == 0 &&
nodes_per_dim[3] == 0) {
switch (mesh.dim()) {
case 1: return std::make_shared<MeshFieldBackendImpl<T, 1, 1>>(mesh);
case 2: return std::make_shared<MeshFieldBackendImpl<T, 2, 1>>(mesh);
case 1: return MakeBackendForComponents<T, 1, 1>(mesh, num_components);
case 2: return MakeBackendForComponents<T, 2, 1>(mesh, num_components);
default: break;
}
} else if (nodes_per_dim[0] == 1 && nodes_per_dim[1] == 1 &&
nodes_per_dim[2] == 0 && nodes_per_dim[3] == 0) {
switch (mesh.dim()) {
case 2: return std::make_shared<MeshFieldBackendImpl<T, 2, 2>>(mesh);
case 3: return std::make_shared<MeshFieldBackendImpl<T, 3, 2>>(mesh);
case 2: return MakeBackendForComponents<T, 2, 2>(mesh, num_components);
case 3: return MakeBackendForComponents<T, 3, 2>(mesh, num_components);
default: break;
}
}
Expand Down
5 changes: 0 additions & 5 deletions src/pcms/field/function_space/lagrange.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,11 +75,6 @@ std::shared_ptr<LagrangeFunctionSpace> LagrangeFunctionSpace::FromMesh(
}
if (backend == Backend::MeshFields) {
#ifdef PCMS_ENABLE_MESHFIELDS
if (num_components != 1) {
throw pcms_error(
"LagrangeFunctionSpace::FromMesh: MeshFields backend only supports "
"single-component fields");
}
std::array<int, 4> nodes_per_dim{};
switch (order) {
case 1: nodes_per_dim = {1, 0, 0, 0}; break;
Expand Down
Loading
Loading