From 7c7774e4e21450854299d7172481667b7cf66bae Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Fri, 11 Sep 2026 13:21:56 -0400 Subject: [PATCH] support distributed fields --- src/pcms/coupler/field_exchange_planner.cpp | 55 ++-- .../coupler/field_layout_communicator.cpp | 2 +- src/pcms/coupler/field_serializer.h | 28 +-- src/pcms/coupler/overlap_mask.h | 20 +- src/pcms/coupler/serializer/xgc.h | 2 +- src/pcms/field/data/mesh_fields.h | 52 +++- src/pcms/field/data/simple.h | 44 +++- src/pcms/field/evaluator/uniform_grid.h | 2 +- src/pcms/field/field.h | 13 + src/pcms/field/field_data.h | 18 ++ src/pcms/field/field_factory.h | 2 +- src/pcms/field/field_layout.h | 234 +++++++++++++++++- src/pcms/field/layout/empty.cpp | 7 +- src/pcms/field/layout/empty.h | 3 +- src/pcms/field/layout/mesh_fields.cpp | 62 ++++- src/pcms/field/layout/mesh_fields.h | 21 +- src/pcms/field/layout/omega_h_entity.cpp | 71 +++++- src/pcms/field/layout/omega_h_entity.h | 18 ++ src/pcms/field/layout/omega_h_lagrange.cpp | 63 ++++- src/pcms/field/layout/omega_h_lagrange.h | 19 +- src/pcms/field/layout/point_cloud.cpp | 7 +- src/pcms/field/layout/point_cloud.h | 3 +- src/pcms/field/layout/uniform_grid.cpp | 8 +- src/pcms/field/layout/uniform_grid.h | 3 +- src/pcms/field/layout/xgc.cpp | 7 +- src/pcms/field/layout/xgc.h | 3 +- src/pcms/field/uniform_grid_binary_field.h | 2 +- src/pcms/pythonapi/bind_field_base.cpp | 15 +- .../omega_h_conservative_projection.cpp | 4 +- .../omega_h_intersection_rhs_integrator.cpp | 2 +- src/pcms/transfer/omega_h_mass_integrator.cpp | 2 +- src/pcms/utility/omega_h_array_utils.h | 21 ++ test/CMakeLists.txt | 10 + test/test_distributed_coupling.cpp | 142 +++++++++++ test/test_distributed_field.cpp | 213 ++++++++++++++++ test/test_field_communication.cpp | 4 +- test/test_proxy_coupling.cpp | 4 +- test/testing.cmake | 4 +- 38 files changed, 1081 insertions(+), 109 deletions(-) create mode 100644 test/test_distributed_coupling.cpp create mode 100644 test/test_distributed_field.cpp diff --git a/src/pcms/coupler/field_exchange_planner.cpp b/src/pcms/coupler/field_exchange_planner.cpp index 88f60d2d..d558287d 100644 --- 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(ent_offsets[d]) && - local_index < static_cast(ent_offsets[d + 1])) { - return d; - } - } - return ent_offsets_len - 1; -} - static size_t GetMessageBlockIndex( LO permutation_entry, Rank1View 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(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(owned.extent(0)); + const LO n_owned = static_cast(coords.extent(0)); std::array 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(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(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( plan.offsets.data(), plan.offsets.size()); std::vector per_rank_offsets(plan.dest_ranks.size()); - for (LO local_index = 0; local_index < static_cast(gids.size()); - ++local_index) { - LO perm_index = plan.permutation[local_index]; + for (LO owned_index = 0; owned_index < static_cast(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((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(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 index 00dabfb7..135ebbe9 100644 --- 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(*overlap_mask) : std::make_unique( - layout.GetGidsHost().size())), + layout.GetNumOwnedDofHolder())), own_mpi_comm_(own_mpi_comm) { gid_comm_ = channel.CreateComm(name_ + "_gids", mpi_comm_); diff --git a/src/pcms/coupler/field_serializer.h b/src/pcms/coupler/field_serializer.h index 723929cf..1918c1d7 100644 --- a/src/pcms/coupler/field_serializer.h +++ b/src/pcms/coupler/field_serializer.h @@ -18,18 +18,12 @@ class FieldSerializer Rank1View buffer, Rank1View 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(data.extent(0)); const LO num_comp = static_cast(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 buffer, Rank1View 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 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 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(sorted.data(), num_dof, num_comp)); + Rank2View(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 index db795b45..e3b7dcb6 100644 --- 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 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(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(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( + 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 index e2658377..a4086bac 100644 --- a/src/pcms/coupler/serializer/xgc.h +++ b/src/pcms/coupler/serializer/xgc.h @@ -90,7 +90,7 @@ class XGCFieldSerializer : public FieldSerializer pcms::GetMPIType(T{}), 0, plane_comm_); xgc_field->SetDOFHolderDataHost(Rank2View( - 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 index 4c27a49d..7da15545 100644 --- 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 #include +#include namespace pcms { @@ -24,9 +26,9 @@ class MeshFieldsFieldData : public FieldData metadata_(metadata), mesh_field_(MakeMeshFieldBackend(*layout_)), host_data_("meshfields_field_data", - static_cast(layout_->OwnedSize())), + static_cast(layout_->LocalSize())), device_data_("meshfields_field_data_device", - static_cast(layout_->OwnedSize())) + static_cast(layout_->LocalSize())) { if (!mesh_field_) { throw pcms_error( @@ -40,14 +42,25 @@ class MeshFieldsFieldData : public FieldData { Kokkos::deep_copy(host_data_, device_data_); return Rank2View(host_data_.data(), - layout_->GetNumOwnedDofHolder(), + layout_->GetNumLocalDofHolder(), layout_->GetNumComponents()); } + Rank2View GetOwnedDOFHolderDataHost() const override + { + return GatherOwnedHostData(*layout_, device_data_, host_data_, + owned_host_data_); + } + + Rank2View GetOwnedDOFHolderData() const override + { + return GatherOwnedDeviceData(*layout_, device_data_, owned_device_data_); + } + void SetDOFHolderDataHost(Rank2View values) override { PCMS_ALWAYS_ASSERT(values.size() == - static_cast(layout_->OwnedSize())); + static_cast(layout_->LocalSize())); CopyHostRank2ViewToDeviceView(device_data_, values); SyncBackend(GetDOFHolderData()); } @@ -58,18 +71,43 @@ class MeshFieldsFieldData : public FieldData // memory is enabled. This may cause issues in multi component cases. See // issue #342 return Rank2View( - device_data_.data(), layout_->GetNumOwnedDofHolder(), + device_data_.data(), layout_->GetNumLocalDofHolder(), layout_->GetNumComponents()); } void SetDOFHolderData(Rank2View values) override { PCMS_ALWAYS_ASSERT(values.size() == - static_cast(layout_->OwnedSize())); + static_cast(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(mesh.nents(dim)) * nodes_per_dim[dim]; + const LO flat_len = num_rows * nc; + const LO flat_off = static_cast(row_offset * static_cast(nc)); + + auto block = Kokkos::subview( + device_data_, Kokkos::make_pair(flat_off, flat_off + flat_len)); + SynchronizeOmegaHBlock(mesh, dim, nc, block); + + row_offset += static_cast(num_rows); + } + + SyncBackend(GetDOFHolderData()); + } + std::shared_ptr> GetMeshFieldBackend() const { return mesh_field_; @@ -103,6 +141,8 @@ class MeshFieldsFieldData : public FieldData FieldMetadata metadata_; std::shared_ptr> mesh_field_; mutable Kokkos::View host_data_; + mutable Kokkos::View owned_host_data_; + mutable Kokkos::View owned_device_data_; Kokkos::View device_data_; }; diff --git a/src/pcms/field/data/simple.h b/src/pcms/field/data/simple.h index 6deda15a..b5cd7be8 100644 --- a/src/pcms/field/data/simple.h +++ b/src/pcms/field/data/simple.h @@ -6,10 +6,14 @@ #include "../field_metadata.h" #include "pcms/utility/arrays.h" #include "pcms/utility/assert.h" +#include +#include "pcms/discretization/discretization/omega_h.hpp" #include "pcms/utility/memory_spaces.h" +#include "pcms/utility/omega_h_array_utils.h" #include #include #include +#include namespace pcms { @@ -29,9 +33,9 @@ class SimpleFieldData : public FieldData : layout_(std::move(layout)), metadata_(metadata), host_data_("simple_field_data", - static_cast(layout_->OwnedSize())), + static_cast(layout_->LocalSize())), device_data_("simple_field_data_device", - static_cast(layout_->OwnedSize())) + static_cast(layout_->LocalSize())) { } @@ -41,35 +45,63 @@ class SimpleFieldData : public FieldData { Kokkos::deep_copy(host_data_, device_data_); return Rank2View(host_data_.data(), - layout_->GetNumOwnedDofHolder(), + layout_->GetNumLocalDofHolder(), layout_->GetNumComponents()); } + Rank2View GetOwnedDOFHolderDataHost() const override + { + return GatherOwnedHostData(*layout_, device_data_, host_data_, + owned_host_data_); + } + + Rank2View GetOwnedDOFHolderData() const override + { + return GatherOwnedDeviceData(*layout_, device_data_, owned_device_data_); + } + void SetDOFHolderDataHost(Rank2View values) override { PCMS_ALWAYS_ASSERT(values.size() == - static_cast(layout_->OwnedSize())); + static_cast(layout_->LocalSize())); CopyHostRank2ViewToDeviceView(device_data_, values); } Rank2View GetDOFHolderData() const override { return Rank2View( - device_data_.data(), layout_->GetNumOwnedDofHolder(), + device_data_.data(), layout_->GetNumLocalDofHolder(), layout_->GetNumComponents()); } void SetDOFHolderData(Rank2View values) override { PCMS_ALWAYS_ASSERT(values.size() == - static_cast(layout_->OwnedSize())); + static_cast(layout_->LocalSize())); CopyDeviceRank2ViewToDeviceView(device_data_, values); } + void SynchronizeGhosts() override + { + auto* oh = dynamic_cast( + layout_->GetDiscretization().get()); + if (!oh) { + return; + } + const int dim = layout_->GetDOFHolderEntityDim(); + if (dim < 0) { + return; + } + SynchronizeOmegaHBlock(oh->GetMesh(), dim, layout_->GetNumComponents(), + device_data_); + } + private: std::shared_ptr layout_; FieldMetadata metadata_; mutable Kokkos::View host_data_; + mutable Kokkos::View owned_host_data_; + mutable Kokkos::View owned_device_data_; Kokkos::View device_data_; }; diff --git a/src/pcms/field/evaluator/uniform_grid.h b/src/pcms/field/evaluator/uniform_grid.h index ed061458..7a8d3208 100644 --- a/src/pcms/field/evaluator/uniform_grid.h +++ b/src/pcms/field/evaluator/uniform_grid.h @@ -81,7 +81,7 @@ class UniformGridPointEvaluator : public PointEvaluator PCMS_ALWAYS_ASSERT(values.extent(0) == static_cast(num_points)); PCMS_ALWAYS_ASSERT(values.extent(1) == static_cast(n_comp)); PCMS_ALWAYS_ASSERT(dof_data.size() == - static_cast(layout_->GetNumOwnedDofHolder() * + static_cast(layout_->GetNumLocalDofHolder() * layout_->GetNumComponents())); auto cell_indices = hint_.cell_indices_; diff --git a/src/pcms/field/field.h b/src/pcms/field/field.h index 6da4261d..97c0bc6b 100644 --- a/src/pcms/field/field.h +++ b/src/pcms/field/field.h @@ -46,6 +46,11 @@ class Field return data_->GetDOFHolderDataHost(); } + Rank2View GetOwnedDOFHolderDataHost() const + { + return data_->GetOwnedDOFHolderDataHost(); + } + void SetDOFHolderDataHost(Rank2View v) { data_->SetDOFHolderDataHost(v); @@ -56,11 +61,19 @@ class Field return data_->GetDOFHolderData(); } + Rank2View GetOwnedDOFHolderData() const + { + return data_->GetOwnedDOFHolderData(); + } + void SetDOFHolderData(Rank2View v) { data_->SetDOFHolderData(v); } + // Synchronize ghost DOF-holder values from their owning ranks. + void SynchronizeGhosts() { data_->SynchronizeGhosts(); } + protected: Field(std::string name, std::shared_ptr layout, std::unique_ptr> data) diff --git a/src/pcms/field/field_data.h b/src/pcms/field/field_data.h index fb3846cd..a4e0a3bf 100644 --- a/src/pcms/field/field_data.h +++ b/src/pcms/field/field_data.h @@ -61,12 +61,30 @@ class FieldData virtual void SetDOFHolderDataHost( Rank2View values) = 0; + // Owned (rank-exclusive) DOF-holder data, compact and owned-indexed. Defaults + // to the local data (owned == local for non-distributed backends). + virtual Rank2View GetOwnedDOFHolderDataHost() const + { + return GetDOFHolderDataHost(); + } + + // Owned (rank-exclusive) DOF-holder data, compact and owned-indexed, on + // device. Defaults to the local device data (owned == local for + // non-distributed backends). + virtual Rank2View GetOwnedDOFHolderData() const + { + return GetDOFHolderData(); + } + // The returned view remains valid until the FieldData is mutated or // destroyed. virtual Rank2View GetDOFHolderData() const = 0; virtual void SetDOFHolderData( Rank2View values) = 0; + // Synchronize ghost DOF-holder values from their owning ranks. + virtual void SynchronizeGhosts() {} + virtual ~FieldData() noexcept = default; }; diff --git a/src/pcms/field/field_factory.h b/src/pcms/field/field_factory.h index 7903c72e..d29cc529 100644 --- a/src/pcms/field/field_factory.h +++ b/src/pcms/field/field_factory.h @@ -17,7 +17,7 @@ namespace detail inline size_t ExpectedFlatFieldDataSize(const FieldLayout& layout) { - return static_cast(layout.GetNumOwnedDofHolder()) * + return static_cast(layout.GetNumLocalDofHolder()) * static_cast(layout.GetNumComponents()); } diff --git a/src/pcms/field/field_layout.h b/src/pcms/field/field_layout.h index b4551e4c..2e6fbd34 100644 --- a/src/pcms/field/field_layout.h +++ b/src/pcms/field/field_layout.h @@ -17,6 +17,20 @@ using EntOffsetsArray = std::array; using ReversePartitionMap = std::map>; +// Returns the mesh entity dimension for the DOF holder at local_index, based on +// the entity-offset array: ent_offsets[d]..ent_offsets[d+1] is the range of DOF +// indices belonging to mesh entity dimension d. +inline int GetMeshEntityDim(LO local_index, const EntOffsetsArray& ent_offsets) +{ + for (int d = 0; d < ent_offsets_len - 1; ++d) { + if (local_index >= static_cast(ent_offsets[d]) && + local_index < static_cast(ent_offsets[d + 1])) { + return d; + } + } + return ent_offsets_len - 1; +} + class FieldLayout { public: @@ -32,13 +46,23 @@ class FieldLayout // number of components int virtual GetNumComponents() const = 0; - // nodes for standard lagrange FEM - LO virtual GetNumOwnedDofHolder() const = 0; + // number of local DOF holders resident on this rank (owned + ghost) + LO virtual GetNumLocalDofHolder() const = 0; + + // number of owned DOF holders (the subset this rank exclusively owns). For + // non-distributed layouts this equals the local count. + LO virtual GetNumOwnedDofHolder() const { return GetNumLocalDofHolder(); } + GO virtual GetNumGlobalDofHolder() const = 0; - // size of buffer that needs to be allocated to represent the field - // # components * NumDOFHolder - LO OwnedSize() const { return GetNumComponents() * GetNumOwnedDofHolder(); }; + // size of buffer needed to hold all local coefficients + // # components * NumLocalDofHolder + LO LocalSize() const { return GetNumComponents() * GetNumLocalDofHolder(); } + + // size of buffer needed to hold owned coefficients + // # components * NumOwnedDofHolder + LO OwnedSize() const { return GetNumComponents() * GetNumOwnedDofHolder(); } + GO GlobalSize() const { return GetNumComponents() * GetNumGlobalDofHolder(); @@ -47,6 +71,9 @@ class FieldLayout virtual Rank1View GetOwnedHost() const = 0; virtual GlobalIDView GetGidsHost() const = 0; + // Device-resident global IDs for all local DOF holders. + virtual GlobalIDView GetGids() const = 0; + // Maps each local DOF holder to its contiguous active index, ordered by GID // within each entity block. Components are not included in the permutation. // For local GIDs [102, 7, 41, 19], the permutation is [3, 0, 2, 1]. Thus @@ -57,16 +84,77 @@ class FieldLayout Kokkos::View GetGlobalToLocalPermutation() const; - // returns true if the field layout is distributed - // if the field layout is distributed, the owned and global dofs are the same + // returns true if the field layout is distributed (holds ghost DOF holders in + // addition to the owned ones); owned and local counts then differ [[nodiscard]] virtual bool IsDistributed() const = 0; virtual EntOffsetsArray GetEntOffsets() const = 0; + // Entity offsets over the OWNED holders (owned-index order), derived from the + // local entity offsets and the owned->local permutation. + EntOffsetsArray GetOwnedEntOffsets() const + { + const auto owned_to_local = GetOwnedToLocalHost(); + if (owned_to_local.size() == 0) { + return GetEntOffsets(); // non-distributed: owned == local + } + EntOffsetsArray offsets{}; + offsets.fill(0); + for (size_t o = 0; o < owned_to_local.size(); ++o) { + const int d = GetMeshEntityDim(owned_to_local(o), GetEntOffsets()); + for (size_t e = static_cast(d) + 1; e < ent_offsets_len; ++e) { + offsets[e] += 1; + } + } + return offsets; + } + virtual CoordinateView GetDOFHolderCoordinates() const = 0; + // Owned (rank-exclusive) global IDs, compact and owned-indexed. Defaults to + // the local array (owned == local for non-distributed layouts). + virtual GlobalIDView GetOwnedGidsHost() const + { + return GetGidsHost(); + } + + // Owned (rank-exclusive) global IDs, compact and owned-indexed, on device. + // Defaults to the local device array (owned == local for non-distributed + // layouts). + virtual GlobalIDView GetOwnedGids() const + { + return GetGids(); + } + + // Owned DOF-holder coordinates, compact and owned-indexed. Defaults to the + // local coordinates. + virtual CoordinateView GetOwnedDOFHolderCoordinates() const + { + return GetDOFHolderCoordinates(); + } + + // Maps owned index (0..GetNumOwnedDofHolder()-1) to its local index. An empty + // view means the identity map (owned == local). + virtual Kokkos::View GetOwnedToLocalHost() const + { + return {}; + } + + // Device analogue of GetOwnedToLocalHost(): maps owned index to its local + // index. An empty view means the identity map (owned == local). + virtual Kokkos::View GetOwnedToLocal() const + { + return {}; + } + virtual int GetDimension() const = 0; + // Entity dimension of the DOF holders (0 = vertex, ..., mesh dim = element), + // used for ghost synchronization. -1 means the layout cannot synchronize + // ghost values via a single mesh entity dimension (e.g. non-mesh or + // multi-dimension layouts). + virtual int GetDOFHolderEntityDim() const { return -1; } + virtual Rank1View GetDOFHolderClassificationDimensionsHost() const = 0; @@ -77,10 +165,6 @@ class FieldLayout protected: // Builds the global-to-local permutation from GetGidsHost()/GetEntOffsets(). - // Derived layouts that support GetGlobalToLocalPermutation() must call this - // from their constructor, once their GIDs and entity offsets are available; - // building eagerly keeps first access free of the data race that lazy - // construction of the mutable cache would introduce. void BuildGlobalToLocalPermutation(); private: @@ -89,5 +173,133 @@ class FieldLayout Kokkos::View global_to_local_; }; +// Compact "owned" (rank-exclusive) views derived from a layout's local arrays. +struct OwnedLayoutData +{ + LO num_owned = 0; + Kokkos::View owned_to_local_host; + Kokkos::View owned_gids_host; + Kokkos::View owned_coords_2d; + Kokkos::View owned_to_local; + Kokkos::View owned_gids; +}; + +template +OwnedLayoutData BuildOwnedLayoutData( + Kokkos::View owned, + GlobalIDView gids, const CoordsView& coords, int dim) +{ + const LO n_local = static_cast(owned.size()); + OwnedLayoutData out; + + out.num_owned = 0; + for (LO i = 0; i < n_local; ++i) { + if (owned(i)) { + ++out.num_owned; + } + } + + out.owned_to_local_host = + Kokkos::View("owned_to_local", out.num_owned); + out.owned_gids_host = + Kokkos::View("owned_gids", out.num_owned); + + LO o = 0; + for (LO i = 0; i < n_local; ++i) { + if (owned(i)) { + out.owned_to_local_host(o) = i; + out.owned_gids_host(o) = static_cast(gids(i)); + ++o; + } + } + + // Gather owned coordinates through a same-layout host mirror of the device + // destination, so the host/device layout mismatch doesn't break deep_copy. + out.owned_coords_2d = + Kokkos::View("owned_coords", out.num_owned, dim); + auto owned_coords_mirror = Kokkos::create_mirror_view(out.owned_coords_2d); + auto coords_host = + Kokkos::create_mirror_view_and_copy(HostMemorySpace(), coords); + for (LO j = 0; j < out.num_owned; ++j) { + const LO i = out.owned_to_local_host(j); + for (int d = 0; d < dim; ++d) { + owned_coords_mirror(j, d) = coords_host(i, d); + } + } + Kokkos::deep_copy(out.owned_coords_2d, owned_coords_mirror); + + // Device-resident copies of the owned->local map and owned GIDs, for the + // device-side owned accessors. + out.owned_to_local = Kokkos::create_mirror_view_and_copy( + DeviceMemorySpace(), out.owned_to_local_host); + out.owned_gids = Kokkos::create_mirror_view_and_copy(DeviceMemorySpace(), + out.owned_gids_host); + + return out; +} + +template +Rank2View GatherOwnedHostData( + const FieldLayout& layout, + const Kokkos::View& device_data, + Kokkos::View& host_data, + Kokkos::View& owned_host_data) +{ + const auto owned_to_local = layout.GetOwnedToLocalHost(); + const LO num_comp = layout.GetNumComponents(); + if (owned_to_local.size() == 0) { + // Non-distributed layout: owned == local. + Kokkos::deep_copy(host_data, device_data); + return Rank2View( + host_data.data(), layout.GetNumLocalDofHolder(), num_comp); + } + const LO num_owned = layout.GetNumOwnedDofHolder(); + Kokkos::deep_copy(host_data, device_data); + owned_host_data = Kokkos::View( + "owned_host_data", + static_cast(num_owned) * static_cast(num_comp)); + for (LO o = 0; o < num_owned; ++o) { + const LO local = owned_to_local(o); + for (LO c = 0; c < num_comp; ++c) { + owned_host_data(static_cast(o) * static_cast(num_comp) + + static_cast(c)) = + host_data(static_cast(local) * static_cast(num_comp) + + static_cast(c)); + } + } + return Rank2View(owned_host_data.data(), num_owned, + num_comp); +} + +template +Rank2View GatherOwnedDeviceData( + const FieldLayout& layout, + const Kokkos::View& device_data, + Kokkos::View& owned_device_data) +{ + const auto owned_to_local = layout.GetOwnedToLocal(); + const LO num_comp = layout.GetNumComponents(); + if (owned_to_local.size() == 0) { + // Non-distributed layout: owned == local. + return Rank2View( + device_data.data(), layout.GetNumLocalDofHolder(), num_comp); + } + const LO num_owned = layout.GetNumOwnedDofHolder(); + owned_device_data = Kokkos::View( + "owned_device_data", + static_cast(num_owned) * static_cast(num_comp)); + Kokkos::parallel_for( + "gather_owned_device_data", + static_cast(num_owned) * static_cast(num_comp), + KOKKOS_LAMBDA(LO i) { + const LO o = i / num_comp; + const LO c = i % num_comp; + const LO local = owned_to_local(o); + owned_device_data(i) = device_data(local * num_comp + c); + }); + return Rank2View(owned_device_data.data(), + num_owned, num_comp); +} + } // namespace pcms #endif // PCMS_FIELD_LAYOUT_H diff --git a/src/pcms/field/layout/empty.cpp b/src/pcms/field/layout/empty.cpp index 83fef485..2e7bfd03 100644 --- a/src/pcms/field/layout/empty.cpp +++ b/src/pcms/field/layout/empty.cpp @@ -28,7 +28,7 @@ int EmptyFieldLayout::GetNumComponents() const return 1; } -LO EmptyFieldLayout::GetNumOwnedDofHolder() const +LO EmptyFieldLayout::GetNumLocalDofHolder() const { return 0; } @@ -48,6 +48,11 @@ GlobalIDView EmptyFieldLayout::GetGidsHost() const return make_const_array_view(gids_host_); } +GlobalIDView EmptyFieldLayout::GetGids() const +{ + return GlobalIDView(gids_.data(), gids_.size()); +} + bool EmptyFieldLayout::IsDistributed() const { return false; diff --git a/src/pcms/field/layout/empty.h b/src/pcms/field/layout/empty.h index 991d58ff..25f9beb1 100644 --- a/src/pcms/field/layout/empty.h +++ b/src/pcms/field/layout/empty.h @@ -17,10 +17,11 @@ class EmptyFieldLayout : public FieldLayout const noexcept override; int GetNumComponents() const override; - LO GetNumOwnedDofHolder() const override; + LO GetNumLocalDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; + GlobalIDView GetGids() const override; bool IsDistributed() const override; EntOffsetsArray GetEntOffsets() const override; CoordinateView GetDOFHolderCoordinates() const override; diff --git a/src/pcms/field/layout/mesh_fields.cpp b/src/pcms/field/layout/mesh_fields.cpp index fa086662..01d2e693 100644 --- a/src/pcms/field/layout/mesh_fields.cpp +++ b/src/pcms/field/layout/mesh_fields.cpp @@ -133,7 +133,7 @@ MeshFieldsAdapterLayout::MeshFieldsAdapterLayout( num_components_(num_components), coordinate_system_(coordinate_system), nodes_per_dim_(nodes_per_dim), - dof_holder_coords_("", GetNumOwnedDofHolder(), mesh_.dim()), + dof_holder_coords_("", GetNumLocalDofHolder(), mesh_.dim()), class_ids_(GetNumEnts()), class_dims_(class_ids_.size()), owned_("", class_dims_.size()), @@ -193,6 +193,7 @@ MeshFieldsAdapterLayout::MeshFieldsAdapterLayout( } } gids_host_ = Omega_h::HostWrite(gids_); + BuildOwnedViews(); int n = class_ids_.size(); classification_dims_host_ = @@ -220,7 +221,7 @@ int MeshFieldsAdapterLayout::GetNumComponents() const return num_components_; } -LO MeshFieldsAdapterLayout::GetNumOwnedDofHolder() const +LO MeshFieldsAdapterLayout::GetNumLocalDofHolder() const { LO count = 0; for (int i = 0; i <= mesh_.dim(); ++i) { @@ -229,6 +230,11 @@ LO MeshFieldsAdapterLayout::GetNumOwnedDofHolder() const return count; } +LO MeshFieldsAdapterLayout::GetNumOwnedDofHolder() const +{ + return num_owned_; +} + GO MeshFieldsAdapterLayout::GetNumGlobalDofHolder() const { LO count = 0; @@ -255,6 +261,11 @@ GlobalIDView MeshFieldsAdapterLayout::GetGidsHost() const return GlobalIDView(gids_host_.data(), gids_host_.size()); } +GlobalIDView MeshFieldsAdapterLayout::GetGids() const +{ + return GlobalIDView(gids_.data(), gids_.size()); +} + CoordinateView MeshFieldsAdapterLayout::GetDOFHolderCoordinates() const { @@ -262,6 +273,53 @@ MeshFieldsAdapterLayout::GetDOFHolderCoordinates() const return CoordinateView{coordinate_system_, coords_view}; } +GlobalIDView MeshFieldsAdapterLayout::GetOwnedGidsHost() const +{ + return GlobalIDView(owned_gids_host_.data(), + owned_gids_host_.size()); +} + +GlobalIDView MeshFieldsAdapterLayout::GetOwnedGids() const +{ + return GlobalIDView(owned_gids_.data(), + owned_gids_.size()); +} + +CoordinateView +MeshFieldsAdapterLayout::GetOwnedDOFHolderCoordinates() const +{ + auto coords_view = MakeConstRank2View(owned_coords_2d_); + return CoordinateView{coordinate_system_, coords_view}; +} + +Kokkos::View +MeshFieldsAdapterLayout::GetOwnedToLocalHost() const +{ + return owned_to_local_host_; +} + +Kokkos::View +MeshFieldsAdapterLayout::GetOwnedToLocal() const +{ + return owned_to_local_; +} + +void MeshFieldsAdapterLayout::BuildOwnedViews() +{ + Kokkos::deep_copy(owned_host_, owned_); + + auto owned = BuildOwnedLayoutData( + owned_host_, + GlobalIDView(gids_host_.data(), gids_host_.size()), + dof_holder_coords_, mesh_.dim()); + num_owned_ = owned.num_owned; + owned_to_local_host_ = owned.owned_to_local_host; + owned_gids_host_ = owned.owned_gids_host; + owned_coords_2d_ = owned.owned_coords_2d; + owned_to_local_ = owned.owned_to_local; + owned_gids_ = owned.owned_gids; +} + bool MeshFieldsAdapterLayout::IsDistributed() const { return true; diff --git a/src/pcms/field/layout/mesh_fields.h b/src/pcms/field/layout/mesh_fields.h index 1c44760c..39040c16 100644 --- a/src/pcms/field/layout/mesh_fields.h +++ b/src/pcms/field/layout/mesh_fields.h @@ -26,15 +26,23 @@ class MeshFieldsAdapterLayout : public FieldLayout int GetNumComponents() const override; // nodes for standard lagrange FEM + LO GetNumLocalDofHolder() const override; LO GetNumOwnedDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; + GlobalIDView GetGids() const override; + GlobalIDView GetOwnedGidsHost() const override; + GlobalIDView GetOwnedGids() const override; CoordinateView GetDOFHolderCoordinates() const override; + CoordinateView GetOwnedDOFHolderCoordinates() + const override; + Kokkos::View GetOwnedToLocalHost() const override; + Kokkos::View GetOwnedToLocal() const override; - // returns true if the field layout is distributed - // if the field layout is distributed, the owned and global dofs are the same + // returns true if the field layout is distributed (holds ghost DOF holders + // in addition to the owned ones); owned and local counts then differ [[nodiscard]] bool IsDistributed() const override; EntOffsetsArray GetEntOffsets() const override; @@ -73,6 +81,15 @@ class MeshFieldsAdapterLayout : public FieldLayout Kokkos::View owned_host_; Kokkos::View classification_dims_host_; Kokkos::View classification_ids_host_; + LO num_owned_ = 0; + Kokkos::View owned_gids_host_; + Kokkos::View owned_coords_2d_; + Kokkos::View owned_to_local_host_; + Kokkos::View owned_gids_; + Kokkos::View owned_to_local_; + + void BuildOwnedViews(); + std::shared_ptr discretization_; }; diff --git a/src/pcms/field/layout/omega_h_entity.cpp b/src/pcms/field/layout/omega_h_entity.cpp index 32f39043..e4c62dc3 100644 --- a/src/pcms/field/layout/omega_h_entity.cpp +++ b/src/pcms/field/layout/omega_h_entity.cpp @@ -74,6 +74,7 @@ OmegaHEntityLayout::OmegaHEntityLayout(Omega_h::Mesh& mesh, int entity_dim, gids_host_ = Omega_h::HostWrite(gids_); Kokkos::deep_copy(owned_host_, owned_); + BuildOwnedViews(); class_dims_ = Omega_h::Read(class_dims_); class_ids_ = Omega_h::Read(class_ids_); @@ -97,11 +98,16 @@ int OmegaHEntityLayout::GetNumComponents() const return num_components_; } -LO OmegaHEntityLayout::GetNumOwnedDofHolder() const +LO OmegaHEntityLayout::GetNumLocalDofHolder() const { return static_cast(coords_.size() / dimension_); } +LO OmegaHEntityLayout::GetNumOwnedDofHolder() const +{ + return num_owned_; +} + GO OmegaHEntityLayout::GetNumGlobalDofHolder() const { return num_global_dof_holder_; @@ -117,17 +123,71 @@ GlobalIDView OmegaHEntityLayout::GetGidsHost() const return GlobalIDView(gids_host_.data(), gids_host_.size()); } +GlobalIDView OmegaHEntityLayout::GetGids() const +{ + return GlobalIDView(gids_.data(), gids_.size()); +} + CoordinateView OmegaHEntityLayout::GetDOFHolderCoordinates() const { using LayoutPolicy = detail::default_layout_for_memory_space_t; Rank2View coords_view( - coords_2d_.data(), GetNumOwnedDofHolder(), dimension_); + coords_2d_.data(), GetNumLocalDofHolder(), dimension_); return CoordinateView{coordinate_system_, coords_view}; } +GlobalIDView OmegaHEntityLayout::GetOwnedGidsHost() const +{ + return GlobalIDView(owned_gids_host_.data(), + owned_gids_host_.size()); +} + +GlobalIDView OmegaHEntityLayout::GetOwnedGids() const +{ + return GlobalIDView(owned_gids_.data(), + owned_gids_.size()); +} + +CoordinateView +OmegaHEntityLayout::GetOwnedDOFHolderCoordinates() const +{ + using LayoutPolicy = + detail::default_layout_for_memory_space_t; + Rank2View coords_view( + owned_coords_2d_.data(), num_owned_, dimension_); + return CoordinateView{coordinate_system_, + coords_view}; +} + +Kokkos::View +OmegaHEntityLayout::GetOwnedToLocalHost() const +{ + return owned_to_local_host_; +} + +Kokkos::View OmegaHEntityLayout::GetOwnedToLocal() + const +{ + return owned_to_local_; +} + +void OmegaHEntityLayout::BuildOwnedViews() +{ + auto owned = BuildOwnedLayoutData( + owned_host_, + GlobalIDView(gids_host_.data(), gids_host_.size()), + coords_2d_, dimension_); + num_owned_ = owned.num_owned; + owned_to_local_host_ = owned.owned_to_local_host; + owned_gids_host_ = owned.owned_gids_host; + owned_coords_2d_ = owned.owned_coords_2d; + owned_to_local_ = owned.owned_to_local; + owned_gids_ = owned.owned_gids; +} + bool OmegaHEntityLayout::IsDistributed() const { return true; @@ -137,7 +197,7 @@ EntOffsetsArray OmegaHEntityLayout::GetEntOffsets() const { EntOffsetsArray offsets{}; offsets.fill(0); - const auto n = static_cast(GetNumOwnedDofHolder()); + const auto n = static_cast(GetNumLocalDofHolder()); for (int i = entity_dim_ + 1; i < ent_offsets_len; ++i) offsets[i] = n; return offsets; @@ -148,6 +208,11 @@ int OmegaHEntityLayout::GetDimension() const return dimension_; } +int OmegaHEntityLayout::GetDOFHolderEntityDim() const +{ + return entity_dim_; +} + Rank1View OmegaHEntityLayout::GetDOFHolderClassificationDimensionsHost() const { diff --git a/src/pcms/field/layout/omega_h_entity.h b/src/pcms/field/layout/omega_h_entity.h index 4c60545a..04c5b28c 100644 --- a/src/pcms/field/layout/omega_h_entity.h +++ b/src/pcms/field/layout/omega_h_entity.h @@ -25,16 +25,25 @@ class OmegaHEntityLayout : public FieldLayout const noexcept override; int GetNumComponents() const override; + LO GetNumLocalDofHolder() const override; LO GetNumOwnedDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; + GlobalIDView GetGids() const override; + GlobalIDView GetOwnedGidsHost() const override; + GlobalIDView GetOwnedGids() const override; CoordinateView GetDOFHolderCoordinates() const override; + CoordinateView GetOwnedDOFHolderCoordinates() + const override; + Kokkos::View GetOwnedToLocalHost() const override; + Kokkos::View GetOwnedToLocal() const override; [[nodiscard]] bool IsDistributed() const override; EntOffsetsArray GetEntOffsets() const override; int GetDimension() const override; + int GetDOFHolderEntityDim() const override; Rank1View GetDOFHolderClassificationDimensionsHost() const override; @@ -59,6 +68,15 @@ class OmegaHEntityLayout : public FieldLayout Kokkos::View owned_host_; Kokkos::View classification_dims_host_; Kokkos::View classification_ids_host_; + LO num_owned_ = 0; + Kokkos::View owned_gids_host_; + Kokkos::View owned_coords_2d_; + Kokkos::View owned_to_local_host_; + Kokkos::View owned_gids_; + Kokkos::View owned_to_local_; + + void BuildOwnedViews(); + std::shared_ptr discretization_; }; diff --git a/src/pcms/field/layout/omega_h_lagrange.cpp b/src/pcms/field/layout/omega_h_lagrange.cpp index c9d7bdb9..8d83b8f8 100644 --- a/src/pcms/field/layout/omega_h_lagrange.cpp +++ b/src/pcms/field/layout/omega_h_lagrange.cpp @@ -126,6 +126,7 @@ OmegaHLagrangeLayout::OmegaHLagrangeLayout(Omega_h::Mesh& mesh, int order, owned_host_ = Kokkos::View("owned_host", owned_.size()); Kokkos::deep_copy(owned_host_, owned_); + BuildOwnedViews(); class_ids_ = Omega_h::Read( mesh_.get_array(entity_dim, "class_id")); @@ -172,6 +173,7 @@ OmegaHLagrangeLayout::OmegaHLagrangeLayout( owned_host_ = Kokkos::View("owned_host", owned_.size()); Kokkos::deep_copy(owned_host_, owned_); + BuildOwnedViews(); class_ids_ = Omega_h::Read( mesh_.get_array(entity_dim, "class_id")); @@ -207,11 +209,16 @@ int OmegaHLagrangeLayout::GetNumComponents() const return num_components_; } -LO OmegaHLagrangeLayout::GetNumOwnedDofHolder() const +LO OmegaHLagrangeLayout::GetNumLocalDofHolder() const { return mesh_.nents(EntityDimForOrder(order_, mesh_.dim())); } +LO OmegaHLagrangeLayout::GetNumOwnedDofHolder() const +{ + return num_owned_; +} + GO OmegaHLagrangeLayout::GetNumGlobalDofHolder() const { return mesh_.nglobal_ents(EntityDimForOrder(order_, mesh_.dim())); @@ -246,6 +253,55 @@ OmegaHLagrangeLayout::GetDOFHolderCoordinates() const coords_view}; } +GlobalIDView OmegaHLagrangeLayout::GetOwnedGidsHost() const +{ + return GlobalIDView(owned_gids_host_.data(), + owned_gids_host_.size()); +} + +GlobalIDView OmegaHLagrangeLayout::GetOwnedGids() const +{ + return GlobalIDView(owned_gids_.data(), + owned_gids_.size()); +} + +CoordinateView +OmegaHLagrangeLayout::GetOwnedDOFHolderCoordinates() const +{ + using LayoutPolicy = + detail::default_layout_for_memory_space_t; + Rank2View coords_view( + owned_coords_2d_.data(), num_owned_, mesh_.dim()); + return CoordinateView{coordinate_system_, + coords_view}; +} + +Kokkos::View +OmegaHLagrangeLayout::GetOwnedToLocalHost() const +{ + return owned_to_local_host_; +} + +Kokkos::View +OmegaHLagrangeLayout::GetOwnedToLocal() const +{ + return owned_to_local_; +} + +void OmegaHLagrangeLayout::BuildOwnedViews() +{ + auto owned = BuildOwnedLayoutData( + owned_host_, + GlobalIDView(gids_host_.data(), gids_host_.size()), + coords_2d_, mesh_.dim()); + num_owned_ = owned.num_owned; + owned_to_local_host_ = owned.owned_to_local_host; + owned_gids_host_ = owned.owned_gids_host; + owned_coords_2d_ = owned.owned_coords_2d; + owned_to_local_ = owned.owned_to_local; + owned_gids_ = owned.owned_gids; +} + bool OmegaHLagrangeLayout::IsDistributed() const { return true; @@ -271,6 +327,11 @@ int OmegaHLagrangeLayout::GetDimension() const return mesh_.dim(); } +int OmegaHLagrangeLayout::GetDOFHolderEntityDim() const +{ + return EntityDimForOrder(order_, mesh_.dim()); +} + Rank1View OmegaHLagrangeLayout::GetDOFHolderClassificationDimensionsHost() const { diff --git a/src/pcms/field/layout/omega_h_lagrange.h b/src/pcms/field/layout/omega_h_lagrange.h index 68154979..b9ac5512 100644 --- a/src/pcms/field/layout/omega_h_lagrange.h +++ b/src/pcms/field/layout/omega_h_lagrange.h @@ -30,19 +30,27 @@ class OmegaHLagrangeLayout : public FieldLayout const noexcept override; int GetNumComponents() const override; + LO GetNumLocalDofHolder() const override; LO GetNumOwnedDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; - GlobalIDView GetGids() const; + GlobalIDView GetGids() const override; + GlobalIDView GetOwnedGidsHost() const override; + GlobalIDView GetOwnedGids() const override; CoordinateView GetDOFHolderCoordinates() const override; + CoordinateView GetOwnedDOFHolderCoordinates() + const override; + Kokkos::View GetOwnedToLocalHost() const override; + Kokkos::View GetOwnedToLocal() const override; [[nodiscard]] bool IsDistributed() const override; EntOffsetsArray GetEntOffsets() const override; int GetDimension() const override; + int GetDOFHolderEntityDim() const override; Rank1View GetDOFHolderClassificationDimensionsHost() const override; @@ -70,6 +78,15 @@ class OmegaHLagrangeLayout : public FieldLayout Omega_h::Read class_dims_; Kokkos::View classification_dims_host_; Kokkos::View classification_ids_host_; + LO num_owned_ = 0; + Kokkos::View owned_gids_host_; + Kokkos::View owned_coords_2d_; + Kokkos::View owned_to_local_host_; + Kokkos::View owned_gids_; + Kokkos::View owned_to_local_; + + void BuildOwnedViews(); + std::shared_ptr discretization_; }; diff --git a/src/pcms/field/layout/point_cloud.cpp b/src/pcms/field/layout/point_cloud.cpp index ffa9ca4d..12e232ac 100644 --- a/src/pcms/field/layout/point_cloud.cpp +++ b/src/pcms/field/layout/point_cloud.cpp @@ -92,7 +92,7 @@ int PointCloudLayout::GetNumComponents() const return components_; } -LO PointCloudLayout::GetNumOwnedDofHolder() const +LO PointCloudLayout::GetNumLocalDofHolder() const { return coords_.extent(0); } @@ -112,6 +112,11 @@ GlobalIDView PointCloudLayout::GetGidsHost() const return GlobalIDView(gids_host_.data(), gids_host_.size()); } +GlobalIDView PointCloudLayout::GetGids() const +{ + return GlobalIDView(gids_.data(), gids_.size()); +} + CoordinateView PointCloudLayout::GetDOFHolderCoordinates() const { diff --git a/src/pcms/field/layout/point_cloud.h b/src/pcms/field/layout/point_cloud.h index 50e6e987..0d14cbad 100644 --- a/src/pcms/field/layout/point_cloud.h +++ b/src/pcms/field/layout/point_cloud.h @@ -22,11 +22,12 @@ class PointCloudLayout : public FieldLayout int GetNumComponents() const override; // nodes for standard lagrange FEM - LO GetNumOwnedDofHolder() const override; + LO GetNumLocalDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; + GlobalIDView GetGids() const override; CoordinateView GetDOFHolderCoordinates() const override; [[nodiscard]] bool IsDistributed() const override; diff --git a/src/pcms/field/layout/uniform_grid.cpp b/src/pcms/field/layout/uniform_grid.cpp index 08f4823d..cc169bf7 100644 --- a/src/pcms/field/layout/uniform_grid.cpp +++ b/src/pcms/field/layout/uniform_grid.cpp @@ -205,7 +205,7 @@ int UniformGridFieldLayout::GetNumComponents() const } template -LO UniformGridFieldLayout::GetNumOwnedDofHolder() const +LO UniformGridFieldLayout::GetNumLocalDofHolder() const { return GetNumDofHolders(); } @@ -229,6 +229,12 @@ GlobalIDView UniformGridFieldLayout::GetGidsHost() const return GlobalIDView(gids_host_.data(), gids_host_.size()); } +template +GlobalIDView UniformGridFieldLayout::GetGids() const +{ + return GlobalIDView(gids_.data(), gids_.size()); +} + template CoordinateView UniformGridFieldLayout::GetDOFHolderCoordinates() const diff --git a/src/pcms/field/layout/uniform_grid.h b/src/pcms/field/layout/uniform_grid.h index f4a78a85..f40068a2 100644 --- a/src/pcms/field/layout/uniform_grid.h +++ b/src/pcms/field/layout/uniform_grid.h @@ -23,11 +23,12 @@ class UniformGridFieldLayout : public FieldLayout const noexcept override; int GetNumComponents() const override; - LO GetNumOwnedDofHolder() const override; + LO GetNumLocalDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; + GlobalIDView GetGids() const override; CoordinateView GetDOFHolderCoordinates() const override; [[nodiscard]] bool IsDistributed() const override; diff --git a/src/pcms/field/layout/xgc.cpp b/src/pcms/field/layout/xgc.cpp index 18729189..1a3ebd59 100644 --- a/src/pcms/field/layout/xgc.cpp +++ b/src/pcms/field/layout/xgc.cpp @@ -124,7 +124,7 @@ int XGCFieldLayout::GetNumComponents() const return 1; } -LO XGCFieldLayout::GetNumOwnedDofHolder() const +LO XGCFieldLayout::GetNumLocalDofHolder() const { return num_plane_nodes_; } @@ -144,6 +144,11 @@ GlobalIDView XGCFieldLayout::GetGidsHost() const return make_const_array_view(gids_host_); } +GlobalIDView XGCFieldLayout::GetGids() const +{ + return GlobalIDView(gids_.data(), gids_.size()); +} + bool XGCFieldLayout::IsDistributed() const { return false; diff --git a/src/pcms/field/layout/xgc.h b/src/pcms/field/layout/xgc.h index a755dcb0..9db9dd82 100644 --- a/src/pcms/field/layout/xgc.h +++ b/src/pcms/field/layout/xgc.h @@ -21,10 +21,11 @@ class XGCFieldLayout : public FieldLayout const noexcept override; int GetNumComponents() const override; - LO GetNumOwnedDofHolder() const override; + LO GetNumLocalDofHolder() const override; GO GetNumGlobalDofHolder() const override; Rank1View GetOwnedHost() const override; GlobalIDView GetGidsHost() const override; + GlobalIDView GetGids() const override; bool IsDistributed() const override; EntOffsetsArray GetEntOffsets() const override; CoordinateView GetDOFHolderCoordinates() const override; diff --git a/src/pcms/field/uniform_grid_binary_field.h b/src/pcms/field/uniform_grid_binary_field.h index bce42439..e9683683 100644 --- a/src/pcms/field/uniform_grid_binary_field.h +++ b/src/pcms/field/uniform_grid_binary_field.h @@ -45,7 +45,7 @@ CreateUniformGridBinaryField(Omega_h::Mesh& mesh, const UniformGrid& grid) auto coord_view = layout->GetDOFHolderCoordinates(); auto coords = coord_view.GetValues(); - LO n = layout->GetNumOwnedDofHolder(); + LO n = layout->GetNumLocalDofHolder(); Kokkos::View coords_d("coords_d", n); Kokkos::parallel_for( diff --git a/src/pcms/pythonapi/bind_field_base.cpp b/src/pcms/pythonapi/bind_field_base.cpp index eca1c98a..bf1a67f2 100644 --- a/src/pcms/pythonapi/bind_field_base.cpp +++ b/src/pcms/pythonapi/bind_field_base.cpp @@ -244,10 +244,17 @@ void bind_create_field_module(py::module& m) .def( "get_num_dof_holders", + [](const Field& self) { + return self.GetLayout().GetNumLocalDofHolder(); + }, + "Number of local DOF holders (owned + ghost) on this rank") + + .def( + "get_num_owned_dof_holders", [](const Field& self) { return self.GetLayout().GetNumOwnedDofHolder(); }, - "Number of owned DOF holders (nodes/elements)") + "Number of owned (rank-exclusive) DOF holders") .def( "get_num_components", @@ -284,7 +291,11 @@ void bind_create_field_module(py::module& m) ptr[i * coords_host.extent(1) + j] = coords_host(i, j); return result; }, - "DOF holder coordinates as a 2D numpy array (num_dof_holders × dim)"); + "DOF holder coordinates as a 2D numpy array (num_dof_holders × dim)") + + .def( + "synchronize_ghosts", [](Field& self) { self.SynchronizeGhosts(); }, + "Synchronize ghost DOF-holder values from their owning ranks"); py::class_>(m, "FunctionSpace") .def( diff --git a/src/pcms/transfer/omega_h_conservative_projection.cpp b/src/pcms/transfer/omega_h_conservative_projection.cpp index 31b1b9e7..be10bd38 100644 --- a/src/pcms/transfer/omega_h_conservative_projection.cpp +++ b/src/pcms/transfer/omega_h_conservative_projection.cpp @@ -64,7 +64,7 @@ OmegaHConservativeProjection::OmegaHConservativeProjection( *rhs_integrator_); target_values_ = Kokkos::View( "conservative_projection_target_values", - target_layout_->GetNumOwnedDofHolder(), target_layout_->GetNumComponents()); + target_layout_->GetNumLocalDofHolder(), target_layout_->GetNumComponents()); } // Defined here so that GalerkinProjectionSolver (forward-declared in the @@ -78,7 +78,7 @@ void OmegaHConservativeProjection::Apply(const Field& source, const auto solution = solver_->Solve(*evaluator_, source); const auto global_to_local = target_layout_->GetGlobalToLocalPermutation(); - const int num_dof_holders = target_layout_->GetNumOwnedDofHolder(); + const int num_dof_holders = target_layout_->GetNumLocalDofHolder(); const int num_components = target_layout_->GetNumComponents(); auto target_values = target_values_; Kokkos::parallel_for( diff --git a/src/pcms/transfer/omega_h_intersection_rhs_integrator.cpp b/src/pcms/transfer/omega_h_intersection_rhs_integrator.cpp index 742b8684..332b9696 100644 --- a/src/pcms/transfer/omega_h_intersection_rhs_integrator.cpp +++ b/src/pcms/transfer/omega_h_intersection_rhs_integrator.cpp @@ -176,7 +176,7 @@ Data BuildDataImpl(const OmegaHLagrangeLayout& source_layout, d.coeffs = std::move(coeffs); d.ndof_per_elem = ndof; d.num_target_dofs = - static_cast(target_layout.GetNumOwnedDofHolder()); + static_cast(target_layout.GetNumLocalDofHolder()); return d; } diff --git a/src/pcms/transfer/omega_h_mass_integrator.cpp b/src/pcms/transfer/omega_h_mass_integrator.cpp index 4538474a..c44c0f6e 100644 --- a/src/pcms/transfer/omega_h_mass_integrator.cpp +++ b/src/pcms/transfer/omega_h_mass_integrator.cpp @@ -94,7 +94,7 @@ Mat BuildOmegaHMassMatrixImpl(Omega_h::Mesh& mesh, { const auto global_to_local = target_layout.GetGlobalToLocalPermutation(); const PetscInt num_dofs = - static_cast(target_layout.GetNumOwnedDofHolder()); + static_cast(target_layout.GetNumLocalDofHolder()); const int nelems = mesh.nelems(); Mat mat = nullptr; diff --git a/src/pcms/utility/omega_h_array_utils.h b/src/pcms/utility/omega_h_array_utils.h index 9b76d97a..63d8ebd9 100644 --- a/src/pcms/utility/omega_h_array_utils.h +++ b/src/pcms/utility/omega_h_array_utils.h @@ -2,8 +2,11 @@ #define PCMS_UTILITY_OMEGA_H_ARRAY_UTILS_H #include +#include #include "pcms/utility/memory_spaces.h" +#include + namespace pcms { @@ -65,6 +68,24 @@ inline Kokkos::View ConvertCoordsTo2D( return coords_2d; } +// Synchronize a block of mesh entity data with Omega_h's internal arrays. +template +void SynchronizeOmegaHBlock(Omega_h::Mesh& mesh, int dim, int nc, + BlockView block) +{ + if constexpr (std::is_same_v) { + Omega_h::Write arr(static_cast(block.size()), + "field_data_block"); + Kokkos::deep_copy(arr.view(), block); + auto synced = mesh.sync_array(dim, Omega_h::Read(arr), nc); + Kokkos::deep_copy(block, synced.view()); + } else { + auto synced = + mesh.sync_array(dim, Omega_h::Read(Omega_h::Write(block)), nc); + Kokkos::deep_copy(block, synced.view()); + } +} + } // namespace pcms #endif // PCMS_UTILITY_OMEGA_H_ARRAY_UTILS_H diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index b51395eb..86e25fb2 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -61,6 +61,16 @@ if(PCMS_ENABLE_OMEGA_H) ${notRendezvous} ${cyclone2p}) endif() + add_exe(test_distributed_field) + if(HOST_NPROC GREATER_EQUAL 2) + mpi_test(test_distributed_field_2p 2 + $) + endif() + add_exe(test_distributed_coupling) + if(HOST_NPROC GREATER_EQUAL 2) + mpi_test(test_distributed_coupling_2p 2 + $) + endif() set(d3d1p ${PCMS_TEST_DATA_DIR}/d3d/d3d-full_9k_sfc.osh/) set(d3d1p_cpn ${PCMS_TEST_DATA_DIR}/d3d/1p.cpn) set(d3d2p_cpn ${PCMS_TEST_DATA_DIR}/d3d/2p.cpn) diff --git a/test/test_distributed_coupling.cpp b/test/test_distributed_coupling.cpp new file mode 100644 index 00000000..df8b8e5e --- /dev/null +++ b/test/test_distributed_coupling.cpp @@ -0,0 +1,142 @@ +#include +#include +#include +#include +#include + +#include "pcms/coupler/field_serializer.h" +#include "pcms/field/data/simple.h" +#include "pcms/field/function_space/lagrange.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/types.h" + +#include +#include +#include + +using pcms::LO; +using pcms::Real; + +namespace +{ +void require(bool cond, const char* msg) +{ + if (!cond) { + int rank = -1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + std::fprintf(stderr, "[rank %d] distributed coupling check failed: %s\n", + rank, msg); + MPI_Abort(MPI_COMM_WORLD, 1); + } +} +} // namespace + +// Exercises the owned-only serialization contract used by distributed coupling: +// only owned (rank-exclusive) DOF holders are put on the wire, and received +// values are scattered back into the local (owned + ghost) field. +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + int result = 0; + { + Kokkos::ScopeGuard kokkos{argc, argv}; + Omega_h::Library lib(&argc, &argv); + auto world = lib.world(); + const int rank = world->rank(); + const int nproc = world->size(); + require(nproc >= 2, "test requires at least 2 MPI ranks"); + + // Distributed 2D simplex box with a ghost layer (owned < local). + Omega_h::Mesh mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, + 0.0, /*nx=*/8, /*ny=*/8, /*nz=*/0, + /*symmetric=*/false); + mesh.set_parting(OMEGA_H_GHOSTED, 1, false); + + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/1, /*num_components=*/1, + pcms::CoordinateSystem::Cartesian, "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + const pcms::FieldLayout& layout = *space->GetLayout(); + + const LO n_owned = layout.GetNumOwnedDofHolder(); + const LO n_local = layout.GetNumLocalDofHolder(); + const LO n_comp = layout.GetNumComponents(); + const auto owned_to_local = layout.GetOwnedToLocalHost(); + require(n_owned > 0, "rank owns at least one DOF holder"); + require(n_local > n_owned, "ghosted mesh has local > owned"); + + // Set local field data: value = local index. + auto field = space->CreateFunction(); + std::vector local_data(static_cast(n_local)); + for (LO i = 0; i < n_local; ++i) { + local_data[static_cast(i)] = static_cast(i); + } + field.SetDOFHolderDataHost( + pcms::Rank2View(local_data.data(), + n_local, n_comp)); + + // Owned-indexed permutation: the last owned holder is "outside the overlap + // region" (perm = -1); the rest get compact buffer slots + // 0..n_participating. + const LO n_participating = n_owned - 1; + std::vector permutation(static_cast(n_owned), -1); + for (LO o = 0; o < n_participating; ++o) { + permutation[static_cast(o)] = o; + } + + // Serialize: only participating owned data is written to the buffer. + pcms::FieldSerializer serializer; + std::vector buffer(static_cast(n_participating) * n_comp); + const int sent = serializer.Serialize( + field.GetData(), layout, pcms::make_array_view(buffer), + pcms::make_const_array_view(permutation)); + require(sent == static_cast(n_participating) * n_comp, + "serialized size matches participating holders"); + + auto owned_field = + pcms::FlattenToRank1View(field.GetOwnedDOFHolderDataHost()); + require(owned_field.size() == static_cast(n_owned), + "owned field size"); + for (LO o = 0; o < n_participating; ++o) { + require(buffer[static_cast(o)] == + owned_field[static_cast(o)], + "serialized value == owned value"); + } + + // Deserialize into a fresh field: participating owned values land at their + // local positions; ghost and non-participating owned slots stay zero. + auto target = space->CreateFunction(); + serializer.Deserialize(target.GetData(), layout, + pcms::make_const_array_view(buffer), + pcms::make_const_array_view(permutation)); + auto target_data = pcms::FlattenToRank1View(target.GetDOFHolderDataHost()); + require(target_data.size() == static_cast(n_local), + "target local size"); + + for (LO o = 0; o < n_participating; ++o) { + const LO local = owned_to_local(o); + require(target_data[static_cast(local)] == + local_data[static_cast(local)], + "participating owned value restored"); + } + + const LO non_participating_local = owned_to_local(n_owned - 1); + require(target_data[static_cast(non_participating_local)] == + Real(0), + "non-participating owned holder is zero"); + + auto owned_mask = layout.GetOwnedHost(); + for (LO i = 0; i < n_local; ++i) { + if (!owned_mask[static_cast(i)]) { + require(target_data[static_cast(i)] == Real(0), + "ghost holder is zero"); + } + } + + if (rank == 0) { + std::printf("distributed coupling test passed (nproc=%d)\n", nproc); + } + } + MPI_Finalize(); + return result; +} diff --git a/test/test_distributed_field.cpp b/test/test_distributed_field.cpp new file mode 100644 index 00000000..051da055 --- /dev/null +++ b/test/test_distributed_field.cpp @@ -0,0 +1,213 @@ +#include +#include +#include +#include +#include + +#include "pcms/field/data/simple.h" +#include "pcms/field/function_space/lagrange.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/types.h" + +#include +#include +#include +#include + +using pcms::GO; +using pcms::LO; +using pcms::Real; + +namespace +{ +void require(bool cond, const char* msg) +{ + if (!cond) { + int rank = -1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + std::fprintf(stderr, "[rank %d] distributed field check failed: %s\n", rank, + msg); + MPI_Abort(MPI_COMM_WORLD, 1); + } +} + +// Copy a device Rank2View (mdspan-like) to a host Kokkos view, staging through +// a same-layout device scratch so host/device layout mismatches are handled. +template +Kokkos::View copy_coords_to_host( + const CoordView& coords, int n, int dim) +{ + Kokkos::View host("coords_host", n, dim); + auto device = Kokkos::create_mirror_view(pcms::DeviceMemorySpace(), host); + Kokkos::parallel_for( + n, KOKKOS_LAMBDA(int i) { + for (int d = 0; d < dim; ++d) + device(i, d) = coords(i, d); + }); + Kokkos::deep_copy(host, device); + return host; +} +} // namespace + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + int result = 0; + { + Kokkos::ScopeGuard kokkos{argc, argv}; + Omega_h::Library lib(&argc, &argv); + auto world = lib.world(); + const int rank = world->rank(); + const int nproc = world->size(); + require(nproc >= 2, "test requires at least 2 MPI ranks"); + + // Distributed 2D simplex box with a ghost layer so each rank has owned and + // ghost (non-owned) vertices. + Omega_h::Mesh mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, + 0.0, /*nx=*/8, /*ny=*/8, /*nz=*/0, + /*symmetric=*/false); + mesh.set_parting(OMEGA_H_GHOSTED, 1, false); + + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/1, /*num_components=*/1, + pcms::CoordinateSystem::Cartesian, "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + const pcms::FieldLayout& layout = *space->GetLayout(); + + const LO n_owned = layout.GetNumOwnedDofHolder(); + const LO n_local = layout.GetNumLocalDofHolder(); + + require(layout.IsDistributed(), "OmegaH layout should be distributed"); + // Count owned vertices from the owned mask on host. Avoid nents_owned(), + // which reads the device-owned array from host in some Omega_h/CUDA builds. + auto owned_mask_omega = Omega_h::HostRead(mesh.owned(0)); + LO n_owned_omega = 0; + for (LO i = 0; i < mesh.nents(0); ++i) { + if (owned_mask_omega[i] != 0) + ++n_owned_omega; + } + require(n_owned == n_owned_omega, "owned count != owned mask count"); + require(n_local == mesh.nents(0), "local count != nents(0)"); + require(layout.GetNumGlobalDofHolder() == mesh.nglobal_ents(0), + "global count != nglobal_ents(0)"); + + // Total owned vertices < total local (owned + ghost) vertices. + GO g_owned = 0, g_local = 0; + const GO owned64 = static_cast(n_owned); + const GO local64 = static_cast(n_local); + MPI_Allreduce(&owned64, &g_owned, 1, MPI_INT64_T, MPI_SUM, MPI_COMM_WORLD); + MPI_Allreduce(&local64, &g_local, 1, MPI_INT64_T, MPI_SUM, MPI_COMM_WORLD); + std::cout << "g_owned: " << g_owned << ", g_local: " << g_local + << std::endl; + require(g_owned < g_local, "global owned !< global local (no ghosts?)"); + + // Owned accessors are sized by the owned count. + auto owned_gids = layout.GetOwnedGidsHost(); + auto local_gids = layout.GetGidsHost(); + auto owned_to_local = layout.GetOwnedToLocalHost(); + std::cout << "owned_gids.size(): " << owned_gids.size() + << ", local_gids.size(): " << local_gids.size() + << ", owned_to_local.size(): " << owned_to_local.size() + << std::endl; + require(owned_gids.size() == static_cast(n_owned), + "owned gids size"); + require(local_gids.size() == static_cast(n_local), + "local gids size"); + require(owned_to_local.size() == static_cast(n_owned), + "owned_to_local size"); + + auto owned_coords = layout.GetOwnedDOFHolderCoordinates().GetValues(); + auto local_coords = layout.GetDOFHolderCoordinates().GetValues(); + std::cout << "owned_coords.extent(0): " << owned_coords.extent(0) + << ", local_coords.extent(0): " << local_coords.extent(0) + << std::endl; + require(owned_coords.extent(0) == static_cast(n_owned), + "owned coords rows"); + require(local_coords.extent(0) == static_cast(n_local), + "local coords rows"); + + // owned_to_local maps each owned holder to an owned local holder, and the + // owned GID equals the local GID at that index. + auto owned_mask = layout.GetOwnedHost(); + require(owned_mask.size() == static_cast(n_local), + "owned mask size"); + for (LO o = 0; o < n_owned; ++o) { + const LO local = owned_to_local(o); + require(local >= 0 && local < n_local, "owned_to_local out of range"); + require(owned_mask[static_cast(local)], + "owned_to_local maps to a non-owned holder"); + require(owned_gids(o) == local_gids(local), "owned gid != local gid"); + } + + // Owned coordinates equal the local coordinates at the mapped index. + auto owned_coords_h = + copy_coords_to_host(owned_coords, static_cast(n_owned), mesh.dim()); + auto local_coords_h = + copy_coords_to_host(local_coords, static_cast(n_local), mesh.dim()); + for (LO o = 0; o < n_owned; ++o) { + const LO local = owned_to_local(o); + for (int d = 0; d < mesh.dim(); ++d) { + require(owned_coords_h(o, d) == local_coords_h(local, d), + "owned coord != local coord"); + } + } + + // Field data: the owned data view is the owned subset of the local view. + auto field = space->CreateFunction(); + std::vector local_data(static_cast(n_local)); + for (LO i = 0; i < n_local; ++i) { + local_data[static_cast(i)] = static_cast(i); + } + field.SetDOFHolderDataHost( + pcms::Rank2View(local_data.data(), + n_local, 1)); + + auto local_field = pcms::FlattenToRank1View(field.GetDOFHolderDataHost()); + auto owned_field = + pcms::FlattenToRank1View(field.GetOwnedDOFHolderDataHost()); + std::cout << "local_field.size(): " << local_field.size() + << ", owned_field.size(): " << owned_field.size() << std::endl; + require(local_field.size() == static_cast(n_local), + "local field size"); + require(owned_field.size() == static_cast(n_owned), + "owned field size"); + for (LO o = 0; o < n_owned; ++o) { + const LO local = owned_to_local(o); + require(owned_field[static_cast(o)] == + local_field[static_cast(local)], + "owned field value != local field value"); + } + + // Ghost synchronization: set each rank's owned values to its rank id, + // synchronize, and verify every ghost value matches its owner's rank id. + auto sync_field = space->CreateFunction(); + std::vector sync_vals(static_cast(n_local), Real(0)); + for (LO i = 0; i < n_local; ++i) { + if (owned_mask[static_cast(i)]) { + sync_vals[static_cast(i)] = static_cast(rank); + } + } + sync_field.SetDOFHolderDataHost( + pcms::Rank2View(sync_vals.data(), + n_local, 1)); + sync_field.SynchronizeGhosts(); + + auto remotes = mesh.ask_owners(0); + auto owners = Omega_h::HostRead(remotes.ranks); + auto sync_data = + pcms::FlattenToRank1View(sync_field.GetDOFHolderDataHost()); + for (LO i = 0; i < n_local; ++i) { + if (!owned_mask[static_cast(i)]) { + require(sync_data[static_cast(i)] == + static_cast(owners[static_cast(i)]), + "ghost value != owner rank after sync"); + } + } + + if (rank == 0) { + std::printf("distributed field test passed (nproc=%d)\n", nproc); + } + } + MPI_Finalize(); + return result; +} diff --git a/test/test_field_communication.cpp b/test/test_field_communication.cpp index dcee3252..d17f2759 100644 --- a/test/test_field_communication.cpp +++ b/test/test_field_communication.cpp @@ -123,7 +123,7 @@ void client1(MPI_Comm comm, Omega_h::Mesh& mesh, std::string comm_name, mesh, order, 1, pcms::CoordinateSystem::Cartesian); auto layout = factory->GetLayout(); auto gids = layout->GetGidsHost(); - const auto n = layout->GetNumOwnedDofHolder(); + const auto n = layout->GetNumLocalDofHolder(); Omega_h::HostWrite ids(n); PCMS_ALWAYS_ASSERT(n == gids.size()); Kokkos::parallel_for( @@ -155,7 +155,7 @@ void client2(MPI_Comm comm, Omega_h::Mesh& mesh, std::string comm_name, mesh, order, 1, pcms::CoordinateSystem::Cartesian); auto layout = factory->GetLayout(); auto gids = layout->GetGidsHost(); - const auto n = layout->GetNumOwnedDofHolder(); + const auto n = layout->GetNumLocalDofHolder(); auto field = factory->CreateFunction(); pcms::FieldLayoutCommunicator layout_comm(comm_name + "2", comm, rdv, channel, diff --git a/test/test_proxy_coupling.cpp b/test/test_proxy_coupling.cpp index 6e166ad6..6989edcd 100644 --- a/test/test_proxy_coupling.cpp +++ b/test/test_proxy_coupling.cpp @@ -26,7 +26,7 @@ void initializeFieldWithGids(const pcms::FieldLayout& layout, pcms::Real multiplier = 1.0) { auto gids = layout.GetGidsHost(); - const auto n = layout.GetNumOwnedDofHolder(); + const auto n = layout.GetNumLocalDofHolder(); Omega_h::HostWrite ids(n); PCMS_ALWAYS_ASSERT(n == gids.size()); @@ -47,7 +47,7 @@ bool validateField(const pcms::FieldLayout& layout, auto gids = layout.GetGidsHost(); auto copied_array = pcms::FlattenToRank1View(field->GetDOFHolderDataHost()); auto owned = layout.GetOwnedHost(); - const auto n = layout.GetNumOwnedDofHolder(); + const auto n = layout.GetNumLocalDofHolder(); PCMS_ALWAYS_ASSERT(copied_array.size() == gids.size()); PCMS_ALWAYS_ASSERT(owned.size() == gids.size()); diff --git a/test/testing.cmake b/test/testing.cmake index 12590a57..6d461410 100644 --- a/test/testing.cmake +++ b/test/testing.cmake @@ -74,12 +74,12 @@ function(mpi_test TESTNAME PROCS EXE) if(NOT ${MPIEXEC_PREFLAGS} MATCHES "none") list(APPEND TEST_COMMAND ${MPIEXEC_PREFLAGS}) endif() - list(APPEND TEST_COMMAND ARGS ${MPIEXEC_NUMPROC_FLAG} ${PROCS}) + list(APPEND TEST_COMMAND ${MPIEXEC_NUMPROC_FLAG} ${PROCS}) if(NOT ${VALGRIND_EXECUTABLE} MATCHES "none") list(APPEND TEST_COMMAND ${VALGRIND_EXECUTABLE} ${VALGRIND_ARGS}) endif() list(APPEND TEST_COMMAND ${EXE} ${ARGN}) - add_test(NAME ${TESTNAME} COMMAND ${TEST_COMMAND }) + add_test(NAME ${TESTNAME} COMMAND ${TEST_COMMAND}) endfunction(mpi_test) function(dual_mpi_test)