From d495d78338e50efaf9a777810e398793993b19d8 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Thu, 24 Sep 2026 18:20:28 -0400 Subject: [PATCH 1/7] 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 | 453 ++++++++++++++++++ test/test_distributed_field.cpp | 233 +++++++++ test/test_field_communication.cpp | 4 +- test/test_proxy_coupling.cpp | 4 +- test/testing.cmake | 4 +- 38 files changed, 1412 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 88f60d2db..d558287dd 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 00dabfb72..135ebbe98 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 723929cf5..1918c1d76 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 db795b45f..e3b7dcb69 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 e26583775..a4086bac9 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 4c27a49d3..7da15545b 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 6deda15ae..b5cd7be8a 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 ed0614585..7a8d3208e 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 6da4261d4..97c0bc6b4 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 fb3846cd0..a4e0a3bf2 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 7903c72e3..d29cc529e 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 b4551e4c1..2e6fbd348 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 83fef4858..2e7bfd03d 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 991d58ff8..25f9beb15 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 fa0866623..01d2e6930 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 1c44760c6..39040c16a 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 32f390437..e4c62dc38 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 4c60545ae..04c5b28c7 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 c9d7bdb9b..8d83b8f8e 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 68154979e..b9ac5512d 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 ffa9ca4d7..12e232ac1 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 50e6e9874..0d14cbad0 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 08f4823d1..cc169bf7c 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 f4a78a854..f40068a2d 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 18729189f..1a3ebd59c 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 a755dcb0b..9db9dd821 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 bce424394..e9683683d 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 eca1c98a1..bf1a67f2e 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 31b1b9e74..be10bd385 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 742b86849..332b96961 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 4538474a2..c44c0f6e8 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 9b76d97a3..63d8ebd96 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 b51395eb0..86e25fb22 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 000000000..c561239c2 --- /dev/null +++ b/test/test_distributed_coupling.cpp @@ -0,0 +1,453 @@ +#include +#include +#include +#include +#include + +#include "pcms/coupler/field_exchange_planner.h" +#include "pcms/coupler/field_serializer.h" +#include "pcms/coupler/overlap_mask.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 +#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 coupling check failed: %s\n", + rank, msg); + MPI_Abort(MPI_COMM_WORLD, 1); + } +} + +// Builds a one-block GID message advertising `gids`, i.e. the structure the +// exchange planner produces for a single destination. +std::vector MakeGidMessage(const std::vector& gids) +{ + require(!gids.empty(), "cannot build an empty GID message"); + std::vector msg( + static_cast(pcms::ent_offsets_len) + gids.size(), 0); + // Every DOF holder in these tests sits on a mesh vertex (entity dim 0), so + // the per-entity-dimension header counts equal the total GID count from slot + // one onwards. + for (int e = 1; e < pcms::ent_offsets_len; ++e) { + msg[static_cast(e)] = static_cast(gids.size()); + } + std::copy(gids.begin(), gids.end(), msg.begin() + pcms::ent_offsets_len); + return msg; +} + +// In-message layout for a one-block message (one sender group) received by +// `rank`. The planner reads only the entries describing this rank's segment of +// the receive buffer: its start and its end. +redev::InMessageLayout MakeInMessageLayout(int rank, int nproc, size_t msg_size) +{ + redev::InMessageLayout layout; + layout.srcRanks.assign(static_cast(nproc), 0); + layout.offset.assign(static_cast(nproc) + 1, 0); + layout.offset[static_cast(rank)] = 0; + layout.offset[static_cast(rank) + 1] = + static_cast(msg_size); + return layout; +} + +// --------------------------------------------------------------------------- +// Checks +// --------------------------------------------------------------------------- + +// Shared state for the distributed coupling checks: the (ghosted) mesh, its +// order-1 Lagrange space and the rank counts. Everything else the checks need +// (counts, GIDs, masks) is derived from the layout. +struct DistributedCouplingContext +{ + Omega_h::Mesh& mesh; + std::shared_ptr space; + int rank; + int nproc; + + const pcms::FieldLayout& Layout() const { return *space->GetLayout(); } + LO NumOwned() const { return Layout().GetNumOwnedDofHolder(); } + LO NumLocal() const { return Layout().GetNumLocalDofHolder(); } + int NumComponents() const { return Layout().GetNumComponents(); } +}; + +// 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. +void check_serializer_round_trip(const DistributedCouplingContext& ctx) +{ + const pcms::FieldLayout& layout = ctx.Layout(); + const LO n_owned = ctx.NumOwned(); + const LO n_local = ctx.NumLocal(); + const LO n_comp = ctx.NumComponents(); + const auto owned_to_local = layout.GetOwnedToLocalHost(); + require(n_owned > 0, "rank owns at least one DOF holder"); + require(n_owned >= 2, "rank owns at least two DOF holders"); + require(n_local > n_owned, "ghosted mesh has local > owned"); + + // Set local field data: value = local index. + auto field = ctx.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 = ctx.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"); + } + } +} + +// Exercises the distributed exchange end to end: an ExchangePlan packed into a +// GID message, GID matching on the receive side, and the payload round trip. +// The redev transport is replaced by MPI collectives (Alltoall/Alltoallv), and +// the exchange partners are the ranks that actually own this rank's ghost DOF +// holders, so the check holds for any partitioning and rank count. +void check_distributed_exchange(const DistributedCouplingContext& ctx) +{ + Omega_h::Mesh& mesh = ctx.mesh; + const pcms::FieldLayout& layout = ctx.Layout(); + const int rank = ctx.rank; + const int nproc = ctx.nproc; + const LO n_owned = ctx.NumOwned(); + const LO n_local = ctx.NumLocal(); + const LO n_comp = ctx.NumComponents(); + const auto owned_to_local = layout.GetOwnedToLocalHost(); + const auto owned_mask = layout.GetOwnedHost(); + auto owned_gids = layout.GetOwnedGidsHost(); + auto local_gids = layout.GetGidsHost(); + auto owner_remotes = mesh.ask_owners(0); + auto owner_ranks = Omega_h::HostRead(owner_remotes.ranks); + + // The real exchange partners: this rank's ghost DOF holders, grouped by the + // rank that owns them. A ghost owned by rank n is a DOF this rank would have + // to obtain from n in a coupled exchange. + std::vector neighbor_set; + std::vector> ghosts_from(static_cast(nproc)); + { + std::set seen; + for (LO i = 0; i < n_local; ++i) { + if (!owned_mask[static_cast(i)]) { + const int owner = static_cast(owner_ranks[static_cast(i)]); + ghosts_from[static_cast(owner)].push_back( + static_cast(local_gids(i))); + seen.insert(owner); + } + } + neighbor_set.assign(seen.begin(), seen.end()); + } + require(!neighbor_set.empty(), + "the ghosted mesh should have at least one neighbour"); + + // Send-side plan: the planner is driven by a synthetic peer partition that + // routes every owned holder to one actual neighbour. The routing is synthetic + // (a single mesh has no peer partition to query), but it drives the same + // partition-based plan construction a coupled application uses. + const int destination = neighbor_set.front(); + std::vector ptn_ranks = {destination, rank}; + std::vector ptn_cuts = {0.0, 2.0}; // heap layout: cuts[1] is the cut + redev::RCBPtn peer_ptn(2, ptn_ranks, ptn_cuts); + redev::Partition exchange_partition{peer_ptn}; + + // Every owned holder participates; the unmapped (-1) path is exercised on the + // receive side, where the advertised list is a strict subset of this rank's + // owned holders. + pcms::OverlapMask overlap_mask(static_cast(n_owned)); + + pcms::GenericFieldExchangePlanner planner; + pcms::ExchangePlan send_plan = + planner.BuildExchangePlan(layout, exchange_partition, &overlap_mask); + + // Plan structure: a single destination (an actual neighbour) described by a + // valid CSR, and a permutation that compacts the owned holders in owned + // order. + require(send_plan.dest_ranks.size() == 1, + "send plan should route to a single destination"); + require(send_plan.dest_ranks[0] == destination, + "send plan destination should be an actual neighbour"); + require(send_plan.offsets.size() == 2, "send plan CSR size"); + require(send_plan.offsets[0] == 0, "send plan CSR start"); + require(send_plan.offsets[1] == static_cast(send_plan.msg_size), + "send plan CSR end != message size"); + require(send_plan.msg_size == static_cast(n_owned), + "every owned holder should participate"); + require(send_plan.permutation.size() == static_cast(n_owned), + "send plan permutation is owned-indexed"); + LO next_slot = 0; + for (LO o = 0; o < n_owned; ++o) { + require(send_plan.permutation[static_cast(o)] == next_slot, + "send plan permutation must compact owned holders in owned order"); + ++next_slot; + } + + // GID message: one block whose header counts the exchanged holders per mesh + // entity dimension and whose payload lists the owned GIDs. + const size_t gid_header = + send_plan.dest_ranks.size() * static_cast(pcms::ent_offsets_len); + std::vector gid_msg(send_plan.msg_size + gid_header, 0); + planner.FillGidMessage(layout, send_plan, + pcms::Rank1View( + gid_msg.data(), gid_msg.size())); + require(gid_msg[0] == 0, "GID message header must start at zero"); + for (int e = 1; e < pcms::ent_offsets_len; ++e) { + require(gid_msg[static_cast(e)] == + static_cast(send_plan.msg_size), + "GID message header must count the vertex holders"); + } + for (LO o = 0; o < n_owned; ++o) { + const auto slot = + static_cast(send_plan.permutation[static_cast(o)]); + require(gid_msg[gid_header + slot] == static_cast(owned_gids(o)), + "GID message payload != owned GIDs"); + } + + // Exchange each ghost GID with its owner. Two-phase all-to-all: first the + // per-neighbour counts, then the GID lists themselves. Afterwards ghost_recv + // holds, grouped by sender rank, every GID this rank owns that some neighbour + // admits holding -- i.e. the owned DOFs this rank must send values for. + std::vector send_counts(static_cast(nproc), 0); + std::vector recv_counts(static_cast(nproc), 0); + std::vector sdispls(static_cast(nproc), 0); + std::vector rdispls(static_cast(nproc), 0); + for (int n = 0; n < nproc; ++n) { + send_counts[static_cast(n)] = + static_cast(ghosts_from[static_cast(n)].size()); + } + MPI_Alltoall(send_counts.data(), 1, MPI_INT, recv_counts.data(), 1, MPI_INT, + MPI_COMM_WORLD); + std::vector ghost_send; + for (int n = 0; n < nproc; ++n) { + sdispls[static_cast(n)] = static_cast(ghost_send.size()); + ghost_send.insert(ghost_send.end(), + ghosts_from[static_cast(n)].begin(), + ghosts_from[static_cast(n)].end()); + } + LO total_recv = 0; + for (int n = 0; n < nproc; ++n) { + rdispls[static_cast(n)] = static_cast(total_recv); + total_recv += recv_counts[static_cast(n)]; + } + std::vector ghost_recv(static_cast(total_recv), 0); + MPI_Alltoallv(ghost_send.data(), send_counts.data(), sdispls.data(), + MPI_INT64_T, ghost_recv.data(), recv_counts.data(), + rdispls.data(), MPI_INT64_T, MPI_COMM_WORLD); + require(total_recv > 0, + "no neighbour admits holding any of this rank's owned DOFs"); + + // The received GID list is what the receive plan is built from: it gathers + // into a single block (one sender's message) every GID some neighbour admits + // holding, then matches those against the holders this rank owns. + const std::vector recv_gid_message = MakeGidMessage(ghost_recv); + pcms::ExchangePlan recv_plan = planner.BuildReceivePlan( + layout, + pcms::GlobalIDView(recv_gid_message.data(), + recv_gid_message.size()), + rank, nproc, MakeInMessageLayout(rank, nproc, recv_gid_message.size())); + + require(recv_plan.permutation.size() == static_cast(n_owned), + "receive plan permutation is owned-indexed"); + require(recv_plan.msg_size == static_cast(total_recv), + "receive plan payload length != advertised GID count"); + require(recv_plan.dest_ranks.size() == 1 && recv_plan.offsets.size() == 2, + "receive plan should describe a single sender"); + require(recv_plan.offsets[0] == 0 && + recv_plan.offsets[1] == static_cast(recv_plan.msg_size), + "receive plan CSR"); + // A rank's owned DOF can be a ghost on several neighbours (e.g. a corner + // vertex shared by four ranks), so the same GID may legitimately be + // advertised more than once. Matching only cares about the distinct set. + const std::set advertised_to_me(ghost_recv.begin(), ghost_recv.end()); + LO n_matched = 0; + for (LO o = 0; o < n_owned; ++o) { + const bool advertised = + advertised_to_me.count(static_cast(owned_gids(o))) > 0; + const LO slot = recv_plan.permutation[static_cast(o)]; + require((slot >= 0) == advertised, + "receive plan matched the wrong owned holders"); + if (advertised) { + ++n_matched; + require(slot < static_cast(recv_plan.msg_size), + "receive plan buffer index out of range"); + } + } + require(n_matched == static_cast(advertised_to_me.size()), + "receive plan must match every advertised GID against an owned " + "holder"); + require(n_matched < n_owned, + "every owned holder was advertised: the unmatched-permutation path " + "was not exercised"); + + // Local packing: Serialize this rank's owned values (the GIDs themselves) + // with the receive plan permutation and round-trip them through Deserialize + // to confirm each matched holder maps onto its slot (checked just below). The + // cross-rank broadcast after it sends the same values back to every + // advertiser via the transpose of the GID exchange. + auto exchange_field = ctx.space->CreateFunction(); + std::vector exchange_vals(static_cast(n_local), Real(0)); + for (LO o = 0; o < n_owned; ++o) { + exchange_vals[static_cast(owned_to_local(o))] = + static_cast(owned_gids(o)); + } + exchange_field.SetDOFHolderDataHost( + pcms::Rank2View(exchange_vals.data(), + n_local, n_comp)); + + pcms::FieldSerializer exchange_serializer; + std::vector payload(recv_plan.msg_size, Real(0)); + const int packed = exchange_serializer.Serialize( + exchange_field.GetData(), layout, pcms::make_array_view(payload), + pcms::make_const_array_view(recv_plan.permutation)); + require(packed == static_cast(recv_plan.msg_size), + "packed payload size != receive plan message size"); + + // Local check of the packing: a matched holder's owned value lands at its + // plan slot, and holders the plan did not match stay untouched. + auto echo_field = ctx.space->CreateFunction(); + exchange_serializer.Deserialize( + echo_field.GetData(), layout, pcms::make_const_array_view(payload), + pcms::make_const_array_view(recv_plan.permutation)); + auto echo_data = pcms::FlattenToRank1View(echo_field.GetDOFHolderDataHost()); + for (LO o = 0; o < n_owned; ++o) { + const LO local = owned_to_local(o); + if (recv_plan.permutation[static_cast(o)] >= 0) { + require(echo_data[static_cast(local)] == + static_cast(owned_gids(o)), + "packed value for a matched holder != its owned value"); + } else { + require(echo_data[static_cast(local)] == Real(0), + "unmatched holder must stay untouched by deserialization"); + } + } + for (LO i = 0; i < n_local; ++i) { + if (!owned_mask[static_cast(i)]) { + require(echo_data[static_cast(i)] == Real(0), + "deserialization must not write ghost holders"); + } + } + + // Cross-rank broadcast: the owner sends its value for every advertised GID + // back to each rank that advertised it, in that rank's own order. This is the + // same broadcast Omega_h's sync_array performs for ghost values -- it must + // serve every advertiser, because a corner DOF may be requested by several + // neighbours. + std::vector return_buf(ghost_recv.size(), Real(0)); + for (size_t k = 0; k < ghost_recv.size(); ++k) { + return_buf[k] = static_cast(ghost_recv[k]); + } + std::vector payload_recv(ghost_send.size(), Real(0)); + MPI_Alltoallv(return_buf.data(), recv_counts.data(), rdispls.data(), + MPI_DOUBLE, payload_recv.data(), send_counts.data(), + sdispls.data(), MPI_DOUBLE, MPI_COMM_WORLD); + for (size_t k = 0; k < ghost_send.size(); ++k) { + require(payload_recv[k] == static_cast(ghost_send[k]), + "values returned for ghost holders != the owning rank's values"); + } +} +} // namespace + +// Exercises the distributed coupling data path on a ghosted mesh: the +// owned-only serialization contract (check_serializer_round_trip) and the +// exchange plan, GID matching and payload round trip a coupler drives +// (check_distributed_exchange). +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + int result = 0; + { + 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"); + + const int n = 4 * nproc; + Omega_h::Mesh mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, + 0.0, /*nx=*/n, /*ny=*/n, /*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 DistributedCouplingContext ctx{mesh, space, rank, nproc}; + + check_serializer_round_trip(ctx); + + check_distributed_exchange(ctx); + + 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 000000000..5e8d40103 --- /dev/null +++ b/test/test_distributed_field.cpp @@ -0,0 +1,233 @@ +#include + +#include +#include + +#include "pcms/field/function_space/lagrange.h" +#include "pcms/field/out_of_bounds_policy.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/types.h" + +#include +#include +#include +#include + +using pcms::LO; +using pcms::Real; + +namespace +{ +constexpr Real kGhostSentinel = -1.0e30; +constexpr Real kFillValue = -2.0e30; + +// Sentinel used to pre-fill the evaluation buffer. It is deliberately different +// from kFillValue: the FILL policy writes kFillValue into every point that +// localization rejects, so if the buffer were pre-filled with the same value, a +// point the evaluator never wrote and a point it rejected would be +// indistinguishable. With a distinct sentinel the three outcomes separate: +// kUnresolved -> the evaluator never wrote this point (kernel/copy did not +// run) kFillValue -> localization rejected this point and the FILL policy +// wrote it anything else -> a real interpolated value +constexpr Real kUnresolved = -3.0e30; + +KOKKOS_INLINE_FUNCTION Real linear_f(Real x, Real y) +{ + return x + 2.0 * y; +} + +template +Kokkos::View copy_coordinates_to_host( + const CoordView& coordinates, LO n, int dim) +{ + using layout = std::conditional_t< + std::is_same_v, + Kokkos::LayoutLeft, Kokkos::LayoutRight>; + Kokkos::View + device_coordinates(coordinates.data_handle(), n, dim); + Kokkos::View host_coordinates("coordinates", n, + dim); + pcms::DeepCopyMismatchLayouts(host_coordinates, device_coordinates); + return host_coordinates; +} + +struct DeviceCoordinates +{ + Kokkos::View storage; + pcms::CoordinateView coordinate_view; +}; + +DeviceCoordinates make_device_coordinates(const std::vector& points) +{ + const LO n = static_cast(points.size() / 2); + Kokkos::View storage("coordinates", n, 2); + Kokkos::View host("coordinates", n, 2); + for (LO i = 0; i < n; ++i) { + host(i, 0) = points[static_cast(i) * 2]; + host(i, 1) = points[static_cast(i) * 2 + 1]; + } + pcms::DeepCopyMismatchLayouts(storage, host); + return {storage, + {pcms::CoordinateSystem::Cartesian, pcms::MakeRank2View(storage)}}; +} + +void require(bool condition, const char* message) +{ + if (!condition) { + int rank = -1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + std::fprintf(stderr, "[rank %d] %s\n", rank, message); + MPI_Abort(MPI_COMM_WORLD, 1); + } +} + +void check_distributed_evaluation(Omega_h::Mesh& mesh, + const pcms::LagrangeFunctionSpace& space) +{ + const auto& layout = *space.GetLayout(); + const LO n_local = layout.GetNumLocalDofHolder(); + int rank = -1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + + require(n_local > layout.GetNumOwnedDofHolder(), + "ghosted mesh did not create ghost DOF holders"); + + auto coords = copy_coordinates_to_host( + layout.GetDOFHolderCoordinates().GetValues(), n_local, mesh.dim()); + auto owned = layout.GetOwnedHost(); + std::vector values(static_cast(n_local), kGhostSentinel); + for (LO i = 0; i < n_local; ++i) { + if (owned[static_cast(i)]) { + values[static_cast(i)] = linear_f(coords(i, 0), coords(i, 1)); + } + } + + auto field = space.CreateFunction(); + field.SetDOFHolderDataHost(pcms::Rank2View( + values.data(), n_local, 1)); + field.SynchronizeGhosts(); + + auto data = pcms::FlattenToRank1View(field.GetDOFHolderDataHost()); + for (LO i = 0; i < n_local; ++i) { + require(std::fabs(data[static_cast(i)] - + linear_f(coords(i, 0), coords(i, 1))) < 1.0e-12, + "ghost synchronization did not produce the affine field"); + } + + const LO n_elements = mesh.nents(mesh.dim()); + const int vertices_per_element = mesh.dim() + 1; + auto element_vertices = Omega_h::HostRead(mesh.ask_elem_verts()); + auto mesh_coordinates = Omega_h::HostRead(mesh.coords()); + std::vector points(static_cast(n_elements) * mesh.dim()); + for (LO element = 0; element < n_elements; ++element) { + for (int dim = 0; dim < mesh.dim(); ++dim) { + Real centroid = 0.0; + for (int vertex = 0; vertex < vertices_per_element; ++vertex) { + const LO vertex_id = + element_vertices[element * vertices_per_element + vertex]; + centroid += mesh_coordinates[vertex_id * mesh.dim() + dim]; + } + points[static_cast(element) * mesh.dim() + dim] = + centroid / vertices_per_element; + } + } + auto device_coords = make_device_coordinates(points); + auto evaluator = + space.CreatePointEvaluator(pcms::EvaluationRequest::FromCoordinates( + device_coords.coordinate_view, + pcms::OutOfBoundsPolicy{pcms::OutOfBoundsMode::FILL, kFillValue})); + + const LO n_points = n_elements; + Kokkos::View result("result", n_points, 1); + Kokkos::deep_copy(result, kUnresolved); + evaluator->Evaluate(field, pcms::MakeRank2View(result)); + auto result_host = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), result); + + // A point still holding kUnresolved was never written by the evaluator (the + // evaluation kernels and/or the device-to-host copy produced nothing); a + // point holding kFillValue was written by the FILL policy, i.e. localization + // rejected it. They are counted separately because the causes differ. + LO n_unlocalized = 0; + LO n_untouched = 0; + LO first_unlocalized = -1; + LO first_untouched = -1; + for (LO i = 0; i < n_points; ++i) { + const Real value = result_host(i, 0); + if (value == kUnresolved) { + ++n_untouched; + if (first_untouched < 0) { + first_untouched = i; + } + continue; + } + if (value == kFillValue) { + ++n_unlocalized; + if (first_unlocalized < 0) { + first_unlocalized = i; + } + continue; + } + const size_t offset = static_cast(i) * 2; + const Real expected = linear_f(points[offset], points[offset + 1]); + require(std::fabs(value - expected) < 1.0e-12, + "distributed evaluation returned the wrong value"); + } + + if (n_untouched != 0) { + std::fprintf(stderr, + "[rank %d] %lld/%lld element centroid results were never " + "written by the evaluator (still at the pre-fill sentinel " + "%.17g, first index %lld): evaluation/copy failure, not a " + "localization failure\n", + rank, static_cast(n_untouched), + static_cast(n_points), + static_cast(kUnresolved), + static_cast(first_untouched)); + } + + if (n_unlocalized != 0) { + const size_t offset = static_cast(first_unlocalized) * 2; + std::fprintf(stderr, + "[rank %d] %lld/%lld element centroids were not localized " + "(first index %lld at point=(%.17g, %.17g))\n", + rank, static_cast(n_unlocalized), + static_cast(n_points), + static_cast(first_unlocalized), + static_cast(points[offset]), + static_cast(points[offset + 1])); + } + + // Distinct messages so the failure class is obvious from the abort line + // alone. + require(n_unlocalized == 0, "local element centroid was not localized"); + require(n_untouched == 0, + "element centroid evaluation result was never written"); + + if (rank == 0) + std::printf("distributed field test passed\n"); +} +} // namespace + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + { + Omega_h::Library library(&argc, &argv); + auto world = library.world(); + require(world->size() >= 2, "test requires at least 2 MPI ranks"); + + const int n = 4 * world->size(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, 0.0, n, n, 0, false); + mesh.set_parting(OMEGA_H_GHOSTED, 1, false); + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, 1, 1, pcms::CoordinateSystem::Cartesian, "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + + check_distributed_evaluation(mesh, *space); + } + MPI_Finalize(); + return 0; +} diff --git a/test/test_field_communication.cpp b/test/test_field_communication.cpp index dcee32526..d17f27593 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 6e166ad60..6989edcd0 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 12590a575..6d461410b 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) From 27a893a35f822467bee999c5052bc8c2bacbff9d Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Sun, 25 Oct 2026 19:27:47 -0400 Subject: [PATCH 2/7] define overlapping predicate --- src/pcms/localization/overlap.h | 426 +++++++++++++++++++++++++ src/pcms/localization/point_search.cpp | 401 ++++++++--------------- src/pcms/localization/point_search.h | 98 ++++++ 3 files changed, 661 insertions(+), 264 deletions(-) create mode 100644 src/pcms/localization/overlap.h diff --git a/src/pcms/localization/overlap.h b/src/pcms/localization/overlap.h new file mode 100644 index 000000000..081463325 --- /dev/null +++ b/src/pcms/localization/overlap.h @@ -0,0 +1,426 @@ +#ifndef PCMS_LOCALIZATION_OVERLAP_H +#define PCMS_LOCALIZATION_OVERLAP_H +#include +#include +#include +#include +#include +#include + +#include + +#include "pcms/utility/bounding_box.h" +#include "pcms/utility/types.h" +#include "pcms/utility/uniform_grid.h" + +namespace pcms +{ + +/// Predicate tags. Each tag names the operand pair with the exact type words +/// from the original function name (Triangle/Bbox, Simplex/Bbox, +/// Simplex/Simplex). A `_Method` suffix is added only when two tags would +/// otherwise collide. +/// +/// Adding a new predicate is a three-step, local change: define a tag here, +/// declare its padding if it bins by bounding box (see below), and add an +/// `Overlap` specialization. +struct TriangleVsBbox +{ + /// Slack in grid cells that `candidate_cells` must add around the box. + /// + /// The predicate accepts a shared boundary plane as an intersection, so a + /// cell that merely *touches* the box is a candidate. One cell is enough: a + /// cell two or more away cannot touch the box at all. + static constexpr LO boundary_pad = 1; +}; +struct SimplexVsBbox +{ + /// \see TriangleVsBbox::boundary_pad + static constexpr LO boundary_pad = 1; +}; +struct SimplexVsSimplex +{ + // No `boundary_pad` +}; + +/// Inclusive per-axis range of grid cells overlapped by a bounding box. +/// +/// The rules that define what a range *means* live here with the type rather +/// than in a consumer. Note that `CellRange{}` is not the empty range: its +/// zero initialization denotes the single cell at axis index 0, so use +/// `Empty()` when no cell should belong to it. +template +struct CellRange +{ + Kokkos::Array lo; + Kokkos::Array hi; + + /// The range no cell belongs to, as published by a rank with no geometry. + [[nodiscard]] static CellRange Empty() noexcept + { + CellRange range{}; + for (int d = 0; d < Dim; ++d) { + range.lo[d] = 0; + range.hi[d] = -1; + } + return range; + } + + /// True when no cell is in the range, so no query can match it. + [[nodiscard]] KOKKOS_INLINE_FUNCTION bool IsEmpty() const noexcept + { + for (int d = 0; d < Dim; ++d) { + if (hi[d] < lo[d]) { + return true; + } + } + return false; + } + + /// True when the cell at per-axis indices `axis_indices` falls in the range. + [[nodiscard]] KOKKOS_INLINE_FUNCTION bool Contains( + const Kokkos::Array& axis_indices) const noexcept + { + // An empty range has hi < lo on every axis, so this is false for it too. + for (int d = 0; d < Dim; ++d) { + if (axis_indices[d] < lo[d] || axis_indices[d] > hi[d]) { + return false; + } + } + return true; + } +}; + +namespace detail +{ + +KOKKOS_INLINE_FUNCTION +AABBox<2> triangle_bbox(const Omega_h::Matrix<2, 3>& coords) +{ + Kokkos::Array max{coords(0, 0), coords(1, 0)}; + Kokkos::Array min{coords(0, 0), coords(1, 0)}; + for (int i = 1; i < 3; ++i) { + max[0] = std::fmax(max[0], coords(0, i)); + max[1] = std::fmax(max[1], coords(1, i)); + min[0] = std::fmin(min[0], coords(0, i)); + min[1] = std::fmin(min[1], coords(1, i)); + } + return {.center = {(max[0] + min[0]) / 2.0, (max[1] + min[1]) / 2.0}, + .half_width = {(max[0] - min[0]) / 2.0, (max[1] - min[1]) / 2.0}}; +} + +template +KOKKOS_INLINE_FUNCTION AABBox simplex_bbox( + const Omega_h::Matrix& coords) +{ + Kokkos::Array max; + Kokkos::Array min; + for (int j = 0; j < dim; ++j) { + max[j] = coords(j, 0); + min[j] = coords(j, 0); + } + for (int i = 1; i < dim + 1; ++i) { + for (int j = 0; j < dim; ++j) { + max[j] = std::fmax(max[j], coords(j, i)); + min[j] = std::fmin(min[j], coords(j, i)); + } + } + + Kokkos::Array center; + Kokkos::Array half_width; + + for (int j = 0; j < dim; ++j) { + center[j] = (max[j] + min[j]) / 2.0; + half_width[j] = (max[j] - min[j]) / 2.0; + } + + return {.center = center, .half_width = half_width}; +} + +// Liang, You-Dong, and B. A. Barsky. "A New Concept and Method for Line +// Clipping." ACM Transactions on Graphics 3, no. 1 (January 1984): 1-22. +// https://doi.org/10.1145/357332.357333. +KOKKOS_INLINE_FUNCTION +bool clipt(Real p, Real q, Real& t0, Real& t1) +{ + if (p < 0) { + auto r = q / p; + if (r > t1) { + return false; + } + if (r > t0) { + t0 = r; + } + } else if (p > 0) { + auto r = q / p; + if (r < t0) { + return false; + } + if (r < t1) { + t1 = r; + } + } + // p==0 + else { + if (q < 0) { + return false; + } + } + return true; +} + +// Liang, You-Dong, and B. A. Barsky. "A New Concept and Method for Line +// Clipping." ACM Transactions on Graphics 3, no. 1 (January 1984): 1-22. +// https://doi.org/10.1145/357332.357333. +KOKKOS_INLINE_FUNCTION +bool line_intersects_bbox(const Omega_h::Vector<2>& p0, + const Omega_h::Vector<2>& p1, const AABBox<2>& bbox) +{ + Real t0 = 0; + Real t1 = 1; + auto xleft = bbox.center[0] - bbox.half_width[0]; + auto deltax = p1[0] - p0[0]; + if (clipt(-deltax, p0[0] - xleft, t0, t1)) { + auto xright = bbox.center[0] + bbox.half_width[0]; + if (clipt(deltax, xright - p0[0], t0, t1)) { + auto deltay = p1[1] - p0[1]; + auto ybottom = bbox.center[1] - bbox.half_width[1]; + if (clipt(-deltay, p0[1] - ybottom, t0, t1)) { + auto ytop = bbox.center[1] + bbox.half_width[1]; + if (clipt(deltay, ytop - p0[1], t0, t1)) { + // full liang-barksy algorithm computes new x-y coordinates here, + // but we just need to check intersections + return true; + } + } + } + } + return false; +} + +template +[[nodiscard]] KOKKOS_INLINE_FUNCTION bool within_bbox( + const Omega_h::Vector coord, const AABBox& bbox) noexcept +{ + for (int i = 0; i < dim; ++i) { + if (coord[i] < bbox.center[i] - bbox.half_width[i]) + return false; + if (coord[i] > bbox.center[i] + bbox.half_width[i]) + return false; + } + return true; +} + +template +[[nodiscard]] KOKKOS_INLINE_FUNCTION bool bbox_verts_within_simplex( + const AABBox& bbox, const Omega_h::Matrix& coords) +{ + // each dimension has a pair of opposing "walls" + // 2D: { [left, right], [top, bottom] } -> { left, right, top, bottom } + // 3D: { [left, right], [top, bottom], [front, back] } -> { left, ..., back } + Kokkos::Array bbox_walls{}; + for (int i = 0; i < dim; i++) { + bbox_walls[i * 2] = bbox.center[i] - bbox.half_width[i]; + bbox_walls[i * 2 + 1] = bbox.center[i] + bbox.half_width[i]; + } + + // 1 << dim == 2 ** dim == num vertices in ndim bounding box / hypercube + constexpr unsigned num_verts = 1 << dim; + + // Each vertex is a just a unique combination of walls + // eg [left, bottom] (2D) or [right, top, front] (3) + for (unsigned i = 0; i < num_verts; ++i) { + // conveniently, i acts a bit field representing the current combination + Omega_h::Vector vert; + for (unsigned j = 0; j < dim; ++j) { + // eg 110 = 6 = [left, top, back] + // eg 01 = 1 = [right, top] + vert[j] = (i >> j) & 1 ? bbox_walls[j * 2] : bbox_walls[j * 2 + 1]; + } + auto xi = Omega_h::barycentric_from_global(vert, coords); + if (Omega_h::is_barycentric_inside(xi)) { + return true; + } + } + return false; +} + +} // namespace detail + +/// An overlap predicate, selected as `Overlap` and invoked through +/// its static `apply()`. +/// +/// This is the single entry point for overlap tests: the primary template is +/// intentionally undefined, so each `Dim x Tag` combination that exists has an +/// explicit specialization whose `apply()` holds the original function body +/// verbatim, and asking for a combination that does not exist is a compile +/// error rather than a silent wrong-dimension call. +/// +/// `Tag` names the operand pair (see the tags above); `Dim` fixes the +/// dimension. Because the tag is a template argument rather than a value, a tag +/// that carries no data is never passed as a dummy parameter. +template +struct Overlap; + +template <> +struct Overlap<2, TriangleVsBbox> +{ + /** + * Check if a triangle element represented by 3 coordinates in two dimensions + * intersects with a bounding box + */ + [[nodiscard]] KOKKOS_FUNCTION static bool apply( + const Omega_h::Matrix<2, 3>& coords, const AABBox<2>& bbox) + { + // triangle and grid cell bounding box intersect + if (intersects(detail::triangle_bbox(coords), bbox)) { + // if any of the triangle verts inside of bbox + if (detail::within_bbox<2>(coords[0], bbox) || + detail::within_bbox<2>(coords[1], bbox) || + detail::within_bbox<2>(coords[2], bbox)) { + return true; + } + // if any of the bbox verts are within the triangle + if (detail::bbox_verts_within_simplex(bbox, coords)) { + return true; + } + // if any of the triangle's edges intersect with the bounding box + if (detail::line_intersects_bbox(coords[0], coords[1], bbox) || + detail::line_intersects_bbox(coords[1], coords[2], bbox) || + detail::line_intersects_bbox(coords[2], coords[0], bbox)) { + return true; + } + } + return false; + } +}; + +template +struct Overlap +{ + [[nodiscard]] KOKKOS_FUNCTION static bool apply( + const Omega_h::Matrix& coords, const AABBox& bbox) + { + return intersects(detail::simplex_bbox(coords), bbox); + // TODO: Add refined cases from triangle_intersects_bbox + } +}; + +template +struct Overlap +{ + /** + * @brief Conservative test for a zero-volume overlap between two simplices. + * + * If verticies of one simplex lies in the closed outer half-space of some + * face plane of the other, their intersection is contained in that plane and + * so has no volume. + * + * @return true when the overlap is provably degenerate and the clip can be + * skipped + */ + [[nodiscard]] OMEGA_H_INLINE static bool apply( + const r3d::Few, Dim + 1>& a, + const r3d::Few, Dim + 1>& b) + { + // The face normal below is a 2D perpendicular or a 3D cross product, so + // this predicate is 2D/3D only even though everything else about it reads as + // dimension-generic. Without this the higher-dimension case is silent: at + // Dim >= 4 `normal[Dim - 1]` is never written and is read as garbage in the + // loop below, and at Dim == 1 writing normal[1] and normal[2] runs off the + // end of the vector. + static_assert(Dim == 2 || Dim == 3, + "simplex-vs-simplex overlap is implemented for 2D and 3D only"); + + // Each iteration treats one simplex as the half-space owner. Both are + // needed: a face of `a` may separate the pair while no face of `b` does. + for (int side = 0; side < 2; ++side) { + const auto& owner = (side == 0) ? a : b; + const auto& other = (side == 0) ? b : a; + + // Face `f` is the face opposite vertex `f`. + for (int f = 0; f < Dim + 1; ++f) { + // Vertices spanning the face, in index order with `f` skipped. + int face_verts[Dim]; + int n_face_verts = 0; + for (int v = 0; v < Dim + 1; ++v) { + if (v != f) { + face_verts[n_face_verts++] = v; + } + } + const auto& origin = owner[face_verts[0]]; + + // Face normal, oriented to point away from the opposite vertex so the + // owner's interior lies on the negative side. + r3d::Vector normal; + if constexpr (Dim == 2) { + const auto edge = owner[face_verts[1]] - origin; + normal[0] = edge[1]; + normal[1] = -edge[0]; + } else { + const auto e1 = owner[face_verts[1]] - origin; + const auto e2 = owner[face_verts[2]] - origin; + normal[0] = e1[1] * e2[2] - e1[2] * e2[1]; + normal[1] = e1[2] * e2[0] - e1[0] * e2[2]; + normal[2] = e1[0] * e2[1] - e1[1] * e2[0]; + } + + Omega_h::Real opposite_side = 0.0; + Omega_h::Real normal_sq = 0.0; + for (int d = 0; d < Dim; ++d) { + opposite_side += normal[d] * (owner[f][d] - origin[d]); + normal_sq += normal[d] * normal[d]; + } + // A degenerate face gives no usable axis. + if (normal_sq == 0.0) { + continue; + } + const Omega_h::Real orientation = (opposite_side > 0.0) ? -1.0 : 1.0; + + bool all_outside = true; + for (int v = 0; v < Dim + 1 && all_outside; ++v) { + Omega_h::Real signed_distance = 0.0; + for (int d = 0; d < Dim; ++d) { + signed_distance += normal[d] * (other[v][d] - origin[d]); + } + all_outside = (orientation * signed_distance) >= 0.0; + } + if (all_outside) { + return true; + } + } + } + return false; + } +}; + +/// Low-level cell range: cells overlapped by `box` on `grid`, widened by `pad` +/// cells per side and clamped to the grid. +template +[[nodiscard]] KOKKOS_INLINE_FUNCTION CellRange overlapping_cells( + const AABBox& box, const UniformGrid& grid, LO pad = 1) +{ + CellRange range; + for (int d = 0; d < Dim; ++d) { + range.lo[d] = Kokkos::max( + grid.AxisCellIndex(d, box.center[d] - box.half_width[d]) - pad, 0); + range.hi[d] = Kokkos::min( + grid.AxisCellIndex(d, box.center[d] + box.half_width[d]) + pad, + grid.divisions[d] - 1); + } + return range; +} + +/// Cells on `grid` that could satisfy a `Pred` test within `bound`: the cells +/// `bound` spans, widened by the predicate's own `boundary_pad` so that a cell +/// which merely touches the bound is included. +template +[[nodiscard]] KOKKOS_INLINE_FUNCTION CellRange candidate_cells( + const AABBox& bound, const UniformGrid& grid) +{ + return overlapping_cells(bound, grid, Pred::boundary_pad); +} + +} // namespace pcms + +#endif // PCMS_LOCALIZATION_OVERLAP_H diff --git a/src/pcms/localization/point_search.cpp b/src/pcms/localization/point_search.cpp index da74f055a..14577c6b1 100644 --- a/src/pcms/localization/point_search.cpp +++ b/src/pcms/localization/point_search.cpp @@ -1,7 +1,10 @@ #include "point_search.h" +#include "pcms/localization/overlap.h" #include #include #include +#include +#include // From // https://en.wikipedia.org/wiki/Distance_from_a_point_to_a_line#Vector_formulation @@ -112,217 +115,63 @@ LO GetOwningElementId(Omega_h::Mesh& mesh, int mesh_dim, int entity_dim, return owner; } -KOKKOS_INLINE_FUNCTION -AABBox<2> triangle_bbox(const Omega_h::Matrix<2, 3>& coords) -{ - Kokkos::Array max{coords(0, 0), coords(1, 0)}; - Kokkos::Array min{coords(0, 0), coords(1, 0)}; - for (int i = 1; i < 3; ++i) { - max[0] = std::fmax(max[0], coords(0, i)); - max[1] = std::fmax(max[1], coords(1, i)); - min[0] = std::fmin(min[0], coords(0, i)); - min[1] = std::fmin(min[1], coords(1, i)); - } - return {.center = {(max[0] + min[0]) / 2.0, (max[1] + min[1]) / 2.0}, - .half_width = {(max[0] - min[0]) / 2.0, (max[1] - min[1]) / 2.0}}; -} - -template -KOKKOS_INLINE_FUNCTION AABBox simplex_bbox( - const Omega_h::Matrix& coords) -{ - Kokkos::Array max; - Kokkos::Array min; - for (int j = 0; j < dim; ++j) { - max[j] = coords(j, 0); - min[j] = coords(j, 0); - } - for (int i = 1; i < dim + 1; ++i) { - for (int j = 0; j < dim; ++j) { - max[j] = std::fmax(max[j], coords(j, i)); - min[j] = std::fmin(min[j], coords(j, i)); - } - } - - Kokkos::Array center; - Kokkos::Array half_width; - - for (int j = 0; j < dim; ++j) { - center[j] = (max[j] + min[j]) / 2.0; - half_width[j] = (max[j] - min[j]) / 2.0; - } - - return {.center = center, .half_width = half_width}; -} - -// Liang, You-Dong, and B. A. Barsky. “A New Concept and Method for Line -// Clipping.” ACM Transactions on Graphics 3, no. 1 (January 1984): 1–22. -// https://doi.org/10.1145/357332.357333. -KOKKOS_INLINE_FUNCTION -bool clipt(Real p, Real q, Real& t0, Real& t1) -{ - if (p < 0) { - auto r = q / p; - if (r > t1) { - return false; - } - if (r > t0) { - t0 = r; - } - } else if (p > 0) { - auto r = q / p; - if (r < t0) { - return false; - } - if (r < t1) { - t1 = r; - } - } - // p==0 - else { - if (q < 0) { - return false; - } - } - return true; -} -// Liang, You-Dong, and B. A. Barsky. “A New Concept and Method for Line -// Clipping.” ACM Transactions on Graphics 3, no. 1 (January 1984): 1–22. -// https://doi.org/10.1145/357332.357333. -KOKKOS_INLINE_FUNCTION -bool line_intersects_bbox(const Omega_h::Vector<2>& p0, - const Omega_h::Vector<2>& p1, const AABBox<2>& bbox) -{ - Real t0 = 0; - Real t1 = 1; - auto xleft = bbox.center[0] - bbox.half_width[0]; - auto deltax = p1[0] - p0[0]; - if (clipt(-deltax, p0[0] - xleft, t0, t1)) { - auto xright = bbox.center[0] + bbox.half_width[0]; - if (clipt(deltax, xright - p0[0], t0, t1)) { - auto deltay = p1[1] - p0[1]; - auto ybottom = bbox.center[1] - bbox.half_width[1]; - if (clipt(-deltay, p0[1] - ybottom, t0, t1)) { - auto ytop = bbox.center[1] + bbox.half_width[1]; - if (clipt(deltay, ytop - p0[1], t0, t1)) { - // full liang-barksy algorithm computes new x-y coordinates here, but - // we just need to check intersections - return true; - } - } - } - } - return false; -} - -template -[[nodiscard]] KOKKOS_INLINE_FUNCTION bool within_bbox( - const Omega_h::Vector coord, const AABBox& bbox) noexcept -{ - for (int i = 0; i < dim; ++i) { - if (coord[i] < bbox.center[i] - bbox.half_width[i]) - return false; - if (coord[i] > bbox.center[i] + bbox.half_width[i]) - return false; - } - return true; -} - -template -[[nodiscard]] KOKKOS_INLINE_FUNCTION bool bbox_verts_within_simplex( - const AABBox& bbox, const Omega_h::Matrix& coords) -{ - // each dimension has a pair of opposing "walls" - // 2D: { [left, right], [top, bottom] } -> { left, right, top, bottom } - // 3D: { [left, right], [top, bottom], [front, back] } -> { left, ..., back } - Kokkos::Array bbox_walls{}; - for (int i = 0; i < dim; i++) { - bbox_walls[i * 2] = bbox.center[i] - bbox.half_width[i]; - bbox_walls[i * 2 + 1] = bbox.center[i] + bbox.half_width[i]; - } - - // 1 << dim == 2 ** dim == num vertices in ndim bounding box / hypercube - constexpr unsigned num_verts = 1 << dim; - - // Each vertex is a just a unique combination of walls - // eg [left, bottom] (2D) or [right, top, front] (3) - for (unsigned i = 0; i < num_verts; ++i) { - // conveniently, i acts a bit field representing the current combination - Omega_h::Vector vert; - for (unsigned j = 0; j < dim; ++j) { - // eg 110 = 6 = [left, top, back] - // eg 01 = 1 = [right, top] - vert[j] = (i >> j) & 1 ? bbox_walls[j * 2] : bbox_walls[j * 2 + 1]; - } - auto xi = Omega_h::barycentric_from_global(vert, coords); - if (Omega_h::is_barycentric_inside(xi)) { - return true; - } - } - return false; -} - /** * Check if a triangle element represented by 3 coordinates in two dimensions - * intersects with a bounding box + * intersects with a bounding box. */ [[nodiscard]] KOKKOS_FUNCTION bool triangle_intersects_bbox( const Omega_h::Matrix<2, 3>& coords, const AABBox<2>& bbox) { - // triangle and grid cell bounding box intersect - if (intersects(triangle_bbox(coords), bbox)) { - // if any of the triangle verts inside of bbox - if (within_bbox<2>(coords[0], bbox) || within_bbox<2>(coords[1], bbox) || - within_bbox<2>(coords[2], bbox)) { - return true; - } - // if any of the bbox verts are within the triangle - if (bbox_verts_within_simplex(bbox, coords)) { - return true; - } - // if any of the triangle's edges intersect with the bounding box - if (line_intersects_bbox(coords[0], coords[1], bbox) || - line_intersects_bbox(coords[1], coords[2], bbox) || - line_intersects_bbox(coords[2], coords[0], bbox)) { - return true; - } - } - return false; -} - -template -[[nodiscard]] KOKKOS_FUNCTION bool simplex_intersects_bbox( - const Omega_h::Matrix& coords, const AABBox& bbox) -{ - return intersects(simplex_bbox(coords), bbox); - // TODO: Add refined cases from triangle_intersects_bbox + return Overlap<2, TriangleVsBbox>::apply(coords, bbox); } namespace detail { +// The candidate-map machinery. Declared in the header, defined here, and +// explicitly instantiated for the supported predicates at the end of this +// file -- see the note there. + /** - * Functor for constructing the mapping from grid cells to intersecting - * triangles. Each row of the resulting CSR structure represents a grid cell and - * each row entry corresponds to an ID of an element that intersects that grid - * cell. - * \Warning since this works on Omega_h meshes, we currently assume each element - * is a 2D simplex (triangle) - * \Warning since this uses Omega_h data which is only available in the - * "Default" Execution space, the should not be used in an alternative EXE space + * The bridge between "an element id" and "the geometry a predicate tests". + * + * This is the only shape-specific part of building a candidate map. Everything + * downstream -- the broad-phase binning and the predicate test -- works on + * whatever `operands()` returns, so supporting a different element shape is a + * new policy with these three members and nothing else changes: + * + * struct MyElements { + * static MyElements FromMesh(Omega_h::Mesh&); + * Operands operands(LO elem) const; // what Pred::apply takes + * AABBox bound(const Operands&) const; // a box containing it + * }; + * + * `bound` is allowed to be conservative: the predicate rejects what it + * over-includes, whereas a bound that is too tight silently drops candidates. */ -/// The bbox-overlap predicate for a simplex against one grid cell. 2D applies -/// the refined edge tests, 3D is a plain bbox overlap; both match what the -/// previous cell-major functors used, so candidate lists are unchanged. template -[[nodiscard]] KOKKOS_INLINE_FUNCTION bool simplex_overlaps_cell( - const Omega_h::Matrix& coords, const AABBox& bbox) +struct MeshSimplexElements { - if constexpr (Dim == 2) { - return triangle_intersects_bbox(coords, bbox); - } else { - return simplex_intersects_bbox(coords, bbox); + Omega_h::LOs elems2verts; + Omega_h::Reals coords; + + [[nodiscard]] static MeshSimplexElements FromMesh(Omega_h::Mesh& mesh) + { + return {mesh.ask_elem_verts(), mesh.coords()}; } -} + + [[nodiscard]] KOKKOS_INLINE_FUNCTION Omega_h::Matrix + operands(LO elem) const + { + const auto verts = Omega_h::gather_verts(elems2verts, elem); + return Omega_h::gather_vectors(coords, verts); + } + + [[nodiscard]] KOKKOS_INLINE_FUNCTION AABBox bound( + const Omega_h::Matrix& ops) const + { + return simplex_bbox(ops); + } +}; /** * Functor for the element -> grid cell map, the transpose of what the search @@ -332,19 +181,18 @@ template * testing every element against a cell, each element only visits the cells its * own bounding box spans. That makes construction O(nelems * cells_per_elem) * instead of O(num_cells * nelems). + * + * Operand-agnostic: `Elements` decides what an element *is* and how to bound it, + * `Pred` decides when it overlaps a cell. Neither is fixed here, so the same + * functor bins any shape a caller can supply a policy for. */ -template +template > struct ElemCellOverlapFunctor { - ElemCellOverlapFunctor(Omega_h::Mesh& mesh, + ElemCellOverlapFunctor(Elements elements, Kokkos::View[1]> grid) - : elems2verts_(mesh.ask_elem_verts()), coords_(mesh.coords()), grid_(grid) + : elements_(elements), grid_(grid) { - if (mesh.dim() != Dim) { - std::cerr << "ElemCellOverlapFunctor: mesh dimension does not match the " - "functor dimension\n"; - std::terminate(); - } } /// Two-pass functor. On the first pass we count the cells this element @@ -352,36 +200,16 @@ struct ElemCellOverlapFunctor KOKKOS_INLINE_FUNCTION LO operator()(LO elem, LO* fill) const { - const auto elem_verts = Omega_h::gather_verts(elems2verts_, elem); - const auto vertex_coords = - Omega_h::gather_vectors(coords_, elem_verts); const auto& grid = grid_(0); + const auto ops = elements_.operands(elem); - // Per-axis inclusive range of cells the element's bbox spans, widened by - // one cell on each side. The overlap predicate treats a shared boundary - // plane as an intersection, so a cell that only touches the bbox counts as - // a candidate and would be missed by the unwidened range -- which happens - // constantly on a structured mesh whose vertices land on grid lines. One - // cell is enough: a cell two or more away cannot touch the bbox at all. - // The predicate below rejects whatever the widening over-includes. - Kokkos::Array lo, hi, idx; - for (int d = 0; d < Dim; ++d) { - Real min = vertex_coords[0][d]; - Real max = vertex_coords[0][d]; - for (int v = 1; v < Dim + 1; ++v) { - min = Kokkos::min(min, vertex_coords[v][d]); - max = Kokkos::max(max, vertex_coords[v][d]); - } - lo[d] = Kokkos::max(grid.AxisCellIndex(d, min) - 1, 0); - hi[d] = - Kokkos::min(grid.AxisCellIndex(d, max) + 1, grid.divisions[d] - 1); - idx[d] = lo[d]; - } + const auto range = candidate_cells(elements_.bound(ops), grid); + Kokkos::Array idx = range.lo; LO num_overlaps = 0; while (true) { const LO cell = grid.CellIndexFromAxisIndices(idx); - if (simplex_overlaps_cell(vertex_coords, grid.GetCellBBOX(cell))) { + if (Overlap::apply(ops, grid.GetCellBBOX(cell))) { if (fill) { fill[num_overlaps] = cell; } @@ -390,10 +218,10 @@ struct ElemCellOverlapFunctor // odometer over the inclusive range [lo, hi] int d = 0; for (; d < Dim; ++d) { - if (++idx[d] <= hi[d]) { + if (++idx[d] <= range.hi[d]) { break; } - idx[d] = lo[d]; + idx[d] = range.lo[d]; } if (d == Dim) { break; @@ -403,8 +231,7 @@ struct ElemCellOverlapFunctor } private: - Omega_h::LOs elems2verts_; - Omega_h::Reals coords_; + Elements elements_; Kokkos::View[1]> grid_; }; @@ -416,18 +243,25 @@ struct ElemCellOverlapFunctor * in.numRows() (the element count), which both over-allocates the row map and * writes out of bounds once num_grid_cells exceeds the element count. */ -template +template Kokkos::Crs construct_intersection_map(Omega_h::Mesh& mesh, Kokkos::View[1]> grid, int num_grid_cells) { + if (mesh.dim() != Dim) { + std::cerr << "construct_intersection_map: mesh dimension does not match the " + "map dimension\n"; + std::terminate(); + } + using CrsT = Kokkos::Crs; const auto nelems = mesh.nelems(); CrsT elem2cells{}; - Kokkos::count_and_fill_crs(elem2cells, nelems, - ElemCellOverlapFunctor{mesh, grid}); + Kokkos::count_and_fill_crs( + elem2cells, nelems, + ElemCellOverlapFunctor{Elements::FromMesh(mesh), grid}); Kokkos::View counts("candidate_counts", num_grid_cells); Kokkos::parallel_for( @@ -483,15 +317,12 @@ construct_intersection_map(Omega_h::Mesh& mesh, * ground truth for what the candidate map should contain, so it is kept and * exposed as an oracle for the equivalence test. */ -template +template > struct ReferenceCellMajorFunctor { - ReferenceCellMajorFunctor(Omega_h::Mesh& mesh, - Kokkos::View[1]> grid) - : elems2verts_(mesh.ask_elem_verts()), - coords_(mesh.coords()), - grid_(grid), - nelems_(mesh.nelems()) + ReferenceCellMajorFunctor(Elements elements, + Kokkos::View[1]> grid, LO nelems) + : elements_(elements), grid_(grid), nelems_(nelems) { } @@ -501,11 +332,8 @@ struct ReferenceCellMajorFunctor const auto cell_bbox = grid_(0).GetCellBBOX(row); LO num_intersections = 0; for (LO elem = 0; elem < nelems_; ++elem) { - const auto elem_verts = - Omega_h::gather_verts(elems2verts_, elem); - const auto vertex_coords = - Omega_h::gather_vectors(coords_, elem_verts); - if (simplex_overlaps_cell(vertex_coords, cell_bbox)) { + const auto ops = elements_.operands(elem); + if (Overlap::apply(ops, cell_bbox)) { if (fill) { fill[num_intersections] = elem; } @@ -516,30 +344,31 @@ struct ReferenceCellMajorFunctor } private: - Omega_h::LOs elems2verts_; - Omega_h::Reals coords_; + Elements elements_; Kokkos::View[1]> grid_; LO nelems_; }; -template +template Kokkos::Crs construct_intersection_map_reference(Omega_h::Mesh& mesh, Kokkos::View[1]> grid, int num_grid_cells) { Kokkos::Crs intersection_map{}; - Kokkos::count_and_fill_crs(intersection_map, num_grid_cells, - ReferenceCellMajorFunctor{mesh, grid}); + Kokkos::count_and_fill_crs( + intersection_map, num_grid_cells, + ReferenceCellMajorFunctor{Elements::FromMesh(mesh), grid, + mesh.nelems()}); return intersection_map; } - Kokkos::Crs construct_intersection_map_reference_2d(Omega_h::Mesh& mesh, Kokkos::View grid, int num_grid_cells) { - return construct_intersection_map_reference<2>(mesh, grid, num_grid_cells); + return construct_intersection_map_reference<2, TriangleVsBbox>( + mesh, grid, num_grid_cells); } Kokkos::Crs @@ -547,17 +376,24 @@ construct_intersection_map_reference_3d(Omega_h::Mesh& mesh, Kokkos::View grid, int num_grid_cells) { - return construct_intersection_map_reference<3>(mesh, grid, num_grid_cells); + return construct_intersection_map_reference<3, SimplexVsBbox>( + mesh, grid, num_grid_cells); } // num_grid_cells should be result of grid.GetNumCells(), take as argument to // avoid extra copy of grid from gpu to cpu +// +// The 2D and 3D overloads name their predicate here rather than reading it from +// a shared mapping: a "2D refined candidate map" is what this function *is*, so +// the choice belongs to its definition and cannot leak into another call site. +// Callers needing a different predicate use the template directly. Kokkos::Crs construct_intersection_map_2d(Omega_h::Mesh& mesh, Kokkos::View grid, int num_grid_cells) { - return construct_intersection_map<2>(mesh, grid, num_grid_cells); + return construct_intersection_map<2, TriangleVsBbox>(mesh, grid, + num_grid_cells); } Kokkos::Crs @@ -565,8 +401,35 @@ construct_intersection_map_3d(Omega_h::Mesh& mesh, Kokkos::View grid, int num_grid_cells) { - return construct_intersection_map<3>(mesh, grid, num_grid_cells); + return construct_intersection_map<3, SimplexVsBbox>(mesh, grid, + num_grid_cells); } +// The supported predicates, instantiated here so that the search classes' +// predicate-taking constructors and the tests can use them from other +// translation units. This is the complete set: naming a predicate that is not +// listed below is a link error, which is the accepted cost of keeping the +// machinery out of the header. Add a line here when adding a predicate that +// callers outside this file need. +template Kokkos::Crs +construct_intersection_map<2, TriangleVsBbox>(Omega_h::Mesh&, + Kokkos::View[1]>, int); +template Kokkos::Crs +construct_intersection_map<2, SimplexVsBbox>(Omega_h::Mesh&, + Kokkos::View[1]>, int); +template Kokkos::Crs +construct_intersection_map<3, SimplexVsBbox>(Omega_h::Mesh&, + Kokkos::View[1]>, int); + +template Kokkos::Crs +construct_intersection_map_reference<2, TriangleVsBbox>( + Omega_h::Mesh&, Kokkos::View[1]>, int); +template Kokkos::Crs +construct_intersection_map_reference<2, SimplexVsBbox>( + Omega_h::Mesh&, Kokkos::View[1]>, int); +template Kokkos::Crs +construct_intersection_map_reference<3, SimplexVsBbox>( + Omega_h::Mesh&, Kokkos::View[1]>, int); + } // namespace detail template @@ -791,6 +654,13 @@ GridPointSearch2D::GridPointSearch2D(Omega_h::Mesh& mesh, LO Nx, LO Ny) GridPointSearch2D::GridPointSearch2D(Omega_h::Mesh& mesh, LO Nx, LO Ny, const PointSearchTolerances& tolerances) : PointLocalizationSearch(tolerances), mesh_(mesh) +{ + InitializeGeometry(mesh, Nx, Ny); + // This search's default predicate: the refined triangle-vs-box test. + InitializeMap(); +} + +void GridPointSearch2D::InitializeGeometry(Omega_h::Mesh& mesh, LO Nx, LO Ny) { auto mesh_bbox = Omega_h::get_bounding_box<2>(&mesh); auto grid_h = Kokkos::create_mirror_view(grid_); @@ -800,11 +670,7 @@ GridPointSearch2D::GridPointSearch2D(Omega_h::Mesh& mesh, LO Nx, LO Ny, .bot_left = {mesh_bbox.min[0], mesh_bbox.min[1]}, .divisions = {Nx, Ny}}; Kokkos::deep_copy(grid_, grid_h); - // Determine inflation radius from tolerances (max of vertex/edge tol) - auto tol_h = Kokkos::create_mirror_view(tolerances_); - Kokkos::deep_copy(tol_h, tolerances_); - candidate_map_ = - detail::construct_intersection_map_2d(mesh, grid_, grid_h(0).GetNumCells()); + num_grid_cells_ = grid_h(0).GetNumCells(); coords_ = mesh.coords(); tris2verts_ = mesh.ask_elem_verts(); tris2edges_adj_ = mesh.ask_down(Omega_h::FACE, Omega_h::EDGE); @@ -932,6 +798,14 @@ GridPointSearch3D::GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz) GridPointSearch3D::GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz, const PointSearchTolerances& tolerances) : PointLocalizationSearch(tolerances), mesh_(mesh) +{ + InitializeGeometry(mesh, Nx, Ny, Nz); + // This search's default predicate: the conservative bounding-box test. + InitializeMap(); +} + +void GridPointSearch3D::InitializeGeometry(Omega_h::Mesh& mesh, LO Nx, LO Ny, + LO Nz) { auto mesh_bbox = Omega_h::get_bounding_box<3>(&mesh); auto grid_h = Kokkos::create_mirror_view(grid_); @@ -949,8 +823,7 @@ GridPointSearch3D::GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz, .divisions = {Nx, Ny, Nz}}; Kokkos::deep_copy(grid_, grid_h); - candidate_map_ = - detail::construct_intersection_map_3d(mesh, grid_, grid_h(0).GetNumCells()); + num_grid_cells_ = grid_h(0).GetNumCells(); coords_ = mesh.coords(); tris2verts_ = mesh.ask_elem_verts(); tris2edges_adj_ = mesh.ask_down(Omega_h::FACE, Omega_h::EDGE); diff --git a/src/pcms/localization/point_search.h b/src/pcms/localization/point_search.h index 602c99fb3..74edd20f3 100644 --- a/src/pcms/localization/point_search.h +++ b/src/pcms/localization/point_search.h @@ -5,10 +5,12 @@ #include #include +#include #include #include #include +#include "pcms/localization/overlap.h" #include "pcms/utility/types.h" #include "pcms/utility/uniform_grid.h" #include "pcms/utility/bounding_box.h" @@ -81,6 +83,33 @@ Kokkos::Crs construct_intersection_map_reference_3d(Omega_h::Mesh& mesh, Kokkos::View grid, int num_grid_cells); + +/// How the candidate map learns an element's geometry, and a conservative bound +/// on it. The default policy for an Omega_h simplex mesh; defined in the source +/// file alongside the rest of the map machinery. A different element shape means +/// a different policy with the same three members, which is the only thing that +/// has to change to bin that shape against the same grid. +template +struct MeshSimplexElements; + +/** + * Build the grid cell -> element candidate map with `Pred`: the element-major + * build, and the point search's hot path at construction time. + */ +template > +Kokkos::Crs +construct_intersection_map(Omega_h::Mesh& mesh, + Kokkos::View[1]> grid, + int num_grid_cells); + +/// O(num_cells * nelems) brute-force counterpart, for tests. \see +/// construct_intersection_map +template > +Kokkos::Crs +construct_intersection_map_reference(Omega_h::Mesh& mesh, + Kokkos::View[1]> grid, + int num_grid_cells); + } // namespace detail [[nodiscard]] KOKKOS_FUNCTION bool triangle_intersects_bbox( @@ -175,6 +204,20 @@ class GridPointSearch2D : public PointLocalizationSearch2D GridPointSearch2D(Omega_h::Mesh& mesh, LO Nx, LO Ny, const PointSearchTolerances& tolerances); + /** + * Build the candidate map with a predicate of the caller's choosing: + * + * GridPointSearch2D search{mesh, nx, ny, tolerances, MyPredicate{}}; + * + * The predicate decides which elements are offered as candidates for a grid + * cell. The query itself is predicate-independent -- it tests containment of + * the point in each candidate -- so a narrower predicate costs a pass over + * more candidates but cannot change which answer is reported. + */ + template + GridPointSearch2D(Omega_h::Mesh& mesh, LO Nx, LO Ny, + const PointSearchTolerances& tolerances, Pred); + /** * Given points in global coordinates, returns for each the entity (vertex, * edge or face, chosen with the vertex and edge tolerances) containing it @@ -187,6 +230,13 @@ class GridPointSearch2D : public PointLocalizationSearch2D Kokkos::View results) const override; private: + /// Grid and adjacency: everything the map needs except the predicate. + void InitializeGeometry(Omega_h::Mesh& mesh, LO Nx, LO Ny); + + /// The candidate map, which is the only predicate-dependent part. + template + void InitializeMap(); + Omega_h::Mesh mesh_; Omega_h::Adj tris2edges_adj_; Omega_h::Adj tris2verts_adj_; @@ -194,6 +244,7 @@ class GridPointSearch2D : public PointLocalizationSearch2D Omega_h::Adj edges2faces_up_; Omega_h::Adj verts2faces_up_; Kokkos::View grid_{"uniform grid"}; + LO num_grid_cells_ = 0; CandidateMapT candidate_map_; Omega_h::LOs tris2verts_; Omega_h::Reals coords_; @@ -211,6 +262,11 @@ class GridPointSearch3D : public PointLocalizationSearch3D GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz, const PointSearchTolerances& tolerances); + /// \see GridPointSearch2D's predicate-taking constructor. + template + GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz, + const PointSearchTolerances& tolerances, Pred); + /** * Given points in global coordinates, returns for each the id of a * tetrahedron containing it and the point's barycentric coordinates in that @@ -234,6 +290,13 @@ class GridPointSearch3D : public PointLocalizationSearch3D Kokkos::View results) const override; private: + /// \see GridPointSearch2D::InitializeGeometry + void InitializeGeometry(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz); + + /// \see GridPointSearch2D::InitializeMap + template + void InitializeMap(); + Omega_h::Mesh mesh_; Omega_h::Adj tris2edges_adj_; Omega_h::Adj tris2verts_adj_; @@ -242,11 +305,46 @@ class GridPointSearch3D : public PointLocalizationSearch3D Omega_h::Adj edges2regions_up_; Omega_h::Adj faces2regions_up_; Kokkos::View[1]> grid_{"uniform grid"}; + LO num_grid_cells_ = 0; CandidateMapT candidate_map_; Omega_h::LOs tris2verts_; Omega_h::Reals coords_; }; +template +GridPointSearch2D::GridPointSearch2D(Omega_h::Mesh& mesh, LO Nx, LO Ny, + const PointSearchTolerances& tolerances, + Pred) + : PointLocalizationSearch2D(tolerances), mesh_(mesh) +{ + InitializeGeometry(mesh, Nx, Ny); + InitializeMap(); +} + +template +void GridPointSearch2D::InitializeMap() +{ + candidate_map_ = + detail::construct_intersection_map<2, Pred>(mesh_, grid_, num_grid_cells_); +} + +template +GridPointSearch3D::GridPointSearch3D(Omega_h::Mesh& mesh, LO Nx, LO Ny, LO Nz, + const PointSearchTolerances& tolerances, + Pred) + : PointLocalizationSearch3D(tolerances), mesh_(mesh) +{ + InitializeGeometry(mesh, Nx, Ny, Nz); + InitializeMap(); +} + +template +void GridPointSearch3D::InitializeMap() +{ + candidate_map_ = + detail::construct_intersection_map<3, Pred>(mesh_, grid_, num_grid_cells_); +} + /// A grid search of either spatial dimension, as owned by components that are /// dimension-agnostic at compile time (e.g. an evaluator factory over an /// Omega_h mesh) and handed to consumers that need to locate points on the same From a8768c9b18fcbc122e6eb988c2f0c0b66ee91899 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Tue, 27 Oct 2026 15:00:07 -0400 Subject: [PATCH 3/7] create distributed evaluation --- src/pcms/field/CMakeLists.txt | 1 + .../evaluator/distributed_point_evaluator.h | 503 ++++++++++ src/pcms/localization/CMakeLists.txt | 4 + src/pcms/localization/candidate_rank_table.h | 273 ++++++ .../localization/distributed_point_search.cpp | 278 ++++++ .../localization/distributed_point_search.h | 231 +++++ src/pcms/transfer/mesh_intersection.hpp | 62 +- src/pcms/utility/CMakeLists.txt | 1 + src/pcms/utility/arrays.h | 73 ++ src/pcms/utility/mpi_utilities.h | 179 ++++ src/pcms/utility/uniform_grid.h | 103 ++ test/CMakeLists.txt | 23 + test/distributed_test_utils.h | 214 ++++ test/test_cell_range.cpp | 270 ++++++ test/test_distributed_evaluate.cpp | 695 +++++++++++++ test/test_distributed_search.cpp | 915 ++++++++++++++++++ 16 files changed, 3765 insertions(+), 60 deletions(-) create mode 100644 src/pcms/field/evaluator/distributed_point_evaluator.h create mode 100644 src/pcms/localization/candidate_rank_table.h create mode 100644 src/pcms/localization/distributed_point_search.cpp create mode 100644 src/pcms/localization/distributed_point_search.h create mode 100644 src/pcms/utility/mpi_utilities.h create mode 100644 test/distributed_test_utils.h create mode 100644 test/test_cell_range.cpp create mode 100644 test/test_distributed_evaluate.cpp create mode 100644 test/test_distributed_search.cpp diff --git a/src/pcms/field/CMakeLists.txt b/src/pcms/field/CMakeLists.txt index db8388b5c..69ae802c1 100644 --- a/src/pcms/field/CMakeLists.txt +++ b/src/pcms/field/CMakeLists.txt @@ -44,6 +44,7 @@ if (PCMS_ENABLE_OMEGA_H) layout/omega_h_entity.h layout/omega_h_lagrange.h evaluator/omega_h_lagrange.h + evaluator/distributed_point_evaluator.h layout/uniform_grid.h uniform_grid_binary_field.h evaluator/uniform_grid.h diff --git a/src/pcms/field/evaluator/distributed_point_evaluator.h b/src/pcms/field/evaluator/distributed_point_evaluator.h new file mode 100644 index 000000000..04dbb3450 --- /dev/null +++ b/src/pcms/field/evaluator/distributed_point_evaluator.h @@ -0,0 +1,503 @@ +#ifndef PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H +#define PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H + +#include "pcms/field/evaluation_request.h" +#include "pcms/field/field.h" +#include "pcms/field/function_space.h" +#include "pcms/field/out_of_bounds_policy.h" +#include "pcms/field/point_evaluator.h" +#include "pcms/localization/distributed_point_search.h" +#include "pcms/localization/point_search.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/assert.h" +#include "pcms/utility/memory_spaces.h" +#include "pcms/utility/mpi_utilities.h" +#include "pcms/utility/uniform_grid.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace pcms +{ + +namespace detail +{ + +/// The concrete local search for a spatial dimension. +template +struct GridSearchOf; +template <> +struct GridSearchOf<2> +{ + using type = GridPointSearch2D; +}; +template <> +struct GridSearchOf<3> +{ + using type = GridPointSearch3D; +}; + +/// A search over the whole decomposition, so every rank builds the same grid. +/// +/// The local evaluator factories build their search from the local element +/// count, which differs per rank. Routing needs the opposite: a grid every rank +/// agrees on, because the rank a point is routed to is read off that grid. +template +typename GridSearchOf::type MakeGlobalGridSearch(Omega_h::Mesh& mesh) +{ + const auto n = DivisionsPerAxisForGlobalMesh(mesh); + if constexpr (Dim == 3) { + return typename GridSearchOf::type(mesh, n, n, n); + } else { + return typename GridSearchOf::type(mesh, n, n); + } +} + +/// The reading an owner writes for a point its own evaluator could not place. +/// +/// Owner-side evaluation has to say "no answer" inside the payload, because the +/// reply is sent before the requester can be told anything else, and because +/// the local evaluator API has no way to report an unplaced point other than +/// through the value it wrote. The out-of-bounds policy supplies the marker: an +/// owner evaluates with FILL and a fill value nothing else can produce. +/// +/// An exception would be the obvious alternative and is not available: an owner +/// evaluates while its peers are blocked waiting for its reply, so throwing +/// there would leave the other ranks in the exchange, not raise an error. +/// +/// The cost is that a field value equal to the marker is read as "no answer". +/// NaN cannot be a coefficient of a meaningful field, and the integral marker +/// is the type's lowest value, so both are safe in practice; it is the reason +/// the confidence is in the test rather than in the type system. +template +[[nodiscard]] inline T UnansweredMarker() +{ + if constexpr (std::is_floating_point_v) { + return std::numeric_limits::quiet_NaN(); + } else { + return std::numeric_limits::lowest(); + } +} + +/// Whether `value` is the marker. Device-callable; `marker` is passed in +/// because `std::numeric_limits` is host-only. +template +[[nodiscard]] KOKKOS_INLINE_FUNCTION bool IsUnanswered(T value, T marker) +{ + if constexpr (std::is_floating_point_v) { + (void)marker; + return value != value; // NaN, without 's overload set + } else { + return value == marker; + } +} + +using DeviceExec = DeviceMemorySpace::execution_space; + +/// Set every entry of `values` to `fill`. +template +void FillAll(Values values, T fill) +{ + const int nc = static_cast(values.extent(1)); + Kokkos::parallel_for( + "DistributedPointEvaluator::Fill", + Kokkos::RangePolicy(0, values.extent(0)), KOKKOS_LAMBDA(LO i) { + for (int c = 0; c < nc; ++c) { + values(i, c) = fill; + } + }); +} + +/// Write row `k` of `rows` to row `targets(k)` of `values`, skipping a row that +/// holds the no-answer marker. Returns how many rows were skipped. +template +LO ScatterAnswered(Values values, Kokkos::View rows, + Kokkos::View targets) +{ + const int nc = static_cast(values.extent(1)); + const T marker = UnansweredMarker(); + LO skipped = 0; + Kokkos::parallel_reduce( + "DistributedPointEvaluator::Scatter", + Kokkos::RangePolicy(0, targets.extent(0)), + KOKKOS_LAMBDA(LO k, LO & n_skipped) { + for (int c = 0; c < nc; ++c) { + if (IsUnanswered(rows(k, c), marker)) { + ++n_skipped; + return; + } + } + const LO i = targets(k); + for (int c = 0; c < nc; ++c) { + values(i, c) = rows(k, c); + } + }, + skipped); + return skipped; +} + +/// `counts` times `factor`, for sending `factor` scalars per point. +inline std::vector Scaled(const std::vector& counts, int factor) +{ + std::vector out(counts.size()); + for (size_t r = 0; r < counts.size(); ++r) { + out[r] = counts[r] * factor; + } + return out; +} + +} // namespace detail + +/** + * \brief Evaluate a field at points that are distributed over ranks. + * + * Any rank can evaluate the field at any point, including points outside its + * own partition, and every rank gets the same answer. + * + * Construction (collective, runs once because the query points are fixed): + * 1. `DistributedPointSearch::Locate` finds the owning rank of every point: + * local search first, then the global search grid gives candidate ranks, + * which search the point locally; the lowest owning rank wins. + * 2. Each point owned by another rank is sent to its owner, once. + * 3. Each rank builds two ordinary local evaluators through + * `FunctionSpace::CreatePointEvaluator`: one for its own points, and one + * for the points other ranks sent it. + * + * `Evaluate` (collective): each rank evaluates both batches locally, and one + * `MPI_Alltoallv` sends the values for the incoming batch back to the ranks + * that asked. Locations are never searched again; values are recomputed on + * every call, from the coefficients as they are now. + * + * Nothing here is specific to a field type: a backend gets distributed + * evaluation as soon as it can evaluate locally. + * + * A rank only ever evaluates points it owns, so a ghost element's (possibly + * stale) coefficients never enter an answer. If an owner's local evaluator + * still cannot place a point -- its own search disagreeing with the routing + * search -- it writes `detail::UnansweredMarker()` and the point is treated + * as not found. + * + * \tparam T Value type. Each backend accepts `CreatePointEvaluator` for + * whichever types it supports. + * \tparam Dim Spatial dimension of the mesh, 2 or 3. + */ +template > +class DistributedPointEvaluator : public PointEvaluator +{ +public: + /** + * \param mesh The mesh points are located on, and whose ownership decides who + * answers for a point. Must outlive this object. + * \param space The field's function space. Every local evaluation goes + * through it. + * \param points Query coordinates, `[num_points][Dim]`, fixed for the life of + * the evaluator. + * \param policy What to do with a point no rank can locate. + * \param tolerance_reach How far outside an element the search may still + * accept a point. + * \param grid_divisions Cells per axis of the global search grid; 0 derives + * it from the mesh. Must be the same on every rank. + */ + DistributedPointEvaluator( + Omega_h::Mesh& mesh, std::shared_ptr space, + Kokkos::View points, + OutOfBoundsPolicy policy = {}, Real tolerance_reach = 1e-12, + LO grid_divisions = 0) + : mesh_(mesh), space_(std::move(space)), policy_(policy) + { + if (!space_) { + throw std::invalid_argument( + "DistributedPointEvaluator: function space is null"); + } + if (mesh_.dim() != Dim) { + throw std::invalid_argument( + "DistributedPointEvaluator: mesh dimension does not match"); + } + n_components_ = space_->GetLayout()->GetNumComponents(); + n_points_ = static_cast(points.extent(0)); + Setup(points, tolerance_reach, grid_divisions); + } + + /// Evaluate the field at every query point. Collective. + /// + /// \param field Supplies the coefficients. Must be compatible with this + /// evaluator's function space. + /// \param values Output, `[num_query_points][num_components]`. + void Evaluate( + const Field& field, + Rank2View values) const override; + + /// How many query points get their value from another rank. + [[nodiscard]] std::size_t LastRoutedPointCount() const noexcept + { + return remote_points_.extent(0); + } + + /// The space whose local evaluator produces the values. + [[nodiscard]] const FunctionSpace& GetFunctionSpace() const noexcept + { + return *space_; + } + +private: + void Setup(Kokkos::View points, + Real tolerance_reach, LO grid_divisions); + + /// A local evaluator over `coords`, or null when there are none. `coords` + /// must stay alive for as long as the evaluator. + [[nodiscard]] std::unique_ptr> MakeLocalEvaluator( + Kokkos::View coords) const; + + /// Run `evaluator` over its `n` points, on the device. + [[nodiscard]] Kokkos::View EvaluateLocally( + const PointEvaluator* evaluator, LO n, const Field& field) const; + + Omega_h::Mesh& mesh_; + std::shared_ptr space_; + OutOfBoundsPolicy policy_; + int n_components_ = 0; + LO n_points_ = 0; + + // Query points this rank owns (indices into the query set), their + // coordinates, and the evaluator over them. + Kokkos::View own_points_; + Kokkos::View own_coords_; + std::unique_ptr> own_evaluator_; + + // Query points another rank owns, grouped by owner rank in ascending order, + // and how many go to each rank. + Kokkos::View remote_points_; + std::vector remote_counts_; + + // Points other ranks sent this rank to evaluate, grouped by requesting rank, + // and how many came from each rank. + std::vector incoming_counts_; + Kokkos::View incoming_coords_; + std::unique_ptr> incoming_evaluator_; +}; + +template +void DistributedPointEvaluator::Setup( + Kokkos::View points, + Real tolerance_reach, LO grid_divisions) +{ + // ---- 1. find the owning rank of every point ---- + // The search is only needed here: once the owners are known, `Evaluate` + // never searches again. + const auto local_search = detail::MakeGlobalGridSearch(mesh_); + DistributedPointSearch search(mesh_, local_search, tolerance_reach, + grid_divisions); + const auto located = search.Locate(points); + + // Split the query indices by owner. The search result is host data, so this + // is host work on indices only; coordinates stay on the device. + const int rank = search.Rank(); + const int nranks = mesh_.comm()->size(); + std::vector> by_owner(static_cast(nranks)); + std::vector own; + for (LO i = 0; i < n_points_; ++i) { + const auto& r = located[static_cast(i)]; + if (r.status != DistributedPointSearch::ResultStatus::LOCATED) { + continue; // in neither list; `Evaluate` fills it + } + if (r.owner_rank == rank) { + own.push_back(i); + } else { + by_owner[static_cast(r.owner_rank)].push_back(i); + } + } + std::vector remote; + remote_counts_.assign(static_cast(nranks), 0); + for (int r = 0; r < nranks; ++r) { + const auto& points_of_r = by_owner[static_cast(r)]; + remote.insert(remote.end(), points_of_r.begin(), points_of_r.end()); + remote_counts_[static_cast(r)] = + static_cast(points_of_r.size()); + } + own_points_ = CopyToDevice(own); + remote_points_ = CopyToDevice(remote); + + // ---- 2. send each remote point to its owner, once ---- + // Only the remote points' coordinates leave the device, because MPI needs + // them on the host. + const auto send = + CopyToHostRowMajor(GatherRows(points, remote_points_)); + std::vector incoming_scalar_counts; + const auto incoming = + ExchangeRecords(mesh_.comm()->get_impl(), detail::Scaled(remote_counts_, Dim), + send, incoming_scalar_counts); + incoming_counts_.resize(incoming_scalar_counts.size()); + for (size_t r = 0; r < incoming_counts_.size(); ++r) { + incoming_counts_[r] = incoming_scalar_counts[r] / Dim; + } + + // ---- 3. one local evaluator per batch, built once ---- + own_coords_ = GatherRows(points, own_points_); // device only + incoming_coords_ = CopyRowMajorToDevice(incoming, Dim); + own_evaluator_ = MakeLocalEvaluator(own_coords_); + incoming_evaluator_ = MakeLocalEvaluator(incoming_coords_); +} + +template +std::unique_ptr> +DistributedPointEvaluator::MakeLocalEvaluator( + Kokkos::View coords) const +{ + if (coords.extent(0) == 0) { + return nullptr; + } + // Only owned points reach this evaluator, so one it cannot place is a + // disagreement between its search and the routing search. FILL with a + // marker lets `Evaluate` treat such a point as not found. + OutOfBoundsPolicy fallback; + fallback.mode = OutOfBoundsMode::FILL; + fallback.fill_value = static_cast(detail::UnansweredMarker()); + + const CoordinateView view(space_->GetCoordinateSystem(), + MakeConstRank2View(coords)); + return space_->CreatePointEvaluator( + EvaluationRequest::FromCoordinates(view, fallback)); +} + +template +Kokkos::View +DistributedPointEvaluator::EvaluateLocally( + const PointEvaluator* evaluator, LO n, const Field& field) const +{ + Kokkos::View values("local values", n, + n_components_); + if (evaluator != nullptr) { + evaluator->Evaluate(field, MakeRank2View(values)); + } + return values; +} + +template +void DistributedPointEvaluator::Evaluate( + const Field& field, + Rank2View values) const +{ + PCMS_ALWAYS_ASSERT(values.extent(0) == static_cast(n_points_)); + PCMS_ALWAYS_ASSERT(values.extent(1) == static_cast(n_components_)); + const int nc = n_components_; + const auto n_own = static_cast(own_points_.extent(0)); + const auto n_remote = static_cast(remote_points_.extent(0)); + + // Owner side: evaluate this rank's own points and the points sent to it. + const auto own_values = EvaluateLocally(own_evaluator_.get(), n_own, field); + const auto incoming_values = EvaluateLocally( + incoming_evaluator_.get(), static_cast(incoming_coords_.extent(0)), + field); + + // Route the incoming values back to the ranks that asked: the only values + // that cross to the host and back, because MPI takes host buffers. One scalar per component per point. + const auto remote_host = ExchangeRecordsWithKnownCounts( + mesh_.comm()->get_impl(), detail::Scaled(incoming_counts_, nc), + CopyToHostRowMajor(incoming_values), detail::Scaled(remote_counts_, nc)); + const auto remote_values = CopyRowMajorToDevice(remote_host, nc); + + // Assemble on the device. Every point starts at the policy's fill value, and + // a point is overwritten only when its owner answered for it. + detail::FillAll(values, static_cast(policy_.fill_value)); + LO missing = n_points_ - n_own - n_remote; // located by no rank + missing += detail::ScatterAnswered(values, own_values, own_points_); + missing += detail::ScatterAnswered(values, remote_values, remote_points_); + + // Thrown after the exchange, so every rank has completed it. + if (policy_.mode == OutOfBoundsMode::ERROR && missing > 0) { + throw pcms_error( + "DistributedPointEvaluator: a query point lies outside the mesh and the " + "out-of-bounds policy is ERROR"); + } +} + +/** + * \brief Make a distributed evaluator for `space`, dispatching on the mesh + * dimension. Collective. + * + * \param points Query coordinates, `[num_points][mesh_dim]`; copied into the + * layout the evaluator uses. + * \param grid_divisions Cells per axis of the global search grid; 0 derives it + * from the mesh. Must be the same on every rank. + * + * The returned evaluator is an ordinary `PointEvaluator`, so it can be used + * wherever a serial one can. + */ +template +[[nodiscard]] std::unique_ptr> MakeDistributedPointEvaluator( + Omega_h::Mesh& mesh, std::shared_ptr space, + Kokkos::View points, + OutOfBoundsPolicy policy = {}, Real tolerance_reach = 1e-12, + LO grid_divisions = 0) +{ + const int mesh_dim = mesh.dim(); + const auto n_points = static_cast(points.extent(0)); + if (mesh_dim != 2 && mesh_dim != 3) { + throw std::invalid_argument( + "MakeDistributedPointEvaluator: only 2D and 3D meshes are supported"); + } + if (points.extent(1) != static_cast(mesh_dim)) { + throw std::invalid_argument( + "MakeDistributedPointEvaluator: query coordinates are not [num_points x " + "mesh_dim]"); + } + + if (mesh_dim == 2) { + Kokkos::View coords("distributed coords", + n_points); + Kokkos::parallel_for( + "copy_distributed_coords", + Kokkos::RangePolicy(0, n_points), + KOKKOS_LAMBDA(LO i) { + coords(i, 0) = points(i, 0); + coords(i, 1) = points(i, 1); + }); + return std::make_unique>( + mesh, std::move(space), coords, policy, tolerance_reach, grid_divisions); + } + Kokkos::View coords("distributed coords", + n_points); + Kokkos::parallel_for( + "copy_distributed_coords", + Kokkos::RangePolicy(0, n_points), + KOKKOS_LAMBDA(LO i) { + coords(i, 0) = points(i, 0); + coords(i, 1) = points(i, 1); + coords(i, 2) = points(i, 2); + }); + return std::make_unique>( + mesh, std::move(space), coords, policy, tolerance_reach, grid_divisions); +} + +/// How many query points of a distributed evaluator get their value from +/// another rank, or 0 when `evaluator` is not one. +/// +/// Diagnostics for callers that hold only the `PointEvaluator` interface, which +/// is the interface a distributed evaluator is meant to be used through. +template +[[nodiscard]] std::size_t LastRoutedPointCount( + const PointEvaluator& evaluator) +{ + if (auto* d2 = + dynamic_cast*>(&evaluator)) { + return d2->LastRoutedPointCount(); + } + if (auto* d3 = + dynamic_cast*>(&evaluator)) { + return d3->LastRoutedPointCount(); + } + return 0; +} + +} // namespace pcms + +#endif // PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H \ No newline at end of file diff --git a/src/pcms/localization/CMakeLists.txt b/src/pcms/localization/CMakeLists.txt index 269b1b776..4ec72a590 100644 --- a/src/pcms/localization/CMakeLists.txt +++ b/src/pcms/localization/CMakeLists.txt @@ -1,6 +1,9 @@ set( PCMS_LOCALIZATION_HEADERS point_search.h + overlap.h + candidate_rank_table.h + distributed_point_search.h queue_visited.hpp adj_search.hpp mls_support_helpers.h @@ -12,6 +15,7 @@ set( set( PCMS_LOCALIZATION_SOURCES point_search.cpp + distributed_point_search.cpp adj_search.cpp mls_support_helpers.cpp mesh_localization.cpp diff --git a/src/pcms/localization/candidate_rank_table.h b/src/pcms/localization/candidate_rank_table.h new file mode 100644 index 000000000..6d7adf894 --- /dev/null +++ b/src/pcms/localization/candidate_rank_table.h @@ -0,0 +1,273 @@ +#ifndef PCMS_LOCALIZATION_CANDIDATE_RANK_TABLE_H +#define PCMS_LOCALIZATION_CANDIDATE_RANK_TABLE_H + +#include "pcms/localization/overlap.h" +#include "pcms/utility/bounding_box.h" +#include "pcms/utility/mpi_utilities.h" +#include "pcms/utility/types.h" +#include "pcms/utility/uniform_grid.h" + +#include +#include +#include + +#include +#include + +namespace pcms +{ + +/** + * \brief Which ranks may own the element containing a query point. + * + * A distributed point search cannot tell from a point's coordinates which rank + * owns it -- that is the whole problem. This table answers the routing question + * without a search: every rank publishes the range of global search cells that + * its own geometry falls in, and a point is then routed to the ranks whose + * published range contains the point's cell. + * + * The search structure is implicit. `UniformGrid` is a handful of scalars, so + * every rank holds the entire global structure without materializing anything + * per cell, and the only thing exchanged is one cell range per rank. Memory is + * therefore O(ranks) per rank and does not grow with the grid resolution, which + * is what lets the caller choose a dense grid for selectivity without a + * corresponding memory cost. + * + * Make sure if a point is contained by an element that `rank` owns, then + * `rank` is among the candidates for that point. The converse does not hold: a + * candidate list is a superset of the ranks that can actually localize the + * point, so the caller must still search and merge the answers. + * + * Two properties of the grid make the guarantee work: + * + * - **The grid must be identical on every rank.** A linear cell ID only names + * the same region of space everywhere if every rank derived the same grid. + * `CreateUniformGridFromMesh` reduces the bounding box globally, so that half + * is automatic, but the resolution must come from a global quantity too -- + * `DivisionsPerAxisForGlobalMesh`, not `DivisionsPerAxisForMesh(nelems())`, + * since `nelems()` counts ghost copies and so varies with the partition. The + * convenience constructor uses the global form; the explicit-grid constructor + * exists so a caller that already owns a shared grid uses that one. + * + * - **The published range must be widened by the containment tolerance.** A + * point search accepts a point up to some distance outside an element (the + * vertex/edge tolerance bands in 2D, the facet plane band in 3D), so a point + * legitimately contained by an element can lie slightly outside that element + * and outside the unpadded bounding box that publishes it. Without the + * widening such a point would be routed to no rank at all and silently + * treated as out of bounds. `tolerance_reach` is that distance. + * + * The bounding box is taken over the rank's *local* coordinates, which include + * ghost vertices. That is a superset of the owned region -- an owned element + * may reference a vertex it does not own -- so it can only ever add candidate + * ranks, never miss one. Keeping it simple this way also makes a rank with no + * local geometry fall out naturally: it publishes an empty range and is never a + * candidate. + * + * Queries are host-side and intended to be called once per point, on the subset + * of points that failed to localize locally, after the search results have been + * copied back from the device. A lookup scans the published ranges, so it is + * O(ranks) per point; the work is bounded by, and much smaller than, the search + * it replaces. + * + * Setup uses a gather rather than a scatter: every rank needs every other + * rank's range, so this is an `allgather` of O(ranks) values, not an all-to-all + * with per-rank payloads. + */ +template +class CandidateRankTable +{ +public: + static constexpr int dim = Dim; + using GridT = UniformGrid; + // Qualified: the target is the namespace-scope `pcms::CellRange`, which an + // unqualified spelling would shadow with this very alias. + using CellRange = pcms::CellRange; + + /** + * Build on the shared routing grid. + * + * The grid is derived from the global element count and the globally reduced + * bounding box, so every rank constructs the same one. + * + * \param mesh The mesh whose elements are being located. + * \param tolerance_reach How far outside an element the point search may + * still accept a point; the published ranges are widened by this much. \see + * PointLocalizationSearch::PointSearchTolerances + */ + explicit CandidateRankTable(Omega_h::Mesh& mesh, Real tolerance_reach = 0.0) + : CandidateRankTable(mesh, SharedRoutingGrid(mesh), tolerance_reach) + { + } + + /** + * Build on a caller-supplied grid. + * + * Use this when the caller already holds the shared grid, so that the grid + * used to route and the grid used for anything else are provably the same. + * + * \param mesh The mesh whose elements are being located. + * \param grid The global search grid; must be identical on every rank. + * \param tolerance_reach \see the convenience constructor. + */ + CandidateRankTable(Omega_h::Mesh& mesh, const GridT& grid, + Real tolerance_reach = 0.0) + : grid_(grid) + { + if (mesh.dim() != Dim) { + throw std::invalid_argument( + "CandidateRankTable: mesh dimension does not match the table"); + } + + auto comm = mesh.comm(); + const int nranks = comm->size(); + local_range_ = LocalCellRange(mesh, tolerance_reach); + + // Every rank needs every other rank's range, so this is an all-ranks gather + // of a fixed 2*Dim values per rank. `Omega_h::Comm::allgather` cannot be + // used: it only exchanges with the ranks of the mesh's communication graph. + // \see pcms/utility/mpi_utilities.h + Kokkos::Array local{}; + for (int d = 0; d < Dim; ++d) { + local[static_cast(d)] = local_range_.lo[d]; + local[static_cast(Dim + d)] = local_range_.hi[d]; + } + const auto gathered = + AllGatherAllRanks(comm->get_impl(), local.data(), 2 * Dim); + + ranges_.assign(static_cast(nranks), CellRange{}); + for (int r = 0; r < nranks; ++r) { + for (int d = 0; d < Dim; ++d) { + ranges_[static_cast(r)].lo[d] = + gathered[static_cast(r) * 2 * Dim + d]; + ranges_[static_cast(r)].hi[d] = + gathered[static_cast(r) * 2 * Dim + Dim + d]; + } + } + } + + /// The shared routing grid for `mesh`, identical on every rank. + [[nodiscard]] static GridT SharedRoutingGrid(Omega_h::Mesh& mesh, + LO min_divisions = 10) + { + return CreateUniformGridFromMesh( + mesh, DivisionsPerAxisForGlobalMesh(mesh, min_divisions)); + } + + /// The routing grid for `mesh` with `divisions` cells per axis, or the + /// default resolution (`SharedRoutingGrid`) when `divisions` is 0. + /// + /// `divisions` must be the same on every rank, or cell ids would name + /// different regions of space on different ranks. + [[nodiscard]] static GridT RoutingGrid(Omega_h::Mesh& mesh, LO divisions) + { + if (divisions > 0) { + return CreateUniformGridFromMesh(mesh, divisions); + } + return SharedRoutingGrid(mesh); + } + + /// The grid a cell ID is interpreted in. Identical on every rank. + [[nodiscard]] const GridT& GetGrid() const noexcept { return grid_; } + + /// Number of ranks in the decomposition, i.e. the size of the table. + [[nodiscard]] int NumRanks() const noexcept + { + return static_cast(ranges_.size()); + } + + /// The range this rank published. Empty when it holds no geometry. + [[nodiscard]] const CellRange& LocalRange() const noexcept + { + return local_range_; + } + + /// The range `rank` published. + [[nodiscard]] const CellRange& RangeForRank(int rank) const + { + return ranges_.at(static_cast(rank)); + } + + /// All published ranges, indexed by rank. + [[nodiscard]] const std::vector& Ranges() const noexcept + { + return ranges_; + } + + /// Ranks that may own the element containing a point in `cell`, ascending. + /// + /// Empty when no rank can own the cell: a cell in a hole of a non-convex + /// domain, or outside the decomposition entirely. That is a legitimate answer + /// and the caller treats it as out of bounds. + [[nodiscard]] std::vector CandidatesForCell(LO cell) const + { + const auto axis_indices = grid_.AxisIndicesFromCellIndex(cell); + std::vector candidates; + candidates.reserve(ranges_.size()); + for (size_t r = 0; r < ranges_.size(); ++r) { + if (ranges_[r].Contains(axis_indices)) { + candidates.push_back(static_cast(r)); + } + } + return candidates; + } + + /// Ranks that may own the element containing `point`, ascending. + /// + /// The point is mapped to its cell first, clamped to the grid, so a point + /// outside the domain is routed by its nearest boundary cell rather than + /// being rejected here. That is deliberate: a point just outside the domain + /// can still be accepted by a boundary element within the containment + /// tolerance, and it is that tolerance the published ranges were widened by. + [[nodiscard]] std::vector CandidatesForPoint( + const Omega_h::Vector& point) const + { + return CandidatesForCell(grid_.ClosestCellID(point)); + } + + /// True when at least one rank may own the element containing `point`. + [[nodiscard]] bool HasCandidates(const Omega_h::Vector& point) const + { + return !CandidatesForPoint(point).empty(); + } + +private: + /// Cell range covering this rank's local geometry, widened by the containment + /// tolerance. Empty when the rank holds no coordinates. + [[nodiscard]] CellRange LocalCellRange(Omega_h::Mesh& mesh, + Real tolerance_reach) const + { + // A rank with no local coordinates owns no elements either, so there is + // nothing for a query to find and the range stays empty. This also keeps + // `find_bounding_box` away from a rank whose reduced bounding box would + // come back inverted: it starts from plus/minus the largest representable + // value, so an empty point set yields max < min. + if (mesh.nverts() == 0) { + return CellRange::Empty(); + } + + // Local coordinates, not owned ones: an owned element can reference a + // vertex it does not own, and including those vertices can only widen the + // range, which is safe. Never allreduces, so this stays per-rank. + const auto bbox = Omega_h::find_bounding_box(mesh.coords()); + + AABBox box{}; + for (int d = 0; d < Dim; ++d) { + box.center[d] = (bbox.min[d] + bbox.max[d]) / 2.0; + // Widen by the tolerance so that a point the search accepts just outside + // this rank's geometry still routes here. \see the class comment. + box.half_width[d] = (bbox.max[d] - bbox.min[d]) / 2.0 + tolerance_reach; + } + // Cells that touch the box count, which is what `candidate_cells` encodes: + // a vertex sitting on a cell boundary makes the touched cell a candidate. + return candidate_cells(box, grid_); + } + + GridT grid_{}; + CellRange local_range_{}; + std::vector ranges_; +}; + +} // namespace pcms + +#endif // PCMS_LOCALIZATION_CANDIDATE_RANK_TABLE_H \ No newline at end of file diff --git a/src/pcms/localization/distributed_point_search.cpp b/src/pcms/localization/distributed_point_search.cpp new file mode 100644 index 000000000..fd93c9f48 --- /dev/null +++ b/src/pcms/localization/distributed_point_search.cpp @@ -0,0 +1,278 @@ +#include "pcms/localization/distributed_point_search.h" +#include "pcms/utility/mpi_utilities.h" + +#include + +#include +#include +#include + +namespace pcms +{ + +namespace +{ + +/// A query point sent to a candidate rank. +template +struct LocateRequest +{ + LO point; // index on the requesting rank + Real coords[Dim]; +}; + +/// A candidate rank's answer for a point it owns. +template +struct LocateReply +{ + LO point; // index on the requesting rank + LO element_id; + GO global_element_id; + Real barycentric[Dim + 1]; +}; + +/// `divisions`, after checking it is the same on every rank. The check is +/// collective, so every rank reaches the same verdict and throws together. +LO AgreedGridDivisions(Omega_h::Mesh& mesh, LO divisions) +{ + const LO lowest = mesh.comm()->allreduce(divisions, Omega_h_Op::OMEGA_H_MIN); + const LO highest = mesh.comm()->allreduce(divisions, Omega_h_Op::OMEGA_H_MAX); + if (lowest != highest) { + throw std::invalid_argument( + "DistributedPointSearch: grid_divisions must be the same on every rank"); + } + if (lowest < 0) { + throw std::invalid_argument( + "DistributedPointSearch: grid_divisions must not be negative"); + } + return lowest; +} + +} // namespace + +template +DistributedPointSearch::DistributedPointSearch( + Omega_h::Mesh& mesh, const LocalSearch& local_search, Real tolerance_reach, + LO grid_divisions) + : mesh_(mesh), + local_search_(local_search), + routing_(mesh, + CandidateRankTable::RoutingGrid( + mesh, AgreedGridDivisions(mesh, grid_divisions)), + tolerance_reach), + tolerance_reach_(tolerance_reach) +{ + if (mesh.dim() != Dim) { + throw std::invalid_argument( + "DistributedPointSearch: mesh dimension does not match the search"); + } + rank_ = mesh.comm()->rank(); + + // Ownership is consulted once per point, so read it once. `nghost_layers() > + // 0` is also the guard for using `Mesh::owned` at all: on a mesh that has not + // been partitioned Omega_h refuses to answer an ownership query on more than + // one rank, and without ghosts every local element is owned anyway. + has_ghost_layer_ = mesh.nghost_layers() > 0; + if (has_ghost_layer_) { + const auto owned = Omega_h::HostRead(mesh.owned(mesh.dim())); + owned_elements_.assign(owned.begin(), owned.end()); + } else { + owned_elements_.assign(static_cast(mesh.nelems()), 1); + } + + // Global ids identify an element across ranks. They are what lets the answer + // to "which element" be checked, and what a caller caches; a mesh without + // them still works, the id is just unavailable. + if (mesh.has_tag(mesh.dim(), "global")) { + const auto globals = Omega_h::HostRead(mesh.globals(mesh.dim())); + global_ids_.assign(globals.begin(), globals.end()); + } +} + +template +std::vector::LocalClassification> +DistributedPointSearch::ClassifyLocally( + Kokkos::View points) const +{ + const auto n_points = static_cast(points.extent(0)); + std::vector classification( + static_cast(n_points)); + + if (n_points == 0) { + return classification; + } + + const auto results = local_search_(points); + const auto owners = local_search_.GetOwningElementIds(results); + + const auto results_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, results); + const auto owners_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, owners); + + for (LO i = 0; i < n_points; ++i) { + auto& out = classification[static_cast(i)]; + const LO element_id = owners_h(i); + // `IsOutside` is the search's own notion of "no element contains this + // point": it covers both an empty candidate cell and a point the search + // could only resolve to a nearest entity. + if (results_h(i).IsOutside() || element_id < 0) { + out.status = LocalStatus::NOT_FOUND_HERE; + out.element_id = -1; + continue; + } + out.element_id = element_id; + out.barycentric = results_h(i).parametric_coords; + out.status = OwnsLocalElement(element_id) ? LocalStatus::LOCATED_HERE + : LocalStatus::LOCATED_ELSEWHERE; + } + return classification; +} + +template +std::vector::Result> +DistributedPointSearch::Locate(Kokkos::View points, + std::vector* cache) +{ + const auto n_points = static_cast(points.extent(0)); + std::vector storage; + std::vector& results = cache != nullptr ? *cache : storage; + // Resizing keeps entries already resolved, so a caller that passes the same + // cache for the same point set only pays for the points still missing. + results.resize(static_cast(n_points)); + + const auto classified = ClassifyLocally(points); + const auto points_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, points); + const auto& grid = routing_.GetGrid(); + const int nranks = mesh_.comm()->size(); + + // ---- 1. local search, 2. candidates from the global grid ---- + std::vector>> requests_by_rank( + static_cast(nranks)); + last_routed_ = 0; + for (LO i = 0; i < n_points; ++i) { + auto& r = results[static_cast(i)]; + if (r.status == ResultStatus::LOCATED) { + continue; // cached + } + const auto& c = classified[static_cast(i)]; + if (c.status == LocalStatus::LOCATED_HERE) { + r.status = ResultStatus::LOCATED; + r.owner_rank = rank_; + r.element_id = c.element_id; + r.global_element_id = GlobalElementId(c.element_id); + r.barycentric = c.barycentric; + continue; + } + LocateRequest request{}; + request.point = i; + Omega_h::Vector point; + for (int d = 0; d < Dim; ++d) { + point[d] = points_h(i, d); + request.coords[d] = points_h(i, d); + } + // Outside the global bounding box no element can contain the point, so it + // stays NOT_FOUND and no rank is asked about it. + if (!grid.IsPointInBounds(point, tolerance_reach_)) { + continue; + } + ++last_routed_; + // This rank already answered for anything it owns, so it never asks + // itself. That also stops a ghost hit from bouncing straight back here. + for (const int candidate : routing_.CandidatesForPoint(point)) { + if (candidate != rank_) { + requests_by_rank[static_cast(candidate)].push_back(request); + } + } + } + + // The exchange is collective, so whether to run it is decided on the global + // total: a rank with nothing to send must still join if any other rank sends. + LO send_total = 0; + for (const auto& requests : requests_by_rank) { + send_total += static_cast(requests.size()); + } + if (mesh_.comm()->allreduce(send_total, Omega_h_Op::OMEGA_H_SUM) == 0) { + return results; + } + + // ---- 3. send to candidates, which search locally ---- + auto comm = mesh_.comm()->get_impl(); + std::vector send_counts(static_cast(nranks)); + std::vector> send; + for (int r = 0; r < nranks; ++r) { + const auto& requests = requests_by_rank[static_cast(r)]; + send_counts[static_cast(r)] = static_cast(requests.size()); + send.insert(send.end(), requests.begin(), requests.end()); + } + std::vector recv_counts; + const auto received = ExchangeRecords(comm, send_counts, send, recv_counts); + + Kokkos::View remote_points("remote query points", + received.size()); + auto remote_h = Kokkos::create_mirror_view(remote_points); + for (size_t t = 0; t < received.size(); ++t) { + for (int d = 0; d < Dim; ++d) { + remote_h(t, d) = received[t].coords[d]; + } + } + Kokkos::deep_copy(remote_points, remote_h); + const auto remote_classified = ClassifyLocally(remote_points); + + // Reply only for points this rank owns. A candidate that cannot place a + // point has nothing to say about it. + std::vector reply_counts(static_cast(nranks), 0); + std::vector> replies; + size_t t = 0; + for (int r = 0; r < nranks; ++r) { + for (int k = 0; k < recv_counts[static_cast(r)]; ++k, ++t) { + const auto& c = remote_classified[t]; + if (c.status != LocalStatus::LOCATED_HERE) { + continue; + } + LocateReply reply{}; + reply.point = received[t].point; + reply.element_id = c.element_id; + reply.global_element_id = GlobalElementId(c.element_id); + for (int d = 0; d <= Dim; ++d) { + reply.barycentric[d] = c.barycentric[d]; + } + replies.push_back(reply); + ++reply_counts[static_cast(r)]; + } + } + + // ---- 4. receive the answers; the lowest owning rank wins ---- + std::vector reply_recv_counts; + const auto answers = + ExchangeRecords(comm, reply_counts, replies, reply_recv_counts); + // Answers arrive grouped by source rank in ascending order, so the first + // answer seen for a point is the lowest rank's. + t = 0; + for (int r = 0; r < nranks; ++r) { + for (int k = 0; k < reply_recv_counts[static_cast(r)]; ++k, ++t) { + const auto& a = answers[t]; + if (a.point < 0 || a.point >= n_points) { + continue; + } + auto& result = results[static_cast(a.point)]; + if (result.status == ResultStatus::LOCATED) { + continue; + } + result.status = ResultStatus::LOCATED; + result.owner_rank = r; + result.element_id = a.element_id; + result.global_element_id = a.global_element_id; + for (int d = 0; d <= Dim; ++d) { + result.barycentric[d] = a.barycentric[d]; + } + } + } + return results; +} + +template class DistributedPointSearch<2>; +template class DistributedPointSearch<3>; + +} // namespace pcms \ No newline at end of file diff --git a/src/pcms/localization/distributed_point_search.h b/src/pcms/localization/distributed_point_search.h new file mode 100644 index 000000000..223c31bdd --- /dev/null +++ b/src/pcms/localization/distributed_point_search.h @@ -0,0 +1,231 @@ +#ifndef PCMS_LOCALIZATION_DISTRIBUTED_POINT_SEARCH_H +#define PCMS_LOCALIZATION_DISTRIBUTED_POINT_SEARCH_H + +#include "pcms/localization/candidate_rank_table.h" +#include "pcms/localization/point_search.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/memory_spaces.h" +#include "pcms/utility/types.h" + +#include +#include + +#include +#include + +namespace pcms +{ + +/** + * \brief Locate points on a mesh that is distributed over several ranks. + * + * The algorithm, for each query point: + * 1. Search locally. A point inside an element this rank owns is done. + * 2. Otherwise, look the point up in the global search grid + * (`CandidateRankTable`) to get the ranks that may own it. + * 3. Send the point to those ranks; each searches it locally. + * 4. Receive the answers. If several ranks own an element containing the + * point (it sits on a partition interface), the lowest rank wins. + * + * A rank answers only for elements it owns, never for ghosts, because ghost + * data belongs to another rank. + * + * This class only locates points. Evaluating a field at them is the job of + * `DistributedPointEvaluator`, which uses the result to know which rank to ask. + */ +template +class DistributedPointSearch +{ +public: + using LocalSearch = PointLocalizationSearch; + + /// What this rank could determine about a point on its own. + enum class LocalStatus + { + /// Contained by an element this rank owns, so this rank can answer. The + /// search is finished for this point. + LOCATED_HERE, + /// Contained by an element this rank holds only as a ghost. The owner has + /// to answer; routing must not send the point back to this rank. + LOCATED_ELSEWHERE, + /// Not contained by anything this rank holds. The owner, if there is one, + /// is somewhere else. + NOT_FOUND_HERE + }; + + struct LocalClassification + { + LocalStatus status = LocalStatus::NOT_FOUND_HERE; + /// Local element id: the owning element for `LOCATED_HERE`, the ghost that + /// contains the point for `LOCATED_ELSEWHERE`, `-1` otherwise. + LO element_id = -1; + /// Barycentric coordinates in `element_id`, zero when not found. + Omega_h::Vector barycentric{}; + }; + + /// The outcome of a distributed locate for one point. + enum class ResultStatus + { + /// Some rank owns an element containing the point. `owner_rank`, + /// `element_id` and `barycentric` describe it. + LOCATED, + /// No rank could place the point: it is outside the mesh, or outside every + /// cell any rank published. The caller decides whether that is an error or + /// a fill. + NOT_FOUND + }; + + struct Result + { + ResultStatus status = ResultStatus::NOT_FOUND; + /// Rank that owns the containing element, `-1` when not located. + int owner_rank = -1; + /// Global id of the containing element on `owner_rank`, `-1` when the mesh + /// carries no global ids or the point was not located. + GO global_element_id = -1; + /// Local element id **on `owner_rank`**. Only meaningful together with + /// `owner_rank`; it is not an index into this rank's mesh. + LO element_id = -1; + /// Barycentric coordinates inside `element_id`. + Omega_h::Vector barycentric{}; + }; + + /** + * \param mesh The mesh to locate points on. + * \param local_search A search over the same mesh. Must outlive this object. + * \param tolerance_reach How far outside an element the local search may + * still accept a point; the routing ranges are widened by this much. + * \param grid_divisions Cells per axis of the global search grid. 0 (the + * default) derives it from the global element count, about one element per + * cell. Must be the same on every rank; the constructor checks this and + * throws on every rank if not. + * + * Constructing this publishes this rank's cell range and gathers the others', + * so the constructor is collective. + */ + DistributedPointSearch(Omega_h::Mesh& mesh, const LocalSearch& local_search, + Real tolerance_reach = 0.0, LO grid_divisions = 0); + + /** + * \brief Locate a batch of points, searching remote ranks where necessary. + * + * Collective: every rank must call it, because the exchange is a collective. + * Points that this rank owns are answered locally at no communication cost; + * only the rest are routed, and only to the ranks whose published cell range + * covers the point. That keeps the common case -- a rank asking about points + * in its own partition -- free of messages. + * + * A point outside the routing grid costs nothing at all. The grid spans the + * global mesh bounding box and every element lies inside that box, so no rank + * can contain such a point: it is reported `NOT_FOUND` without being routed. + * The box is widened by the search tolerance, so a point the local search + * would still accept just outside a boundary face is not rejected here. + * + * When several candidate ranks report a containing element -- a point on a + * partition interface is contained by an element on each side -- the answer + * from the **lowest rank id** wins. Any of them is geometrically right; the + * rule just makes the outcome deterministic and identical on every rank. + * + * \param points Query coordinates, one per row. + * \param cache Optional, indexed by point. Entries that are already `LOCATED` + * are reused and not searched again, which is what makes repeated + * localization of the same point set cheap. The caller owns the cache + * and must clear it whenever the point coordinates or the mesh change: + * an entry is only valid for the point and mesh it was computed from. + * Points that could not be placed are not cached, so they are retried. + * \return One result per row, in the same order. + */ + [[nodiscard]] std::vector Locate( + Kokkos::View points, + std::vector* cache = nullptr); + + /// How many points the most recent `Locate` sent to another rank. + /// + /// Diagnostics. Zero means every point was answered here, either because this + /// rank owns it, because it lies outside the mesh, or because the cache + /// already held its location. + [[nodiscard]] std::size_t LastRoutedPointCount() const noexcept + { + return last_routed_; + } + + /** + * \brief Classify a batch by what this rank can answer by itself. + * + * A point is `LOCATED_HERE` only when the search reports an element that this + * rank owns. A point the search resolves to a ghost is `LOCATED_ELSEWHERE` + * rather than `LOCATED_HERE` even though the ghost geometrically contains it: + * the ghost's data belongs to another rank, so the owner has to answer. That + * happens for a sizeable fraction of points near a partition interface -- + * roughly a quarter of element-centroid queries on a two-rank mesh -- and for + * points on shared vertices or edges, where the reported element is resolved + * through local adjacency and can land on either side. + * + * \param points Query coordinates, one per row. + * \return One classification per row, in the same order. + */ + [[nodiscard]] std::vector ClassifyLocally( + Kokkos::View points) const; + + /// The mesh this search locates points on. + [[nodiscard]] Omega_h::Mesh& GetMesh() const noexcept { return mesh_; } + + /// The local search the distributed search is built on. + [[nodiscard]] const LocalSearch& GetLocalSearch() const noexcept + { + return local_search_; + } + + /// True when this rank holds ghost elements, i.e. when ownership can differ + /// from "every local element". + [[nodiscard]] bool HasGhostLayer() const noexcept { return has_ghost_layer_; } + + /// Whether this rank owns the local element `element_id`. + [[nodiscard]] bool OwnsLocalElement(LO element_id) const + { + return element_id >= 0 && + static_cast(element_id) < owned_elements_.size() && + owned_elements_[static_cast(element_id)] != 0; + } + + /// The routing table used to pick candidate ranks. + [[nodiscard]] const CandidateRankTable& GetRoutingTable() const noexcept + { + return routing_; + } + + /// This rank's id in the mesh's communicator. + [[nodiscard]] int Rank() const noexcept { return rank_; } + +private: + /// Global id of a local element, or -1 when the mesh carries no global ids. + [[nodiscard]] GO GlobalElementId(LO element_id) const + { + if (global_ids_.empty() || element_id < 0 || + static_cast(element_id) >= global_ids_.size()) { + return -1; + } + return global_ids_[static_cast(element_id)]; + } + + Omega_h::Mesh& mesh_; + const LocalSearch& local_search_; + CandidateRankTable routing_; + int rank_ = 0; + /// Per-local-element ownership, indexed by local element id. All true when + /// the mesh has no ghost layer, since then every local element is owned. + std::vector owned_elements_; + /// Per-local-element global ids, empty when the mesh has no "global" tag. + std::vector global_ids_; + /// How far outside its geometry the local search still accepts a point. + Real tolerance_reach_ = 0; + bool has_ghost_layer_ = false; + std::size_t last_routed_ = 0; +}; + +extern template class DistributedPointSearch<2>; +extern template class DistributedPointSearch<3>; + +} // namespace pcms + +#endif // PCMS_LOCALIZATION_DISTRIBUTED_POINT_SEARCH_H \ No newline at end of file diff --git a/src/pcms/transfer/mesh_intersection.hpp b/src/pcms/transfer/mesh_intersection.hpp index 9d5ead08d..18bd5e882 100644 --- a/src/pcms/transfer/mesh_intersection.hpp +++ b/src/pcms/transfer/mesh_intersection.hpp @@ -2,11 +2,11 @@ #define PCMS_TRANSFER_MESH_INTERSECTION_HPP #include +#include #include #include #include #include -#include #include #include #include @@ -51,65 +51,7 @@ template const r3d::Few, Dim + 1>& a, const r3d::Few, Dim + 1>& b) { - // Each iteration treats one simplex as the half-space owner. Both are - // needed: a face of `a` may separate the pair while no face of `b` does. - for (int side = 0; side < 2; ++side) { - const auto& owner = (side == 0) ? a : b; - const auto& other = (side == 0) ? b : a; - - // Face `f` is the face opposite vertex `f`. - for (int f = 0; f < Dim + 1; ++f) { - // Vertices spanning the face, in index order with `f` skipped. - int face_verts[Dim]; - int n_face_verts = 0; - for (int v = 0; v < Dim + 1; ++v) { - if (v != f) { - face_verts[n_face_verts++] = v; - } - } - const auto& origin = owner[face_verts[0]]; - - // Face normal, oriented to point away from the opposite vertex so the - // owner's interior lies on the negative side. - r3d::Vector normal; - if constexpr (Dim == 2) { - const auto edge = owner[face_verts[1]] - origin; - normal[0] = edge[1]; - normal[1] = -edge[0]; - } else { - const auto e1 = owner[face_verts[1]] - origin; - const auto e2 = owner[face_verts[2]] - origin; - normal[0] = e1[1] * e2[2] - e1[2] * e2[1]; - normal[1] = e1[2] * e2[0] - e1[0] * e2[2]; - normal[2] = e1[0] * e2[1] - e1[1] * e2[0]; - } - - Omega_h::Real opposite_side = 0.0; - Omega_h::Real normal_sq = 0.0; - for (int d = 0; d < Dim; ++d) { - opposite_side += normal[d] * (owner[f][d] - origin[d]); - normal_sq += normal[d] * normal[d]; - } - // A degenerate face gives no usable axis. - if (normal_sq == 0.0) { - continue; - } - const Omega_h::Real orientation = (opposite_side > 0.0) ? -1.0 : 1.0; - - bool all_outside = true; - for (int v = 0; v < Dim + 1 && all_outside; ++v) { - Omega_h::Real signed_distance = 0.0; - for (int d = 0; d < Dim; ++d) { - signed_distance += normal[d] * (other[v][d] - origin[d]); - } - all_outside = (orientation * signed_distance) >= 0.0; - } - if (all_outside) { - return true; - } - } - } - return false; + return Overlap::apply(a, b); } /** diff --git a/src/pcms/utility/CMakeLists.txt b/src/pcms/utility/CMakeLists.txt index 0016de28f..7c98a54ba 100644 --- a/src/pcms/utility/CMakeLists.txt +++ b/src/pcms/utility/CMakeLists.txt @@ -8,6 +8,7 @@ set( eqdsk.h mesh_geometry.h memory_spaces.h + mpi_utilities.h mpi_type.h omega_h_array_utils.h types.h diff --git a/src/pcms/utility/arrays.h b/src/pcms/utility/arrays.h index 4f06347d6..f29bfc647 100644 --- a/src/pcms/utility/arrays.h +++ b/src/pcms/utility/arrays.h @@ -6,6 +6,8 @@ #include +#include + namespace pcms { @@ -502,5 +504,76 @@ void CopyRank1ViewToDevice(Rank1View dest, KOKKOS_LAMBDA(LO i) { dest(i) = src_tmp_d(i); }); } + +// ---- std::vector <-> device copies ---- +// +// Each performs at most one host/device transfer. A row-major `std::vector` +// is the layout MPI buffers use; the device side uses the default layout for +// the device memory space, which may be column-major. + +/// `host` copied to a new device view. +template +Kokkos::View CopyToDevice(const std::vector& host) +{ + Kokkos::View device("CopyToDevice", host.size()); + const Kokkos::View + host_view(host.data(), host.size()); + Kokkos::deep_copy(device, host_view); + return device; +} + +/// A device matrix as a row-major host vector, `extent(0) * extent(1)` long. +template +std::vector> CopyToHostRowMajor( + const Kokkos::View& device) +{ + using U = std::remove_const_t; + std::vector host(device.extent(0) * device.extent(1)); + Kokkos::View + host_view(host.data(), device.extent(0), device.extent(1)); + DeepCopyMismatchLayouts(host_view, device); + return host; +} + +/// A row-major host vector of `n_cols` columns as a new device matrix. +template +Kokkos::View CopyRowMajorToDevice( + const std::vector& host, int n_cols) +{ + const size_t n_rows = + n_cols > 0 ? host.size() / static_cast(n_cols) : 0; + const Kokkos::View + host_view(host.data(), n_rows, n_cols); + Kokkos::View device("CopyRowMajorToDevice", n_rows, + n_cols); + DeepCopyMismatchLayouts(device, host_view); + return device; +} + +/// The rows `rows(k)` of a rank-2 device view, gathered into a new device +/// matrix on the device (no transfer). +template +Kokkos::View +GatherRows(const SrcView& src, Kokkos::View rows) +{ + static_assert(Kokkos::is_view::value && SrcView::rank() == 2, + "GatherRows takes a rank-2 Kokkos::View"); + using T = typename SrcView::non_const_value_type; + const int n_cols = static_cast(src.extent(1)); + Kokkos::View out("GatherRows", rows.extent(0), + n_cols); + Kokkos::parallel_for( + "GatherRows", + Kokkos::RangePolicy(0, rows.extent(0)), + KOKKOS_LAMBDA(LO k) { + for (int c = 0; c < n_cols; ++c) { + out(k, c) = src(rows(k), c); + } + }); + return out; +} + } // namespace pcms #endif // PCMS_COUPLING_ARRAYS_H diff --git a/src/pcms/utility/mpi_utilities.h b/src/pcms/utility/mpi_utilities.h new file mode 100644 index 000000000..9004b5849 --- /dev/null +++ b/src/pcms/utility/mpi_utilities.h @@ -0,0 +1,179 @@ +#ifndef PCMS_MPI_UTILITIES_H +#define PCMS_MPI_UTILITIES_H + +// Small wrappers over the MPI collectives used to exchange data between all +// ranks of a communicator. + +#include "pcms/utility/mpi_type.h" + +#include + +#include +#include +#include +#include + +namespace pcms +{ + +// ---- gathers ---- + +/** + * \brief The displacements matching a per-rank count vector, for the `v` + * collectives (`MPI_Allgatherv`, `MPI_Alltoallv`). + * + * \param counts How many entries each rank contributes, indexed by rank. + * \return The displacement of each rank's block in the concatenated buffer. + */ +[[nodiscard]] inline std::vector DisplacementsOf( + const std::vector& counts) +{ + std::vector displs(counts.size(), 0); + for (size_t r = 1; r < counts.size(); ++r) { + displs[r] = displs[r - 1] + counts[r - 1]; + } + return displs; +} + +/** + * \brief Gather one scalar from every rank, indexed by rank, on every rank. + * + * There is no Omega_h equivalent to reach for here. `Omega_h::Comm::allgather` + * is a *neighbor* gather: it exchanges with the ranks of the mesh's + * communication graph (`Comm::sources()`/`destinations()`) rather than with all + * ranks, and returns one entry per neighbor, ordered by that neighbor list. It + * is therefore wrong -- and quietly wrong, returning a short buffer -- for + * anything keyed by rank. + * + * `Omega_h::Comm::allreduce` does not have this problem: it is a plain + * `MPI_Allreduce` on the communicator. So a quantity that is the same on every + * rank can go through Omega_h, while a per-rank quantity cannot. + * + * \param comm The communicator, for example `mesh.comm()->get_impl()`. + * \param value The value this rank contributes. + * \return `size`-many values, where element `r` is rank `r`'s contribution. + */ +template +[[nodiscard]] std::vector AllGatherAllRanks(MPI_Comm comm, const T& value) +{ + int size = 0; + MPI_Comm_size(comm, &size); + std::vector gathered(static_cast(size)); + MPI_Allgather(&value, 1, GetMPIType(value), gathered.data(), 1, + GetMPIType(value), comm); + return gathered; +} + +/** + * \brief Gather `count` scalars from every rank, each rank's block contiguous. + * + * The result is rank-major: rank `r`'s values are at `[r * count, (r+1) * + * count)`. Use this over the single-value overload when a rank contributes + * several values, so one exchange replaces several. + * + * \param comm The communicator, for example `mesh.comm()->get_impl()`. + * \param values `count` values this rank contributes. + * \param count How many values each rank contributes. + * \return `size * count` values, rank-major. + */ +template +[[nodiscard]] std::vector AllGatherAllRanks(MPI_Comm comm, const T* values, + int count) +{ + int size = 0; + MPI_Comm_size(comm, &size); + std::vector gathered(static_cast(size) * + static_cast(count)); + MPI_Allgather(values, count, GetMPIType(T{}), gathered.data(), count, + GetMPIType(T{}), comm); + return gathered; +} + +/** + * \brief Gather a variable number of scalars from every rank, concatenated in + * rank order, on every rank. + * + * \param comm The communicator, for example `mesh.comm()->get_impl()`. + * \param values The values this rank contributes; may be empty. + * \param counts Optional. Filled with the number of values each rank sent. + * \return All ranks' values, rank 0's first. + */ +template +[[nodiscard]] std::vector AllGatherVAllRanks( + MPI_Comm comm, const std::vector& values, + std::vector* counts = nullptr) +{ + const auto per_rank = + AllGatherAllRanks(comm, static_cast(values.size())); + const auto displs = DisplacementsOf(per_rank); + std::vector gathered(static_cast( + std::accumulate(per_rank.begin(), per_rank.end(), 0))); + MPI_Allgatherv(values.data(), static_cast(values.size()), + GetMPIType(T{}), gathered.data(), per_rank.data(), + displs.data(), GetMPIType(T{}), comm); + if (counts != nullptr) { + *counts = per_rank; + } + return gathered; +} + +// ---- all-to-all exchanges ---- + +/** + * \brief Exchange records whose per-rank counts are already known on both + * sides. + * + * `send` holds the records for rank 0 first, then rank 1, and so on, with + * `send_counts[r]` records for rank `r`. The result is laid out the same way, + * by source rank, with `recv_counts[r]` records from rank `r`. + * + * Records are sent as raw bytes, so `Rec` must be trivially copyable. Every + * rank runs the same binary, so struct padding only costs a few bytes. + */ +template +[[nodiscard]] std::vector ExchangeRecordsWithKnownCounts( + MPI_Comm comm, const std::vector& send_counts, + const std::vector& send, const std::vector& recv_counts) +{ + static_assert(std::is_trivially_copyable_v, + "ExchangeRecords sends records as raw bytes"); + const auto to_bytes = [](const std::vector& counts) { + std::vector bytes(counts.size()); + for (size_t r = 0; r < counts.size(); ++r) { + bytes[r] = counts[r] * static_cast(sizeof(Rec)); + } + return bytes; + }; + const auto send_bytes = to_bytes(send_counts); + const auto recv_bytes = to_bytes(recv_counts); + const auto send_displs = DisplacementsOf(send_bytes); + const auto recv_displs = DisplacementsOf(recv_bytes); + + std::vector recv(static_cast( + std::accumulate(recv_counts.begin(), recv_counts.end(), 0))); + MPI_Alltoallv(send.data(), send_bytes.data(), send_displs.data(), MPI_BYTE, + recv.data(), recv_bytes.data(), recv_displs.data(), MPI_BYTE, + comm); + return recv; +} + +/** + * \brief Exchange records, first telling each rank how many to expect. + * + * Same layout as `ExchangeRecordsWithKnownCounts`. `recv_counts` is filled + * with the number of records received from each rank. + */ +template +[[nodiscard]] std::vector ExchangeRecords( + MPI_Comm comm, const std::vector& send_counts, + const std::vector& send, std::vector& recv_counts) +{ + recv_counts.assign(send_counts.size(), 0); + MPI_Alltoall(send_counts.data(), 1, MPI_INT, recv_counts.data(), 1, MPI_INT, + comm); + return ExchangeRecordsWithKnownCounts(comm, send_counts, send, recv_counts); +} + +} // namespace pcms + +#endif // PCMS_MPI_UTILITIES_H diff --git a/src/pcms/utility/uniform_grid.h b/src/pcms/utility/uniform_grid.h index e6730fcec..81b7a612b 100644 --- a/src/pcms/utility/uniform_grid.h +++ b/src/pcms/utility/uniform_grid.h @@ -1,6 +1,7 @@ #ifndef PCMS_COUPLING_UNIFORM_GRID_H #define PCMS_COUPLING_UNIFORM_GRID_H #include "pcms/utility/bounding_box.h" +#include "Omega_h_comm.hpp" #include "Omega_h_vector.hpp" #include "Omega_h_bbox.hpp" #include "Omega_h_mesh.hpp" @@ -65,6 +66,20 @@ struct UniformGrid return GetCellIndex(indexes); } + /// Per-axis cell indices in coordinate order (x,y,z) for a linear cell ID. + /// + /// The exact inverse of `CellIndexFromAxisIndices`. Provided so callers that + /// work with linear cell IDs -- a routing table keyed by cell, for example -- + /// do not have to reimplement the grid's internal index ordering, which is + /// reversed relative to coordinates and easy to get wrong silently. + [[nodiscard]] KOKKOS_INLINE_FUNCTION Kokkos::Array + AxisIndicesFromCellIndex(LO idx) const + { + auto indices = GetDimensionedIndex(idx); + reverse(indices); + return indices; + } + [[nodiscard]] KOKKOS_INLINE_FUNCTION AABBox GetCellBBOX(LO idx) const { auto index = GetDimensionedIndex(idx); @@ -98,6 +113,24 @@ struct UniformGrid return true; } + /// As `IsPointInBounds`, but a point within `tolerance` of the box on every + /// side still counts as inside. Used to reject query points that no element + /// can contain: the tolerance matches the band where the search still accepts + /// a point just outside a boundary face. + [[nodiscard]] KOKKOS_INLINE_FUNCTION bool IsPointInBounds( + const Omega_h::Vector& point, Real tolerance) const + { + for (size_t i = 0; i < dim; ++i) { + Real coord = point[i]; + Real grid_min = bot_left[i] - tolerance; + Real grid_max = bot_left[i] + edge_length[i] + tolerance; + if (coord < grid_min || coord > grid_max) { + return false; + } + } + return true; + } + [[nodiscard]] KOKKOS_INLINE_FUNCTION Kokkos::Array GetDimensionedIndex(LO idx) const { @@ -161,6 +194,76 @@ LO DivisionsPerAxisForMesh(LO nelems, LO min_divisions = 10) return std::max(per_axis, min_divisions); } +/** + * \brief Number of elements this rank owns, counted on the device. + * + * Used instead of `Omega_h::Mesh::nents_owned`, which is unusable in a + * Kokkos+CUDA build without OpenMP: its reduction functor reads a device view + * from a host loop (`Omega_h_reduce.hpp`, the non-OpenMP `transform_reduce`), + * which trips Kokkos' memory-space access check. + * + * \param mesh A mesh that carries ownership information, i.e. one that has been + * partitioned. Callers must establish that before calling, typically from + * `nghost_layers() > 0`. + * \return The number of elements this rank owns + */ +inline LO OwnedElementCount(Omega_h::Mesh& mesh) +{ + const auto ent_dim = mesh.dim(); + auto owned = mesh.owned(ent_dim); + LO count = 0; + Kokkos::parallel_reduce( + "pcms::OwnedElementCount", owned.size(), + KOKKOS_LAMBDA(LO i, LO & update) { update += owned[i] ? 1 : 0; }, count); + return count; +} + +/** + * \brief Number of elements in the whole decomposition. + * + * This is the quantity a grid resolution must be derived from so that every + * rank constructs the *same* grid. `mesh.nelems()` cannot be used for that + * purpose on its own: it counts the ghost copies a rank holds as well as its + * owned elements, so it varies with the partition. Summing it does not vary, + * because every rank reduces the same total, which is the property that + * matters. + * + * The count is exact. On a mesh with no ghost layers each element lives on + * exactly one rank, so the reduced local count is already the global count. On + * a ghosted mesh the ghost copies make that local count an overestimate, so the + * owned elements are counted instead -- ghosting implies the mesh was + * partitioned and therefore carries the ownership information that needs. + * + * \param mesh The mesh to count elements for + * \return The global number of elements, identical on every rank + */ +inline LO GlobalElementCount(Omega_h::Mesh& mesh) +{ + const auto local_count = + mesh.nghost_layers() > 0 ? OwnedElementCount(mesh) : mesh.nelems(); + return mesh.comm()->allreduce(local_count, Omega_h_Op::OMEGA_H_SUM); +} + +/** + * \brief Divisions per axis giving roughly one element per grid cell, derived + * from the *global* element count. + * + * The distributed counterpart of `DivisionsPerAxisForMesh(mesh.nelems())`, + * which silently produces a different grid on every rank because `nelems()` + * includes ghosts. Reach for this function whenever the resulting grid has to + * describe the whole decomposition rather than one partition. + * + * \tparam dim Spatial dimension of the grid + * \param mesh The mesh the grid is built over + * \param min_divisions Lower bound on the number of cells per axis + * \return The same value on every rank + */ +template +LO DivisionsPerAxisForGlobalMesh(Omega_h::Mesh& mesh, LO min_divisions = 10) +{ + return DivisionsPerAxisForMesh(GlobalElementCount(mesh), min_divisions); +} + /** * \brief Create a uniform grid layout from an Omega_h mesh as a bounding box * diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 8d44f3a43..e8ee62fc6 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -71,6 +71,28 @@ if(PCMS_ENABLE_OMEGA_H) mpi_test(test_distributed_coupling_2p 2 $) endif() + add_exe(test_distributed_evaluate) + if(HOST_NPROC GREATER_EQUAL 2) + mpi_test(test_distributed_evaluate_2p 2 + $) + endif() + if(HOST_NPROC GREATER_EQUAL 3) + mpi_test(test_distributed_evaluate_3p 3 + $) + endif() + if(HOST_NPROC GREATER_EQUAL 4) + mpi_test(test_distributed_evaluate_4p 4 + $) + endif() + add_exe(test_distributed_search) + if(HOST_NPROC GREATER_EQUAL 2) + mpi_test(test_distributed_search_2p 2 + $) + endif() + if(HOST_NPROC GREATER_EQUAL 3) + mpi_test(test_distributed_search_3p 3 + $) + 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) @@ -424,6 +446,7 @@ if(Catch2_FOUND) test_error_handling.cpp test_eqdsk.cpp test_uniform_grid.cpp + test_cell_range.cpp test_field_evaluation.cpp test_field_interpolation.cpp test_field_copy.cpp diff --git a/test/distributed_test_utils.h b/test/distributed_test_utils.h new file mode 100644 index 000000000..995f57fb0 --- /dev/null +++ b/test/distributed_test_utils.h @@ -0,0 +1,214 @@ +#ifndef PCMS_TEST_DISTRIBUTED_TEST_UTILS_H +#define PCMS_TEST_DISTRIBUTED_TEST_UTILS_H + +// Shared helpers for the MPI tests of distributed search and evaluation. +// +// These tests run under mpirun without Catch2, so a failed check aborts every +// rank: a rank that just returned would leave the others blocked in the next +// collective. + +#include + +#include +#include +#include + +#include "pcms/localization/point_search.h" +#include "pcms/utility/mesh_geometry.h" +#include "pcms/utility/mpi_utilities.h" + +#include +#include +#include +#include + +namespace distributed_test +{ + +using pcms::LO; +using pcms::Real; + +/// Abort every rank with `message` unless `condition` holds. +inline void require(bool condition, const char* message) +{ + if (!condition) { + int rank = -1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + std::fprintf(stderr, "[rank %d] check failed: %s\n", rank, message); + MPI_Abort(MPI_COMM_WORLD, 1); + } +} + +/// Abort unless every rank holds the same integer value. +inline void require_same_on_all_ranks(Omega_h::CommPtr comm, LO value, + const char* message) +{ + const auto gathered = pcms::AllGatherAllRanks(comm->get_impl(), value); + for (size_t i = 0; i < gathered.size(); ++i) { + if (gathered[i] != gathered[0]) { + std::fprintf( + stderr, "[rank %d] %s: rank %zu holds %lld, rank 0 holds %lld\n", + comm->rank(), message, i, static_cast(gathered[i]), + static_cast(gathered[0])); + MPI_Abort(MPI_COMM_WORLD, 1); + } + } +} + +/// Abort unless every rank holds the same real value, to `tolerance`. +inline void require_same_real_on_all_ranks(Omega_h::CommPtr comm, Real value, + Real tolerance, const char* message) +{ + const auto gathered = pcms::AllGatherAllRanks(comm->get_impl(), value); + for (size_t i = 0; i < gathered.size(); ++i) { + if (std::fabs(gathered[i] - gathered[0]) > tolerance) { + std::fprintf(stderr, + "[rank %d] %s: rank %zu holds %.17g, rank 0 holds %.17g\n", + comm->rank(), message, i, gathered[i], gathered[0]); + MPI_Abort(MPI_COMM_WORLD, 1); + } + } +} + +/// The unit square (2D) or cube (3D) as a simplex mesh, partitioned over +/// `world` with one ghost layer. The resolution grows with the rank count so +/// every rank holds some elements. +template +Omega_h::Mesh BuildGhostedBox(Omega_h::CommPtr world) +{ + const LO n = (Dim == 2 ? 4 : 2) * world->size(); + const Real z = Dim == 2 ? 0.0 : 1.0; + const LO nz = Dim == 2 ? 0 : n; + auto mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, z, n, n, nz, + false); + mesh.set_parting(OMEGA_H_GHOSTED, 1, false); + return mesh; +} + +/// The grid point search for `Dim`, with `n` divisions per axis. +template +struct LocalSearchOf; +template <> +struct LocalSearchOf<2> +{ + using type = pcms::GridPointSearch2D; +}; +template <> +struct LocalSearchOf<3> +{ + using type = pcms::GridPointSearch3D; +}; + +template +typename LocalSearchOf::type MakeLocalSearch(Omega_h::Mesh& mesh, LO n) +{ + if constexpr (Dim == 3) { + return typename LocalSearchOf::type(mesh, n, n, n); + } else { + return typename LocalSearchOf::type(mesh, n, n); + } +} + +/// Points gathered from every rank onto every rank, in a deterministic order +/// (contributing rank, then local index). +struct GatheredPoints +{ + std::vector coords; // [total * dim], interleaved + std::vector contributor; // [total], the rank each point came from + std::vector global_id; // [total], when gathered per element + int total = 0; +}; + +/// Gather `local` (interleaved, `dim` values per point) from every rank. +inline GatheredPoints GatherGlobalPoints(Omega_h::Mesh& mesh, + const std::vector& local, + int dim) +{ + GatheredPoints out; + std::vector value_counts; + out.coords = + pcms::AllGatherVAllRanks(mesh.comm()->get_impl(), local, &value_counts); + out.total = static_cast(out.coords.size()) / dim; + for (size_t r = 0; r < value_counts.size(); ++r) { + out.contributor.insert(out.contributor.end(), + static_cast(value_counts[r] / dim), + static_cast(r)); + } + return out; +} + +/// Every owned element centroid in the decomposition, with its global id. +inline GatheredPoints GatherOwnedCentroids(Omega_h::Mesh& mesh) +{ + const auto dim = mesh.dim(); + const auto owned = Omega_h::HostRead(mesh.owned(dim)); + const auto centroids = + Omega_h::HostRead(pcms::get_entity_centroids(mesh, dim)); + std::vector local; + std::vector local_globals; + const bool has_globals = mesh.has_tag(dim, "global"); + const auto globals = has_globals + ? Omega_h::HostRead(mesh.globals(dim)) + : Omega_h::HostRead(); + for (LO e = 0; e < mesh.nelems(); ++e) { + if (owned[e] == 0) { + continue; + } + for (int d = 0; d < dim; ++d) { + local.push_back(centroids[e * dim + d]); + } + if (has_globals) { + local_globals.push_back(globals[e]); + } + } + auto out = GatherGlobalPoints(mesh, local, dim); + if (has_globals) { + out.global_id = + pcms::AllGatherVAllRanks(mesh.comm()->get_impl(), local_globals); + } + return out; +} + +/// Every vertex of every owned element, so the points land on the faces and +/// edges elements share with their neighbours. +inline GatheredPoints GatherOwnedElementVertices(Omega_h::Mesh& mesh) +{ + const auto dim = mesh.dim(); + const auto owned = Omega_h::HostRead(mesh.owned(dim)); + const auto verts = Omega_h::HostRead(mesh.ask_verts_of(dim)); + const auto coords = Omega_h::HostRead(mesh.coords()); + std::vector local; + for (LO e = 0; e < mesh.nelems(); ++e) { + if (owned[e] == 0) { + continue; + } + for (int v = 0; v <= dim; ++v) { + const LO vert = verts[e * (dim + 1) + v]; + for (int d = 0; d < dim; ++d) { + local.push_back(coords[vert * dim + d]); + } + } + } + return GatherGlobalPoints(mesh, local, dim); +} + +/// Interleaved coordinates as a device point view, `[n][Dim]`. +template +Kokkos::View ToPointView(const std::vector& interleaved, + const char* label = "points") +{ + const auto n = static_cast(interleaved.size() / Dim); + Kokkos::View points(std::string(label), n); + auto points_h = Kokkos::create_mirror_view(points); + for (LO i = 0; i < n; ++i) { + for (int d = 0; d < Dim; ++d) { + points_h(i, d) = interleaved[static_cast(i) * Dim + d]; + } + } + Kokkos::deep_copy(points, points_h); + return points; +} + +} // namespace distributed_test + +#endif // PCMS_TEST_DISTRIBUTED_TEST_UTILS_H diff --git a/test/test_cell_range.cpp b/test/test_cell_range.cpp new file mode 100644 index 000000000..24e75c1ed --- /dev/null +++ b/test/test_cell_range.cpp @@ -0,0 +1,270 @@ +#include +#include +#include +#include + +// Serial tests of the cell-range layer the routing table is built on: which +// grid cells a bounding box covers (`candidate_cells`, `overlapping_cells`) and +// the invariants of `CellRange`. + +using pcms::AABBox; +using pcms::candidate_cells; +using pcms::CellRange; +using pcms::LO; +using pcms::Overlap; +using pcms::overlapping_cells; +using pcms::Real; +using pcms::SimplexVsBbox; +using pcms::TriangleVsBbox; +using pcms::Uniform2DGrid; +using pcms::UniformGrid; + +TEST_CASE("bbox vs grid cell range") +{ + Uniform2DGrid grid{ + .edge_length = {10, 10}, .bot_left = {0, 0}, .divisions = {10, 10}}; + + SECTION("box inside one cell still covers padded neighbors") + { + AABBox<2> box{.center = {0.5, 0.5}, .half_width = {0.1, 0.1}}; + const auto range = overlapping_cells<2>(box, grid); + // default pad = 1 clamps at the grid border + REQUIRE(range.lo[0] == 0); + REQUIRE(range.lo[1] == 0); + REQUIRE(range.hi[0] == 1); + REQUIRE(range.hi[1] == 1); + } + + SECTION("no padding gives the exact span") + { + AABBox<2> box{.center = {2.5, 2.5}, .half_width = {0.4, 0.4}}; + const auto range = overlapping_cells<2>(box, grid, 0); + REQUIRE(range.lo[0] == 2); + REQUIRE(range.lo[1] == 2); + REQUIRE(range.hi[0] == 2); + REQUIRE(range.hi[1] == 2); + } + + SECTION("box spanning several cells") + { + AABBox<2> box{.center = {2.5, 2.5}, .half_width = {1.0, 1.0}}; + const auto range = overlapping_cells<2>(box, grid, 0); + REQUIRE(range.lo[0] == 1); + REQUIRE(range.lo[1] == 1); + REQUIRE(range.hi[0] == 3); + REQUIRE(range.hi[1] == 3); + } +} + +namespace +{ +/// True when grid cell `(cx, cy)` and `box` meet, each treated as closed. +/// +/// A shared boundary counts, which is the semantics the overlap predicate has +/// and therefore the set `candidate_cells` has to cover. +[[nodiscard]] bool cell_touches_box(const Uniform2DGrid& grid, LO cx, LO cy, + const AABBox<2>& box) +{ + const Kokkos::Array cell{cx, cy}; + for (int d = 0; d < 2; ++d) { + const Real width = grid.edge_length[d] / grid.divisions[d]; + const Real cell_lo = grid.bot_left[d] + cell[d] * width; + const Real cell_hi = cell_lo + width; + const Real box_lo = box.center[d] - box.half_width[d]; + const Real box_hi = box.center[d] + box.half_width[d]; + if (cell_hi < box_lo || cell_lo > box_hi) { + return false; + } + } + return true; +} +} // namespace + +/// The property `candidate_cells` exists to guarantee: its range contains every +/// cell a closed-boundary predicate can accept. A range that is too narrow +/// drops candidates silently -- the element is never offered to the query and +/// nothing downstream can tell -- so this is checked as a coverage property +/// rather than against the exact bounds, which are allowed to over-include. +TEST_CASE( + "candidate_cells covers every cell a closed-boundary predicate accepts") +{ + const Uniform2DGrid grid{ + .edge_length = {10, 10}, .bot_left = {0, 0}, .divisions = {10, 10}}; + + // Centres at, near and between grid lines, and extents that land exactly on a + // grid line: that is when an unwidened range loses the touching cell. + for (const Real center : {0.5, 2.5, 5.0, 9.5}) { + for (const Real half : {0.0, 0.25, 0.5, 1.0, 3.0}) { + const AABBox<2> box{.center = {center, center}, + .half_width = {half, half}}; + const auto range = candidate_cells<2, TriangleVsBbox>(box, grid); + for (LO cx = 0; cx < grid.divisions[0]; ++cx) { + for (LO cy = 0; cy < grid.divisions[1]; ++cy) { + if (cell_touches_box(grid, cx, cy, box)) { + REQUIRE(range.Contains(Kokkos::Array{cx, cy})); + } + } + } + } + } +} + +TEST_CASE("candidate_cells bins a supplied bound") +{ + const Uniform2DGrid grid{ + .edge_length = {10, 10}, .bot_left = {0, 0}, .divisions = {10, 10}}; + + // A degenerate triangle whose bounding box is the single grid line x = y = 3: + // the cells on either side of that line only touch it, so the padding must + // come from the predicate rather than from the caller. + const Omega_h::Matrix<2, 3> degenerate{{3.0, 3.0}, {3.0, 3.0}, {3.0, 3.0}}; + // `candidate_cells` bins a bound it is given; the element side supplies it. + const auto range = candidate_cells<2, TriangleVsBbox>( + pcms::detail::simplex_bbox<2>(degenerate), grid); + REQUIRE(range.Contains(Kokkos::Array{3, 3})); + REQUIRE(range.Contains(Kokkos::Array{2, 2})); +} + +TEST_CASE("CellRange states its own invariants") +{ + const auto empty = pcms::CellRange<2>::Empty(); + REQUIRE(empty.IsEmpty()); + REQUIRE(!empty.Contains(Kokkos::Array{0, 0})); + + const pcms::CellRange<2> unit{.lo = {2, 2}, .hi = {2, 2}}; + REQUIRE(!unit.IsEmpty()); + REQUIRE(unit.Contains(Kokkos::Array{2, 2})); + REQUIRE(!unit.Contains(Kokkos::Array{3, 2})); + + // The empty range must not be what value initialization produces: a + // zero-initialized CellRange denotes the single cell at axis index 0. + const pcms::CellRange<2> zeroed{}; + REQUIRE(!zeroed.IsEmpty()); + REQUIRE(zeroed.Contains(Kokkos::Array{0, 0})); +} + +// What dimension does the overlap + grid-cell-range layer actually support? +// +// `Dim` is a template parameter throughout and every loop runs to `Dim`, so +// nothing in these functions is written for 2 or 3 specifically. That is easy +// to lose by accident -- one hardcoded `[2]` or `<3, 3>` and the layer silently +// becomes 2D/3D only -- so it is pinned here at dimensions no mesh can +// represent, which is exactly where a hidden assumption would show up. +// +// Findings, and the boundary they draw: +// +// generic, and covered here AABBox, intersects, simplex_bbox, CellRange, +// overlapping_cells, candidate_cells, +// UniformGrid, Overlap +// generic over Dim, but the MeshSimplexElements, ElemCellOverlapFunctor, +// map builder needs a mesh construct_intersection_map -- all take `Dim`, +// but are fed an Omega_h mesh, which is <= 3D, so +// dimensions above 3 cannot be reached through +// them +// NOT generic GridPointSearch2D / GridPointSearch3D, which +// are +// concrete classes with hand-written +// per-dimension query logic (VERTEX/EDGE/FACE for +// 2D, REGION for 3D) +// NOT generic, was silent Overlap: nominally +// Dim-generic but its face normal is a 2D +// perpendicular or a 3D cross product. Now +// guarded by a static_assert, so asking for +// another dimension is a compile error instead of +// reading uninitialized memory. +TEST_CASE("overlap and cell-range layers are dimension-generic") +{ + SECTION("1D") + { + const UniformGrid<1> grid{ + .edge_length = {4.0}, .bot_left = {0.0}, .divisions = {4}}; + const AABBox<1> box{.center = {1.5}, .half_width = {0.25}}; + // Cell 1, padded by 1, clamped at the low edge. + const auto range = candidate_cells<1, SimplexVsBbox>(box, grid); + REQUIRE(range.lo[0] == 0); + REQUIRE(range.hi[0] == 2); + REQUIRE(range.Contains(Kokkos::Array{1})); + } + + SECTION("4D") + { + const UniformGrid<4> grid{.edge_length = {4.0, 4.0, 4.0, 4.0}, + .bot_left = {0.0, 0.0, 0.0, 0.0}, + .divisions = {4, 4, 4, 4}}; + const AABBox<4> box{.center = {1.5, 1.5, 1.5, 1.5}, + .half_width = {0.25, 0.25, 0.25, 0.25}}; + const auto range = candidate_cells<4, SimplexVsBbox>(box, grid); + for (int d = 0; d < 4; ++d) { + REQUIRE(range.lo[d] == 0); + REQUIRE(range.hi[d] == 2); + } + REQUIRE(range.Contains(Kokkos::Array{1, 1, 1, 1})); + // Same axis order as 2D: CellIndexFromAxisIndices round-trips. + const auto indices = Kokkos::Array{1, 1, 1, 1}; + const auto linear = grid.CellIndexFromAxisIndices(indices); + const auto back = grid.AxisIndicesFromCellIndex(linear); + for (int d = 0; d < 4; ++d) { + REQUIRE(back[d] == indices[d]); + } + } + + SECTION("5D range invariants") + { + const auto empty = CellRange<5>::Empty(); + REQUIRE(empty.IsEmpty()); + REQUIRE(!empty.Contains(Kokkos::Array{0, 0, 0, 0, 0})); + + const CellRange<5> unit{.lo = {3, 3, 3, 3, 3}, .hi = {3, 3, 3, 3, 3}}; + REQUIRE(!unit.IsEmpty()); + REQUIRE(unit.Contains(Kokkos::Array{3, 3, 3, 3, 3})); + REQUIRE(!unit.Contains(Kokkos::Array{3, 3, 3, 3, 4})); + } + + SECTION("SimplexVsBbox predicate in 4D") + { + const AABBox<4> box{.center = {1.0, 1.0, 1.0, 1.0}, + .half_width = {0.5, 0.5, 0.5, 0.5}}; + + // A degenerate simplex: all 5 vertices at a point inside the box. + Omega_h::Matrix<4, 5> inside{}; + for (int v = 0; v < 5; ++v) { + for (int d = 0; d < 4; ++d) { + inside(d, v) = 1.0; + } + } + REQUIRE(Overlap<4, SimplexVsBbox>::apply(inside, box)); + + // Far outside on every axis. + Omega_h::Matrix<4, 5> outside{}; + for (int v = 0; v < 5; ++v) { + for (int d = 0; d < 4; ++d) { + outside(d, v) = 100.0; + } + } + REQUIRE_FALSE(Overlap<4, SimplexVsBbox>::apply(outside, box)); + + // Straddling the box on one axis only: the conservative predicate is true. + Omega_h::Matrix<4, 5> straddling{}; + for (int v = 0; v < 5; ++v) { + for (int d = 0; d < 4; ++d) { + straddling(d, v) = 1.0; + } + } + straddling(3, 0) = 0.5; + straddling(3, 1) = 1.9; + REQUIRE(Overlap<4, SimplexVsBbox>::apply(straddling, box)); + } + + SECTION("SimplexVsBbox predicate in 5D") + { + const AABBox<5> box{.center = {2.0, 2.0, 2.0, 2.0, 2.0}, + .half_width = {1.0, 1.0, 1.0, 1.0, 1.0}}; + Omega_h::Matrix<5, 6> inside{}; + for (int v = 0; v < 6; ++v) { + for (int d = 0; d < 5; ++d) { + inside(d, v) = 2.0; + } + } + REQUIRE(Overlap<5, SimplexVsBbox>::apply(inside, box)); + } +} diff --git a/test/test_distributed_evaluate.cpp b/test/test_distributed_evaluate.cpp new file mode 100644 index 000000000..05a9f13cf --- /dev/null +++ b/test/test_distributed_evaluate.cpp @@ -0,0 +1,695 @@ +// Distributed evaluation, against the serial evaluator and against the closed +// form. +// +// The distributed evaluator produces values by asking each rank's local +// evaluator about the points that rank owns, so the comparison that matters is +// between a distributed answer and a serial one. The serial evaluator is local: +// it can only answer for points its own partition plus ghosts contain, and +// applies its policy to the rest. So it is asked with a sentinel fill value, +// and a sentinel reading means "not defined here" rather than a disagreement. +// To keep full coverage anyway, every point is additionally checked against the +// closed form, which a linear field on linear elements reproduces exactly. +// +// The last case runs a backend that has no distributed code at all -- the +// uniform-grid evaluator, with a different localization hint and a different +// value rule from the Lagrange one -- through the same distributed evaluator, +// to show that a backend needs nothing but local evaluation to be distributed. +// +// The query points are every owned element centroid of the whole decomposition, +// so each rank is asked about points it does not own and the remote path is +// exercised. Nothing here calls SynchronizeGhosts: a rank that reads only what +// it owns must still produce values for everyone else's points. + +#include "distributed_test_utils.h" + +#include + +#include +#include +#include + +#include "pcms/field/evaluator/distributed_point_evaluator.h" +#include "pcms/field/field.h" +#include "pcms/field/field_data.h" +#include "pcms/field/function_space/lagrange.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/mesh_geometry.h" +#include "pcms/utility/mpi_utilities.h" +#include "pcms/utility/mpi_type.h" +#include "pcms/utility/uniform_grid.h" + +#include +#include +#include +#include +#include +#include + +using distributed_test::require; +using pcms::LO; +using pcms::Real; + +namespace +{ + +/// The fill value the serial reference is asked to use for points it cannot +/// reach. Distinct enough to be unmistakable in a comparison. +constexpr Real kNotDefinedHere = -1.0e300; + +/// The affine test field: exactly representable on linear elements. +KOKKOS_INLINE_FUNCTION Real linear_f(Real x, Real y) +{ + return x + 2.0 * y; +} + +/// Set the DOF holder coefficients from an analytic function. +/// +/// A local copy of the test utility, because this test links neither Catch2 nor +/// the Catch2-based header that utility lives in. +/// +/// `scale` multiplies the analytic function, so a caller can change the field's +/// coefficients without touching its geometry. A Lagrange evaluation is linear +/// in the coefficients, so the value at a fixed point scales with them. +void SetLinearField(pcms::FieldData& data, + const pcms::FieldLayout& layout, Real scale = 1.0) +{ + const auto dof_coords = layout.GetDOFHolderCoordinates().GetValues(); + const LO n = static_cast(dof_coords.extent(0)); + Kokkos::View values("field values", n, 1); + Kokkos::parallel_for( + "set_linear_field", + Kokkos::RangePolicy(0, n), + KOKKOS_LAMBDA(LO i) { + values(i, 0) = scale * linear_f(dof_coords(i, 0), dof_coords(i, 1)); + }); + data.SetDOFHolderData(pcms::MakeRank2View(values)); +} + +/// Query coordinates on the device with their coordinate-system wrapper. +struct QueryCoords +{ + Kokkos::View storage; + pcms::CoordinateView view; +}; + +QueryCoords MakeQueryCoords(const std::vector& interleaved, int dim) +{ + const LO n = static_cast(interleaved.size()) / dim; + Kokkos::View host("query host", n, dim); + for (LO i = 0; i < n; ++i) { + for (int d = 0; d < dim; ++d) { + host(i, d) = + interleaved[static_cast(i) * static_cast(dim) + + static_cast(d)]; + } + } + Kokkos::View device("query device", n, dim); + pcms::DeepCopyMismatchLayouts(device, host); + return QueryCoords{ + device, pcms::CoordinateView( + pcms::CoordinateSystem::Cartesian, pcms::MakeRank2View(device))}; +} + +/// Distributed evaluation must equal the serial evaluator wherever the serial +/// evaluator is defined, and must equal the closed form everywhere. +template +void check_agrees_with_serial(Omega_h::Mesh& mesh, int order) +{ + const auto dim = mesh.dim(); + const auto rank = mesh.comm()->rank(); + + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, order, 1, pcms::CoordinateSystem::Cartesian, "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto field = space->CreateFunction(); + SetLinearField(field.GetData(), *space->GetLayout()); + + const auto gathered = distributed_test::GatherOwnedCentroids(mesh); + require(gathered.total > 0, "no owned element centroids were gathered"); + + auto device_coords = MakeQueryCoords(gathered.coords, dim); + const LO n_points = static_cast(gathered.total); + + // Serial reference, through the public evaluator API, with a fill value that + // marks the points it cannot reach. + pcms::OutOfBoundsPolicy serial_policy; + serial_policy.mode = pcms::OutOfBoundsMode::FILL; + serial_policy.fill_value = kNotDefinedHere; + auto serial = + space->CreatePointEvaluator(pcms::EvaluationRequest::FromCoordinates( + device_coords.view, serial_policy)); + Kokkos::View reference("reference", n_points, + 1); + serial->Evaluate(field, pcms::MakeRank2View(reference)); + const auto reference_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), reference); + + // Distributed, from the same space and points. The space, not a layout, is + // what is handed over: the distributed evaluator reaches a backend only + // through the local evaluator the space creates. + auto distributed = pcms::MakeDistributedPointEvaluator( + mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + + Kokkos::View got("got", n_points, 1); + distributed->Evaluate(field, pcms::MakeRank2View(got)); + const std::size_t routed = pcms::LastRoutedPointCount(*distributed); + const auto got_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), got); + + LO n_compared = 0; + LO n_not_local = 0; + LO n_mismatch = 0; + LO n_bad_vs_closed_form = 0; + Real worst = 0; + Real worst_vs_closed_form = 0; + LO shown = 0; + for (LO i = 0; i < n_points; ++i) { + const Real value = got_h(i, 0); + const Real reference_value = reference_h(i, 0); + const size_t base = static_cast(i) * static_cast(dim); + const Real truth = + linear_f(gathered.coords[base], gathered.coords[base + 1]); + + // Independent of the serial evaluator, so this covers the routed points + // too, which the serial one may not be able to answer at all. + const Real closed_form_error = std::fabs(value - truth); + worst_vs_closed_form = closed_form_error > worst_vs_closed_form + ? closed_form_error + : worst_vs_closed_form; + if (closed_form_error > 1e-10) { + ++n_bad_vs_closed_form; + } + + if (reference_value == kNotDefinedHere) { + ++n_not_local; + continue; + } + ++n_compared; + const Real difference = std::fabs(value - reference_value); + worst = difference > worst ? difference : worst; + if (difference > 1e-10) { + ++n_mismatch; + if (shown < 4) { + std::printf( + " point (%g,%g): distributed=%.9f serial=%.9f exact=%.9f\n", + static_cast(gathered.coords[base]), + static_cast(gathered.coords[base + 1]), + static_cast(value), static_cast(reference_value), + static_cast(truth)); + ++shown; + } + } + } + + const LO t_compared = + mesh.comm()->allreduce(n_compared, Omega_h_Op::OMEGA_H_SUM); + const LO t_not_local = + mesh.comm()->allreduce(n_not_local, Omega_h_Op::OMEGA_H_SUM); + const LO t_mismatch = + mesh.comm()->allreduce(n_mismatch, Omega_h_Op::OMEGA_H_SUM); + const LO t_bad = + mesh.comm()->allreduce(n_bad_vs_closed_form, Omega_h_Op::OMEGA_H_SUM); + const LO t_points = mesh.comm()->allreduce(n_points, Omega_h_Op::OMEGA_H_SUM); + const LO t_routed = + mesh.comm()->allreduce(static_cast(routed), Omega_h_Op::OMEGA_H_SUM); + const Real t_worst = mesh.comm()->allreduce(worst, Omega_h_Op::OMEGA_H_MAX); + const Real t_worst_closed = + mesh.comm()->allreduce(worst_vs_closed_form, Omega_h_Op::OMEGA_H_MAX); + + if (rank == 0) { + std::printf( + "%dD order %d: %lld points, %lld routed to an owner, %lld comparable " + "locally (%lld beyond it), %lld mismatches (worst %.3g)\n", + Dim, order, static_cast(t_points), + static_cast(t_routed), static_cast(t_compared), + static_cast(t_not_local), static_cast(t_mismatch), + static_cast(t_worst)); + std::printf(" against the closed form: %lld bad, worst %.3g\n", + static_cast(t_bad), + static_cast(t_worst_closed)); + } + + require(t_mismatch == 0, + "distributed evaluation disagrees with the serial evaluator"); + require(t_bad == 0, + "distributed evaluation does not reproduce a linear field"); + // A comparison is only meaningful if the two paths actually overlapped, and + // the remote path is only covered if something had to be routed. + require(t_compared > 0, "no point was localizable by the serial evaluator"); + require(t_not_local > 0, + "the serial evaluator reached every point, so the sentinel never " + "marked anything and the comparison is untested at its boundary"); + require(t_routed > 0, + "every point was already local, so the remote path was not used"); + + // A second Evaluate on the same points has to recompute the values, because + // the coefficients change even when the location does not. Scaling the field + // is how the two are told apart: a value left over from the first call, or + // read out of the payload's zero fill, comes back unscaled. Recomputing is + // not free for the points another rank owns -- only its owner holds the + // coefficients -- so the same points are put to their owners again, and the + // routed count repeats. + const Real scale = 3.0; + SetLinearField(field.GetData(), *space->GetLayout(), scale); + Kokkos::View second("second", n_points, 1); + distributed->Evaluate(field, pcms::MakeRank2View(second)); + const LO routed_again = + static_cast(pcms::LastRoutedPointCount(*distributed)); + const auto second_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), second); + + LO n_not_scaled = 0; + for (LO i = 0; i < n_points; ++i) { + const Real expected = scale * got_h(i, 0); + if (std::fabs(second_h(i, 0) - expected) > + 1e-10 * (1.0 + std::fabs(expected))) { + ++n_not_scaled; + } + } + const LO t_not_scaled = + mesh.comm()->allreduce(n_not_scaled, Omega_h_Op::OMEGA_H_SUM); + const LO t_routed_again = + mesh.comm()->allreduce(routed_again, Omega_h_Op::OMEGA_H_SUM); + + if (rank == 0) { + std::printf( + " second pass on the same points: %lld points re-asked of their owners, " + "%lld stale values after the field was scaled by %g\n", + static_cast(t_routed_again), + static_cast(t_not_scaled), static_cast(scale)); + } + // The same points need an owner's answer on both passes: the location did not + // change, but the value has to come from the coefficients as they are now. + require(t_routed_again == t_routed, + "a repeated evaluation did not re-ask the owners of the same points"); + require(t_not_scaled == 0, + "a repeated evaluation returned a value computed from the old " + "coefficients"); +} + +/// Order 0 is piecewise constant, so a point on a shared edge has as many +/// candidate values as there are elements touching it, and they differ. Exactly +/// one answer is correct for a distributed evaluator, and it has to be the same +/// answer on every rank, or two ranks disagree about a shared boundary. +/// +/// This checks the arbitration rule on such points: the result is identical on +/// every rank, it equals a real element's coefficient rather than an average or +/// a stale value, and the query set really did contain contested points. +template +void check_boundary_arbitration(Omega_h::Mesh& mesh) +{ + const auto dim = mesh.dim(); + const auto rank = mesh.comm()->rank(); + const int nranks = mesh.comm()->size(); + auto comm = mesh.comm()->get_impl(); + + // Order 0: one coefficient per element, set from the element's own centroid, + // so neighbouring elements hold different values. + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/0, 1, pcms::CoordinateSystem::Cartesian, "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto field = space->CreateFunction(); + SetLinearField(field.GetData(), *space->GetLayout()); + + const auto gathered = + distributed_test::GatherOwnedElementVertices(mesh); + require(gathered.total > 0, "no owned element vertices were gathered"); + auto device_coords = MakeQueryCoords(gathered.coords, dim); + const LO n_points = static_cast(gathered.total); + + // Every rank's own answer, with the sentinel for the points it cannot reach. + pcms::OutOfBoundsPolicy serial_policy; + serial_policy.mode = pcms::OutOfBoundsMode::FILL; + serial_policy.fill_value = kNotDefinedHere; + auto serial = + space->CreatePointEvaluator(pcms::EvaluationRequest::FromCoordinates( + device_coords.view, serial_policy)); + Kokkos::View reference("reference", n_points, + 1); + serial->Evaluate(field, pcms::MakeRank2View(reference)); + const auto reference_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), reference); + + auto distributed = pcms::MakeDistributedPointEvaluator( + mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + Kokkos::View got("got", n_points, 1); + distributed->Evaluate(field, pcms::MakeRank2View(got)); + const auto got_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), got); + + // Gather everyone's distributed answer and everyone's serial answer, so both + // can be compared point by point across ranks. + const auto gather_all = [&](const Real* local) { + std::vector all(static_cast(nranks) * + static_cast(n_points)); + std::vector counts(static_cast(nranks), 0); + std::vector displs(static_cast(nranks), 0); + for (int r = 0; r < nranks; ++r) { + counts[static_cast(r)] = static_cast(n_points); + displs[static_cast(r)] = static_cast(n_points) * r; + } + MPI_Allgatherv(local, static_cast(n_points), pcms::GetMPIType(Real{}), + all.data(), counts.data(), displs.data(), + pcms::GetMPIType(Real{}), comm); + return all; + }; + const auto all_distributed = gather_all(got_h.data()); + const auto all_serial = gather_all(reference_h.data()); + + LO n_rank_disagreement = 0; + LO n_not_an_element = 0; + LO n_contested = 0; + for (LO i = 0; i < n_points; ++i) { + const Real value = got_h(i, 0); + + // The same point must be answered identically everywhere. + for (int r = 1; r < nranks; ++r) { + const Real other = + all_distributed[static_cast(r) * static_cast(n_points) + + static_cast(i)]; + if (std::fabs(other - value) > 1e-12) { + ++n_rank_disagreement; + } + } + + // The value must be some element's coefficient, so it has to appear among + // the serial answers of the ranks that could locate it. + bool found = false; + LO distinct = 0; + Real first = 0; + for (int r = 0; r < nranks; ++r) { + const Real candidate = + all_serial[static_cast(r) * static_cast(n_points) + + static_cast(i)]; + if (candidate == kNotDefinedHere) { + continue; + } + if (std::fabs(candidate - value) <= 1e-12) { + found = true; + } + if (distinct == 0 || std::fabs(candidate - first) > 1e-12) { + first = candidate; + ++distinct; + } + } + if (!found) { + ++n_not_an_element; + } + // Two ranks offering different coefficients for one point is what makes the + // arbitration rule load-bearing. + if (distinct > 1) { + ++n_contested; + } + } + + const auto reduce_sum = [&](LO value) { + return mesh.comm()->allreduce(value, Omega_h_Op::OMEGA_H_SUM); + }; + const LO t_disagreement = reduce_sum(n_rank_disagreement); + const LO t_not_an_element = reduce_sum(n_not_an_element); + const LO t_contested = reduce_sum(n_contested); + const LO t_points = reduce_sum(n_points); + + if (rank == 0) { + std::printf( + "%dD order 0 arbitration: %lld boundary points, %lld contested, %lld " + "rank disagreements, %lld not an element coefficient\n", + Dim, static_cast(t_points), + static_cast(t_contested), + static_cast(t_disagreement), + static_cast(t_not_an_element)); + } + + // A vacuous pass would be a test that never saw a contested point. + require(t_contested > 0, + "no query point was contested, so the arbitration rule is untested"); + require(t_disagreement == 0, "ranks disagree about a shared boundary point"); + require(t_not_an_element == 0, + "a distributed value is not any element's coefficient"); +} + +/// A backend with no distributed code of its own still gets distributed +/// evaluation, because the distributed evaluator reaches a backend only through +/// the local evaluator its `FunctionSpace` creates. +/// +/// The backend here is the uniform-grid one, which is as different from +/// Lagrange as it gets: its localization hint is a cell index plus parametric +/// coordinates, not an element plus barycentric weights, and its value rule is +/// multilinear interpolation over a structured grid, not a Lagrange combination +/// over a simplex. None of that reaches the distributed layer, which hands the +/// coordinates of the points a rank owns to that rank's own evaluator and +/// carries the result back as bytes. +/// +/// The grid is global and identical on every rank, so an owner's answer is the +/// answer and the two paths must agree exactly. The remote path is still +/// exercised, and the closed form still holds, because multilinear +/// interpolation reproduces an affine field exactly on a tensor-product grid. +template +void check_backend_without_distributed_code(Omega_h::Mesh& mesh) +{ + const auto dim = mesh.dim(); + const auto rank = mesh.comm()->rank(); + + // Both types are spelled out rather than deduced: the grid's type depends on + // `Dim`, and a dependent `auto` would make `space` dependent too, which stops + // the compiler from recognizing `CreateFunction` below as a + // template-id before instantiation. + using Grid = pcms::UniformGrid(Dim)>; + const Grid grid = + pcms::CreateUniformGridFromMesh(Dim)>(mesh, 8); + const std::shared_ptr space = + pcms::LagrangeFunctionSpace::FromUniformGrid( + grid, 1, pcms::CoordinateSystem::Cartesian, /*order=*/1); + auto field = space->CreateFunction(); + SetLinearField(field.GetData(), *space->GetLayout()); + + const auto gathered = distributed_test::GatherOwnedCentroids(mesh); + require(gathered.total > 0, "no owned element centroids were gathered"); + auto device_coords = MakeQueryCoords(gathered.coords, dim); + const LO n_points = static_cast(gathered.total); + + // The serial evaluator for this backend, with a sentinel for the points it + // cannot reach. + pcms::OutOfBoundsPolicy serial_policy; + serial_policy.mode = pcms::OutOfBoundsMode::FILL; + serial_policy.fill_value = kNotDefinedHere; + auto serial = + space->CreatePointEvaluator(pcms::EvaluationRequest::FromCoordinates( + device_coords.view, serial_policy)); + Kokkos::View reference("reference", n_points, + 1); + serial->Evaluate(field, pcms::MakeRank2View(reference)); + const auto reference_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), reference); + + // The same backend, distributed. + auto distributed = pcms::MakeDistributedPointEvaluator( + mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + Kokkos::View got("got", n_points, 1); + distributed->Evaluate(field, pcms::MakeRank2View(got)); + const std::size_t routed = pcms::LastRoutedPointCount(*distributed); + const auto got_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), got); + + LO n_mismatch = 0; + LO n_serial_reached = 0; + LO n_bad_closed_form = 0; + Real worst = 0; + Real worst_closed_form = 0; + for (LO i = 0; i < n_points; ++i) { + const Real value = got_h(i, 0); + const size_t base = static_cast(i) * static_cast(dim); + const Real truth = + linear_f(gathered.coords[base], gathered.coords[base + 1]); + const Real closed_form_error = std::fabs(value - truth); + worst_closed_form = closed_form_error > worst_closed_form + ? closed_form_error + : worst_closed_form; + if (closed_form_error > 1e-10) { + ++n_bad_closed_form; + } + + if (reference_h(i, 0) == kNotDefinedHere) { + continue; + } + ++n_serial_reached; + const Real difference = std::fabs(value - reference_h(i, 0)); + worst = difference > worst ? difference : worst; + if (difference > 1e-10) { + ++n_mismatch; + } + } + + const LO t_mismatch = + mesh.comm()->allreduce(n_mismatch, Omega_h_Op::OMEGA_H_SUM); + const LO t_serial_reached = + mesh.comm()->allreduce(n_serial_reached, Omega_h_Op::OMEGA_H_SUM); + const LO t_bad = + mesh.comm()->allreduce(n_bad_closed_form, Omega_h_Op::OMEGA_H_SUM); + const LO t_points = mesh.comm()->allreduce(n_points, Omega_h_Op::OMEGA_H_SUM); + const LO t_routed = + mesh.comm()->allreduce(static_cast(routed), Omega_h_Op::OMEGA_H_SUM); + const Real t_worst = mesh.comm()->allreduce(worst, Omega_h_Op::OMEGA_H_MAX); + const Real t_worst_closed = + mesh.comm()->allreduce(worst_closed_form, Omega_h_Op::OMEGA_H_MAX); + + if (rank == 0) { + std::printf( + "%dD uniform-grid backend, no distributed code: %lld points, %lld " + "routed, %lld comparable, %lld mismatches (worst %.3g)\n", + Dim, static_cast(t_points), static_cast(t_routed), + static_cast(t_serial_reached), + static_cast(t_mismatch), static_cast(t_worst)); + std::printf(" against the closed form: %lld bad, worst %.3g\n", + static_cast(t_bad), + static_cast(t_worst_closed)); + } + + require(t_mismatch == 0, + "the distributed path disagrees with the second backend's own " + "evaluator"); + require(t_bad == 0, + "the second backend does not reproduce a linear field when " + "distributed"); + require(t_routed > 0, + "every point was already local, so the remote path was not used with " + "the second backend"); + // The sentinel must not be hiding anything: this backend's grid spans the + // whole mesh, so its own evaluator reaches every query point. + require(t_serial_reached == t_points, + "the second backend's serial evaluator could not reach every query " + "point, so the comparison is incomplete"); +} + +/// The policy is applied to exactly the points nothing answered for, and it is +/// applied after the exchange, so ERROR cannot leave the other ranks waiting. +/// +/// Two things are worth pinning down here. A point no rank can locate is the +/// only one the fill touches, so a located point is never overwritten by it +/// even when the location and the payload disagree (which is what the no-answer +/// marker exists to detect). And ERROR has to be raised once the exchange is +/// finished: the point set is the same on every rank, so every rank raises and +/// none is left blocked in a collective. +template +void check_out_of_bounds_policy(Omega_h::Mesh& mesh) +{ + const auto dim = mesh.dim(); + + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/1, 1, pcms::CoordinateSystem::Cartesian, "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto field = space->CreateFunction(); + SetLinearField(field.GetData(), *space->GetLayout()); + + const auto gathered = distributed_test::GatherOwnedCentroids(mesh); + require(gathered.total > 0, "no owned element centroids were gathered"); + + // The interior points, plus one far outside the mesh. The outside point is + // outside the routing grid too, so it is never routed. + std::vector coords = gathered.coords; + for (int d = 0; d < dim; ++d) { + coords.push_back(-10.0); + } + auto device_coords = MakeQueryCoords(coords, dim); + const LO n_points = static_cast(gathered.total) + 1; + const LO outside_point = n_points - 1; + const Real fill_value = -1234.5; + + pcms::OutOfBoundsPolicy fill_policy; + fill_policy.mode = pcms::OutOfBoundsMode::FILL; + fill_policy.fill_value = fill_value; + auto filled = pcms::MakeDistributedPointEvaluator( + mesh, space, device_coords.storage, fill_policy); + Kokkos::View got("got", n_points, 1); + filled->Evaluate(field, pcms::MakeRank2View(got)); + const auto got_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), got); + + LO n_interior_bad = 0; + LO n_outside_filled = 0; + for (LO i = 0; i < n_points; ++i) { + if (i == outside_point) { + if (got_h(i, 0) == fill_value) { + ++n_outside_filled; + } + continue; + } + const size_t base = static_cast(i) * static_cast(dim); + const Real truth = + linear_f(gathered.coords[base], gathered.coords[base + 1]); + if (std::fabs(got_h(i, 0) - truth) > 1e-10) { + ++n_interior_bad; + } + } + + const LO t_interior_bad = + mesh.comm()->allreduce(n_interior_bad, Omega_h_Op::OMEGA_H_SUM); + const LO t_outside_filled = + mesh.comm()->allreduce(n_outside_filled, Omega_h_Op::OMEGA_H_SUM); + const LO t_points = mesh.comm()->allreduce(n_points, Omega_h_Op::OMEGA_H_SUM); + + if (mesh.comm()->rank() == 0) { + std::printf( + "%dD FILL policy: %lld points, %lld filled for the point outside the " + "mesh, %lld interior answers wrong\n", + Dim, static_cast(t_points), + static_cast(t_outside_filled), + static_cast(t_interior_bad)); + } + + require(t_interior_bad == 0, + "a point inside the mesh was treated as out of bounds"); + // The query set is replicated, so every rank has the one outside point and + // every rank must have filled it. + require(t_outside_filled == static_cast(mesh.comm()->size()), + "the point outside the mesh was not filled on every rank"); + + // ERROR, on the same point set: every rank has the point, so every rank + // raises, after the exchange rather than inside it. + pcms::OutOfBoundsPolicy error_policy; + error_policy.mode = pcms::OutOfBoundsMode::ERROR; + bool threw = false; + try { + auto strict = pcms::MakeDistributedPointEvaluator( + mesh, space, device_coords.storage, error_policy); + Kokkos::View strict_values("strict", + n_points, 1); + strict->Evaluate(field, pcms::MakeRank2View(strict_values)); + } catch (const pcms::pcms_error&) { + threw = true; + } + require(threw, "the ERROR policy did not report a point outside the mesh"); +} + +} // namespace + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + { + Omega_h::Library library(&argc, &argv); + auto world = library.world(); + require(world->size() >= 2, "test requires at least 2 MPI ranks"); + + { + auto mesh = distributed_test::BuildGhostedBox<2>(world); + check_agrees_with_serial<2>(mesh, /*order=*/1); + check_agrees_with_serial<2>(mesh, /*order=*/0); + check_boundary_arbitration<2>(mesh); + check_backend_without_distributed_code<2>(mesh); + check_out_of_bounds_policy<2>(mesh); + } + + { + auto mesh = distributed_test::BuildGhostedBox<3>(world); + check_agrees_with_serial<3>(mesh, /*order=*/1); + check_out_of_bounds_policy<3>(mesh); + } + + if (world->rank() == 0) { + std::printf("distributed evaluation test passed\n"); + } + } + MPI_Finalize(); + return 0; +} diff --git a/test/test_distributed_search.cpp b/test/test_distributed_search.cpp new file mode 100644 index 000000000..b0658a8fe --- /dev/null +++ b/test/test_distributed_search.cpp @@ -0,0 +1,915 @@ +// Distributed point search, layer by layer. +// +// 1. The routing table (`CandidateRankTable`): the global search grid is the +// same on every rank, and a point inside an element a rank owns always +// lists that rank as a candidate. +// 2. The local pass (`ClassifyLocally`): a rank answers only for elements it +// owns, never for ghosts. +// 3. The full search (`Locate`): every point is attributed to its owning rank +// and element, including points that had to be routed, and a warm cache +// routes nothing. +#include "distributed_test_utils.h" + +#include "pcms/localization/candidate_rank_table.h" +#include "pcms/localization/distributed_point_search.h" +#include "pcms/utility/uniform_grid.h" + +#include + +#include +#include +#include +#include +#include +#include + +using distributed_test::require; +using distributed_test::require_same_on_all_ranks; +using distributed_test::require_same_real_on_all_ranks; +using pcms::LO; +using pcms::Real; + +namespace +{ + +// ---- 1. routing table ---- + +/// Points of a simplex on a barycentric lattice of denominator `denom`. +/// +/// Enumerating barycentric weights covers every vertex, every edge and face +/// interior point, and the interior, in one sweep, which is what the routing +/// guarantee has to hold for. +template +std::vector> SimplexSamples( + const std::array, Dim + 1>& vertices, int denom) +{ + std::vector> samples; + std::array weight{}; + while (true) { + int head_sum = 0; + for (int v = 0; v < Dim; ++v) { + head_sum += weight[v]; + } + if (head_sum <= denom) { + weight[Dim] = denom - head_sum; + std::array point{}; + for (int v = 0; v <= Dim; ++v) { + for (int d = 0; d < Dim; ++d) { + point[d] += static_cast(weight[v]) * vertices[v][d]; + } + } + for (int d = 0; d < Dim; ++d) { + point[d] /= static_cast(denom); + } + samples.push_back(point); + } + // Odometer over the first Dim weights; the last is determined by them. + int axis = 0; + for (; axis < Dim; ++axis) { + if (++weight[axis] <= denom) { + break; + } + weight[axis] = 0; + } + if (axis == Dim) { + break; + } + } + return samples; +} + +/// The routing grid must be the shared one, or a cell id means something +/// different on each rank and every lookup below is meaningless. +/// +/// Its resolution therefore has to come from the global element count. The +/// local count, `mesh.nelems()`, includes ghost copies and differs per rank, +/// so this only means something on a mesh that really has ghosts. +template +void check_shared_grid(Omega_h::Mesh& mesh, + const pcms::CandidateRankTable& table) +{ + auto comm = mesh.comm(); + const auto& grid = table.GetGrid(); + const auto expected = pcms::CandidateRankTable::SharedRoutingGrid(mesh); + + const auto global_count = pcms::GlobalElementCount(mesh); + require(comm->allreduce(mesh.nelems(), Omega_h_Op::OMEGA_H_SUM) > + global_count, + "no rank holds ghost elements, so the shared-grid check is vacuous"); + require_same_on_all_ranks(comm, global_count, + "global element count differs between ranks"); + + for (int d = 0; d < Dim; ++d) { + require_same_on_all_ranks(comm, grid.divisions[d], + "grid divisions differ between ranks"); + require_same_real_on_all_ranks(comm, grid.bot_left[d], 1e-12, + "grid bottom-left differs between ranks"); + require_same_real_on_all_ranks(comm, grid.edge_length[d], 1e-12, + "grid extent differs between ranks"); + // ...and the table is on that grid rather than one of its own making. + require(grid.divisions[d] == expected.divisions[d], + "table did not use the shared routing grid"); + require(grid.bot_left[d] == expected.bot_left[d], + "table did not use the shared routing grid"); + require(grid.edge_length[d] == expected.edge_length[d], + "table did not use the shared routing grid"); + } + + if (comm->rank() == 0) { + std::printf("routing grid: %lld x %lld cells over [%.3f, %.3f]\n", + static_cast(grid.divisions[0]), + static_cast(grid.divisions[Dim - 1]), + grid.bot_left[0], grid.bot_left[0] + grid.edge_length[0]); + } +} + +/// Helper: this rank's element vertices, and which of them it owns. +struct ElementGeom +{ + std::vector, 4>> vertices; // [elem][vert][dim] + std::vector owned; +}; + +ElementGeom GatherElementGeom(Omega_h::Mesh& mesh) +{ + const int dim = mesh.dim(); + const int nverts_per_elem = dim + 1; + const auto n_elems = mesh.nelems(); + const auto elem_verts = Omega_h::HostRead(mesh.ask_elem_verts()); + const auto coords = Omega_h::HostRead(mesh.coords()); + const auto owned = Omega_h::HostRead(mesh.owned(dim)); + + ElementGeom geom; + geom.vertices.resize(static_cast(n_elems)); + geom.owned.resize(static_cast(n_elems)); + for (LO e = 0; e < n_elems; ++e) { + geom.owned[static_cast(e)] = owned[e] != 0; + for (int v = 0; v < nverts_per_elem; ++v) { + const LO vert = elem_verts[e * nverts_per_elem + v]; + for (int d = 0; d < 3; ++d) { + geom.vertices[static_cast(e)][static_cast(v)] + [static_cast(d)] = + d < dim ? coords[vert * dim + d] : 0.0; + } + } + } + return geom; +} + +/// The guarantee: a point inside an element this rank owns must route back to +/// this rank. Sampled over a barycentric lattice so vertices, edge and face +/// interiors, and the element interior are all covered. +template +void check_owned_points_route_to_owner( + Omega_h::Mesh& mesh, const pcms::CandidateRankTable& table, int denom) +{ + auto comm = mesh.comm(); + const int rank = comm->rank(); + const auto geom = GatherElementGeom(mesh); + + LO n_owned = 0; + LO n_probes = 0; + LO n_uncovered = 0; + std::array first_uncovered{}; + for (size_t e = 0; e < geom.vertices.size(); ++e) { + if (!geom.owned[e]) { + continue; + } + ++n_owned; + std::array, Dim + 1> vertices{}; + for (int v = 0; v <= Dim; ++v) { + for (int d = 0; d < Dim; ++d) { + vertices[static_cast(v)][static_cast(d)] = + geom.vertices[e][static_cast(v)][static_cast(d)]; + } + } + for (const auto& sample : SimplexSamples(vertices, denom)) { + ++n_probes; + Omega_h::Vector point; + for (int d = 0; d < Dim; ++d) { + point[d] = sample[static_cast(d)]; + } + const auto candidates = table.CandidatesForPoint(point); + if (std::find(candidates.begin(), candidates.end(), rank) == + candidates.end()) { + if (n_uncovered == 0) { + for (int d = 0; d < 3; ++d) { + first_uncovered[static_cast(d)] = + sample[static_cast(d < Dim ? d : 0)]; + } + } + ++n_uncovered; + } + } + } + + const LO total_probes = comm->allreduce(n_probes, Omega_h_Op::OMEGA_H_SUM); + const LO total_uncovered = + comm->allreduce(n_uncovered, Omega_h_Op::OMEGA_H_SUM); + const LO total_owned = comm->allreduce(n_owned, Omega_h_Op::OMEGA_H_SUM); + + if (n_uncovered > 0) { + std::fprintf(stderr, "[rank %d] first uncovered probe (%.17g, %.17g)\n", + rank, first_uncovered[0], first_uncovered[1]); + } + if (rank == 0) { + std::printf( + "%dD: %lld owned elements, %lld probes, %lld not routed to their owner\n", + Dim, static_cast(total_owned), + static_cast(total_probes), + static_cast(total_uncovered)); + } + require(total_probes > 0, "no probes were generated"); + require(total_uncovered == 0, + "a point inside an owned element did not route to its owner"); +} + +/// The routing decision must be a function of the point alone: every rank holds +/// the same table, so every rank must produce the same candidate list for the +/// same point. A gather that mismatched ranks would show up here. +template +void check_candidate_lists_agree_across_ranks( + Omega_h::Mesh& mesh, const pcms::CandidateRankTable& table) +{ + auto comm = mesh.comm(); + const int nranks = table.NumRanks(); + const auto& grid = table.GetGrid(); + + // Cell centers are a deterministic, rank-independent point set because the + // grid is shared. + const LO num_cells = grid.GetNumCells(); + const LO stride = std::max(1, num_cells / 64); + std::vector cells; + for (LO cell = 0; cell < num_cells; cell += stride) { + cells.push_back(cell); + } + + const size_t width = cells.size() * static_cast(nranks); + std::vector local(width, -1); + for (size_t i = 0; i < cells.size(); ++i) { + const auto bbox = grid.GetCellBBOX(cells[i]); + Omega_h::Vector point; + for (int d = 0; d < Dim; ++d) { + point[d] = bbox.center[static_cast(d)]; + } + const auto candidates = table.CandidatesForPoint(point); + for (size_t j = 0; j < candidates.size(); ++j) { + local[i * static_cast(nranks) + j] = candidates[j]; + } + } + + std::vector all(width * static_cast(nranks), 0); + MPI_Allgather(local.data(), static_cast(width), MPI_INT, all.data(), + static_cast(width), MPI_INT, comm->get_impl()); + for (int r = 1; r < nranks; ++r) { + for (size_t k = 0; k < width; ++k) { + if (all[k] != all[static_cast(r) * width + k]) { + std::fprintf(stderr, + "[rank %d] rank %d computed a different candidate list " + "for sampled cell %zu\n", + comm->rank(), r, k / static_cast(nranks)); + MPI_Abort(MPI_COMM_WORLD, 1); + } + } + } + + if (comm->rank() == 0) { + std::printf("%dD: %zu sampled cells give identical candidate lists on all " + "%d ranks\n", + Dim, cells.size(), nranks); + } +} + +/// The published range must be exactly what the local bounding box implies. +/// +/// Recomputed here from first principles -- box corners binned by the grid's +/// own per-axis rule, with one cell of padding -- rather than by calling the +/// table's helper, so this catches a wrong pad, a wrong axis order, or a range +/// built from the wrong box. It also validates the gather: the range a rank +/// sees for rank `r` must equal what rank `r` computed for itself, which is +/// exactly what a neighbour-scoped gather silently gets wrong by returning a +/// short buffer. +template +void check_published_ranges(Omega_h::Mesh& mesh, + const pcms::CandidateRankTable& table, + Real tolerance_reach) +{ + auto comm = mesh.comm(); + const auto& grid = table.GetGrid(); + + typename pcms::CandidateRankTable::CellRange expected{}; + if (mesh.nverts() == 0) { + expected = pcms::CandidateRankTable::CellRange::Empty(); + } else { + const auto bbox = Omega_h::find_bounding_box(mesh.coords()); + for (int d = 0; d < Dim; ++d) { + const Real lo_coord = bbox.min[d] - tolerance_reach; + const Real hi_coord = bbox.max[d] + tolerance_reach; + expected.lo[d] = std::max(grid.AxisCellIndex(d, lo_coord) - 1, 0); + expected.hi[d] = std::min(grid.AxisCellIndex(d, hi_coord) + 1, + grid.divisions[d] - 1); + } + } + + for (int d = 0; d < Dim; ++d) { + require(table.LocalRange().lo[d] == expected.lo[d], + "published range's lower bound does not match the local bbox"); + require(table.LocalRange().hi[d] == expected.hi[d], + "published range's upper bound does not match the local bbox"); + } + + // Independent expectations, gathered and compared against the table: this is + // the check that a rank-indexed gather went to all ranks rather than only to + // the mesh's communication-graph neighbours. + Kokkos::Array local{}; + for (int d = 0; d < Dim; ++d) { + local[static_cast(d)] = expected.lo[d]; + local[static_cast(Dim + d)] = expected.hi[d]; + } + const auto all = + pcms::AllGatherAllRanks(comm->get_impl(), local.data(), 2 * Dim); + for (int r = 0; r < table.NumRanks(); ++r) { + for (int d = 0; d < Dim; ++d) { + require(table.RangeForRank(r).lo[d] == + all[static_cast(r) * 2 * Dim + d], + "rank r's gathered range does not match rank r's own range"); + require(table.RangeForRank(r).hi[d] == + all[static_cast(r) * 2 * Dim + Dim + d], + "rank r's gathered range does not match rank r's own range"); + } + } + + if (comm->rank() == 0) { + std::printf("%dD: published ranges verified against each rank's own bbox " + "and against a rank-indexed gather\n", + Dim); + } +} + +/// Widening the published ranges by the containment tolerance may only ever add +/// candidate ranks, and must actually add them once the tolerance exceeds a +/// cell size. Also checks the case the widening exists for: a point just +/// outside an element, which the search may still accept, must route to that +/// element's owner. +template +void check_tolerance_widening(Omega_h::Mesh& mesh) +{ + auto comm = mesh.comm(); + const int rank = comm->rank(); + + // Use a grid fine enough that a rank's bare range does not already cover + // every cell. On a coarse grid with few ranks both ranges can span the whole + // domain, and then widening cannot add a candidate and the property is + // unobservable. + const auto shared = pcms::CandidateRankTable::SharedRoutingGrid(mesh); + const LO fine_divisions = 8 * shared.divisions[0]; + const auto fine_grid = + pcms::CreateUniformGridFromMesh(mesh, fine_divisions); + const Real cell_size = fine_grid.edge_length[0] / fine_grid.divisions[0]; + + // Widen by several cells so the effect is unambiguous, and probe a face at + // half that distance: more than the single cell of padding a bare range + // already has, so a bare range must miss the probe while a widened one + // catches it. + const Real reach = 4.0 * cell_size; + const pcms::CandidateRankTable strict(mesh, fine_grid, 0.0); + const pcms::CandidateRankTable wide(mesh, fine_grid, reach); + + LO pairs_strict = 0; + LO pairs_wide = 0; + LO cells_grown = 0; + LO subset_violations = 0; + for (LO cell = 0; cell < strict.GetGrid().GetNumCells(); ++cell) { + const auto s = strict.CandidatesForCell(cell); + const auto w = wide.CandidatesForCell(cell); + pairs_strict += static_cast(s.size()); + pairs_wide += static_cast(w.size()); + if (w.size() > s.size()) { + ++cells_grown; + } + // Both lists are ascending, so this is a subset test. + if (!std::includes(w.begin(), w.end(), s.begin(), s.end())) { + ++subset_violations; + } + } + + // The case the widening exists for: a point the search may still accept + // because it lies just outside this rank's geometry. Probe just beyond each + // face of the rank's own bounding box, at half the tolerance, so it is inside + // the tolerance-widened box and must route back here. Probbing relative to + // the box rather than to an element makes this deterministic: an interior + // element's face can sit far inside the box, where the bare range already + // covers the probe and no widening is needed. + LO n_face_probes = 0; + LO n_missed_wide = 0; + LO n_missed_strict = 0; + if (mesh.nverts() > 0) { + const auto bbox = Omega_h::find_bounding_box(mesh.coords()); + for (int d = 0; d < Dim; ++d) { + for (int side = 0; side < 2; ++side) { + Omega_h::Vector probe; + for (int e = 0; e < Dim; ++e) { + probe[e] = (bbox.min[e] + bbox.max[e]) / 2.0; + } + // Half a tolerance outside a face of the box: within `reach` of the + // rank's geometry, so the widened range must cover it. + probe[d] = + side == 0 ? bbox.min[d] - reach / 2.0 : bbox.max[d] + reach / 2.0; + ++n_face_probes; + const auto wide_candidates = wide.CandidatesForPoint(probe); + if (std::find(wide_candidates.begin(), wide_candidates.end(), rank) == + wide_candidates.end()) { + ++n_missed_wide; + } + const auto strict_candidates = strict.CandidatesForPoint(probe); + if (std::find(strict_candidates.begin(), strict_candidates.end(), + rank) == strict_candidates.end()) { + ++n_missed_strict; + } + } + } + } + + const LO total_pairs_strict = + comm->allreduce(pairs_strict, Omega_h_Op::OMEGA_H_SUM); + const LO total_pairs_wide = + comm->allreduce(pairs_wide, Omega_h_Op::OMEGA_H_SUM); + const LO total_grown = comm->allreduce(cells_grown, Omega_h_Op::OMEGA_H_SUM); + const LO total_violations = + comm->allreduce(subset_violations, Omega_h_Op::OMEGA_H_SUM); + const LO total_probes = + comm->allreduce(n_face_probes, Omega_h_Op::OMEGA_H_SUM); + const LO total_missed_wide = + comm->allreduce(n_missed_wide, Omega_h_Op::OMEGA_H_SUM); + const LO total_missed_strict = + comm->allreduce(n_missed_strict, Omega_h_Op::OMEGA_H_SUM); + + if (rank == 0) { + std::printf( + "%dD tolerance widening (reach %.4g = %.1f cells): %lld -> %lld " + "candidate " + "pairs, %lld grown cells; %lld probes outside a face, missed by widened " + "ranges %lld, by bare ranges %lld\n", + Dim, static_cast(reach), static_cast(reach / cell_size), + static_cast(total_pairs_strict), + static_cast(total_pairs_wide), + static_cast(total_grown), static_cast(total_probes), + static_cast(total_missed_wide), + static_cast(total_missed_strict)); + } + + require(total_violations == 0, "widening removed a candidate rank"); + // The widened range must contain the bare range on every axis, and the + // widened candidate list must contain the bare one. Both hold whether or not + // the partition happens to saturate the grid. + for (int d = 0; d < Dim; ++d) { + require(wide.LocalRange().lo[d] <= strict.LocalRange().lo[d], + "widening moved a range's lower bound outward"); + require(wide.LocalRange().hi[d] >= strict.LocalRange().hi[d], + "widening shrank a range's upper bound"); + } + require(total_missed_wide == 0, + "a point just outside an owned element was not routed to its owner"); + + /// Widening is unobservable exactly when every non-empty rank's range already + /// spans the whole grid: there is no cell left to add, so the candidate lists + /// cannot grow no matter how large the tolerance. A range can span the grid + /// because the partition interleaves (each rank's coordinate extent covers + /// the domain) or because the grid is coarse relative to the decomposition. + bool any_rank_can_grow = false; + for (int r = 0; r < strict.NumRanks(); ++r) { + const auto& range = strict.RangeForRank(r); + if (range.IsEmpty()) { + continue; + } + bool spans = true; + for (int d = 0; d < Dim; ++d) { + if (range.lo[d] != 0 || range.hi[d] != fine_grid.divisions[d] - 1) { + spans = false; + } + } + if (!spans) { + any_rank_can_grow = true; + } + } + if (rank == 0) { + std::printf("%dD widening is %s on this partition\n", Dim, + any_rank_can_grow ? "observable" : "unobservable"); + } + if (any_rank_can_grow) { + require(total_pairs_wide > total_pairs_strict, + "widening by several cells did not add any candidate"); + require(total_grown > 0, "widening did not grow any cell's candidate list"); + } +} + +/// A rank with no local geometry can find nothing, so it must publish an empty +/// range and never be a routing candidate. +template +void check_empty_ranks(Omega_h::Mesh& mesh, + const pcms::CandidateRankTable& table) +{ + auto comm = mesh.comm(); + const int nranks = table.NumRanks(); + const auto nverts_all = + pcms::AllGatherAllRanks(comm->get_impl(), static_cast(mesh.nverts())); + + LO n_empty = 0; + for (int r = 0; r < nranks; ++r) { + if (nverts_all[r] != 0) { + continue; + } + ++n_empty; + require(table.RangeForRank(r).IsEmpty(), + "a rank with no geometry published a non-empty range"); + for (LO cell = 0; cell < table.GetGrid().GetNumCells(); ++cell) { + const auto candidates = table.CandidatesForCell(cell); + require(std::find(candidates.begin(), candidates.end(), r) == + candidates.end(), + "an empty rank is a routing candidate"); + } + } + + // `n_empty` is already global: every rank holds the same gathered view, so it + // must not be reduced again. + if (comm->rank() == 0) { + std::printf("%dD: %lld rank(s) with no local geometry\n", Dim, + static_cast(n_empty)); + } +} + +// ---- 2. local pass, 3. full search ---- + +/// A centroid lies strictly inside its own element, so the geometrically right +/// answer is that element. That makes ownership of the answer unambiguous: an +/// owned element must be answerable here, a ghost must not be. +template +void check_local_classification(Omega_h::Mesh& mesh) +{ + using Status = pcms::DistributedPointSearch::LocalStatus; + + const auto n_elems = mesh.nelems(); + const auto n_owned = pcms::OwnedElementCount(mesh); + require(mesh.nghost_layers() > 0, "mesh is not ghosted"); + + const auto divisions = pcms::DivisionsPerAxisForGlobalMesh(mesh); + const auto local_search = distributed_test::MakeLocalSearch(mesh, + divisions); + const pcms::DistributedPointSearch search{mesh, local_search}; + require(search.HasGhostLayer(), "distributed search missed the ghost layer"); + + const auto centroids = + Omega_h::HostRead(pcms::get_entity_centroids(mesh, Dim)); + Kokkos::View points("centroids", n_elems); + auto points_h = Kokkos::create_mirror_view(points); + for (LO i = 0; i < n_elems; ++i) { + for (int d = 0; d < Dim; ++d) { + points_h(i, d) = centroids[i * Dim + d]; + } + } + Kokkos::deep_copy(points, points_h); + + const auto classified = search.ClassifyLocally(points); + require(classified.size() == static_cast(n_elems), + "one classification per point expected"); + + LO n_here = 0; + LO n_elsewhere = 0; + LO n_missing = 0; + LO n_owned_wrong = 0; + LO n_ghost_wrong = 0; + for (LO i = 0; i < n_elems; ++i) { + const auto& c = classified[static_cast(i)]; + const bool owned = search.OwnsLocalElement(i); + switch (c.status) { + case Status::LOCATED_HERE: + ++n_here; + // Only an owned element may be answered here, and it must be the + // element the point came from. + if (!owned || c.element_id != i) { + ++n_owned_wrong; + } + break; + case Status::LOCATED_ELSEWHERE: + ++n_elsewhere; + // A ghost answer must name a non-owned element, and element_id must be + // a valid local element. + if (owned || c.element_id < 0 || c.element_id >= n_elems) { + ++n_ghost_wrong; + } + break; + case Status::NOT_FOUND_HERE: ++n_missing; break; + } + } + + const LO total_here = mesh.comm()->allreduce(n_here, Omega_h_Op::OMEGA_H_SUM); + const LO total_elsewhere = + mesh.comm()->allreduce(n_elsewhere, Omega_h_Op::OMEGA_H_SUM); + const LO total_missing = + mesh.comm()->allreduce(n_missing, Omega_h_Op::OMEGA_H_SUM); + const LO total_owned_wrong = + mesh.comm()->allreduce(n_owned_wrong, Omega_h_Op::OMEGA_H_SUM); + const LO total_ghost_wrong = + mesh.comm()->allreduce(n_ghost_wrong, Omega_h_Op::OMEGA_H_SUM); + const LO total_owned = + mesh.comm()->allreduce(n_owned, Omega_h_Op::OMEGA_H_SUM); + // A rank whose partition is isolated holds no ghosts and simply contributes + // nothing here; what matters is that the decomposition as a whole exercised + // the ghost classification. That is checked below via `total_elsewhere`. + + if (mesh.comm()->rank() == 0) { + std::printf( + "classification of %lld owned element centroids: %lld answerable here, " + "%lld answerable elsewhere (ghost), %lld not found, %lld/%lld errors\n", + static_cast(total_owned), static_cast(total_here), + static_cast(total_elsewhere), + static_cast(total_missing), + static_cast(total_owned_wrong + total_ghost_wrong), + static_cast(total_here + total_elsewhere + total_missing)); + } + + require( + total_owned_wrong == 0, + "an element not owned by this rank was classified as answerable here"); + require(total_ghost_wrong == 0, + "a ghost answer did not name a non-owned element"); + require(total_here == total_owned, + "not every owned element was answerable on its owner"); + require( + total_missing == 0, + "a local element centroid was not found at all by the rank holding it"); + require(total_elsewhere > 0, + "no point was classified as answerable elsewhere, so the ownership " + "distinction is not being exercised"); +} + +/// Every rank locates every owned element centroid in the decomposition. +/// +/// Each point is therefore answered by the rank that owns it, locally, and by +/// every other rank, through routing. Both paths must agree on the owning rank +/// and on the element's global id -- that is the cross-path consistency the +/// merge rule exists to guarantee. +template +void check_distributed_locate(Omega_h::Mesh& mesh) +{ + using ResultStatus = pcms::DistributedPointSearch::ResultStatus; + + const int rank = mesh.comm()->rank(); + require(mesh.has_tag(Dim, "global"), + "test needs element global ids to check which element was found"); + + // Every owned element centroid in the decomposition, plus one point far + // outside the domain on every axis. + auto gathered = distributed_test::GatherOwnedCentroids(mesh); + const int total = gathered.total; + const auto& global_ids = gathered.global_id; + const auto& expected_owner = gathered.contributor; + const LO n_local = static_cast( + std::count(expected_owner.begin(), expected_owner.end(), rank)); + for (int d = 0; d < Dim; ++d) { + gathered.coords.push_back(2.0); + } + const LO n_points = static_cast(total) + 1; + const auto points = distributed_test::ToPointView( + gathered.coords, "global element centroids"); + + const auto divisions = pcms::DivisionsPerAxisForGlobalMesh(mesh); + const auto local_search = distributed_test::MakeLocalSearch(mesh, + divisions); + pcms::DistributedPointSearch search{mesh, local_search}; + + const auto results = search.Locate(points); + require(results.size() == static_cast(n_points), + "one result per point expected"); + + LO n_wrong_owner = 0; + LO n_wrong_element = 0; + LO n_not_found = 0; + LO n_bad_bary = 0; + for (int i = 0; i < total; ++i) { + const auto& r = results[static_cast(i)]; + if (r.status != ResultStatus::LOCATED) { + ++n_not_found; + continue; + } + if (r.owner_rank != expected_owner[static_cast(i)]) { + ++n_wrong_owner; + } + if (r.global_element_id != global_ids[static_cast(i)]) { + ++n_wrong_element; + } + // A centroid is strictly inside its element, so every barycentric weight is + // positive, and the weights sum to one. This also catches a result + // assembled from the wrong slots. + Real weight = 0; + for (int d = 0; d <= Dim; ++d) { + weight += r.barycentric[d]; + if (r.barycentric[d] <= 0) { + ++n_bad_bary; + } + } + if (std::fabs(weight - 1.0) > 1e-8) { + ++n_bad_bary; + } + } + // A point no rank placed must not name an owner or an element, so a caller + // cannot mistake it for a location. + const auto& outside = results[static_cast(total)]; + const bool outside_not_found = outside.status == ResultStatus::NOT_FOUND && + outside.owner_rank == -1 && + outside.element_id == -1 && + outside.global_element_id == -1; + + const LO t_not_found = + mesh.comm()->allreduce(n_not_found, Omega_h_Op::OMEGA_H_SUM); + const LO t_wrong_owner = + mesh.comm()->allreduce(n_wrong_owner, Omega_h_Op::OMEGA_H_SUM); + const LO t_wrong_element = + mesh.comm()->allreduce(n_wrong_element, Omega_h_Op::OMEGA_H_SUM); + const LO t_bad_bary = + mesh.comm()->allreduce(n_bad_bary, Omega_h_Op::OMEGA_H_SUM); + const LO t_points = mesh.comm()->allreduce(n_points, Omega_h_Op::OMEGA_H_SUM); + const LO t_routed_away = mesh.comm()->allreduce( + static_cast(total) - n_local, + Omega_h_Op::OMEGA_H_SUM); + + if (rank == 0) { + std::printf( + "located %lld gathered points (%lld queries per rank left their rank): " + "%lld not found, %lld wrong owner, %lld wrong element, %lld bad " + "barycentric\n", + static_cast(t_points), static_cast(t_routed_away), + static_cast(t_not_found), + static_cast(t_wrong_owner), + static_cast(t_wrong_element), + static_cast(t_bad_bary)); + } + + require(outside_not_found, + "a point outside the domain was reported as located, or its result " + "named an owner or element"); + require(t_not_found == 0, "a point inside the domain was not located"); + require(t_wrong_owner == 0, + "a point was attributed to the wrong owning rank"); + require(t_wrong_element == 0, "a point resolved to the wrong element"); + require(t_bad_bary == 0, + "a located point has non-positive barycentric weight"); + require(t_routed_away > 0, + "no query left its rank, so routing was not exercised"); + + // Determinism and caching: the same batch twice must give the same answer, + // and a second pass over a warm cache must not re-search anything it + // resolved. + std::vector::Result> cache; + const auto pass_one = search.Locate(points, &cache); + const auto routed_one = search.LastRoutedPointCount(); + const auto pass_two = search.Locate(points, &cache); + const auto routed_two = search.LastRoutedPointCount(); + + LO n_unresolved_results = 0; + for (const auto& r : pass_one) { + if (r.status != ResultStatus::LOCATED) { + ++n_unresolved_results; + } + } + + LO n_unstable = 0; + for (size_t i = 0; i < results.size(); ++i) { + if (pass_two[i].status != results[i].status || + pass_two[i].owner_rank != results[i].owner_rank || + pass_two[i].global_element_id != results[i].global_element_id) { + ++n_unstable; + } + } + + if (rank == 0) { + std::printf("routed %zu points on the first pass, %zu on the second " + "(only unresolved points are retried, and points outside the " + "mesh bounding box are never routed)\n", + routed_one, routed_two); + } + require(routed_one > 0, "the first pass routed nothing"); + // The only point still unresolved on a warm cache is the one outside the + // grid, and that one needs no message, so a warm pass routes nothing at all. + require(n_unresolved_results > 0, "the out-of-domain probe was located"); + require(routed_two == 0, "a cached pass still routed points"); + require(routed_two < routed_one, "caching did not reduce the routing work"); + require(mesh.comm()->allreduce(n_unstable, Omega_h_Op::OMEGA_H_SUM) == 0, + "a cached pass gave a different answer"); + + // A batch of points outside the global bounding box needs no communication at + // all: no element can contain them, so not a single query is routed. + // One past each end of the first axis, then one past the far corner. + Kokkos::View far_points("far outside", 3); + auto far_h = Kokkos::create_mirror_view(far_points); + for (int d = 0; d < Dim; ++d) { + far_h(0, d) = (d == 0) ? 2.0 : 0.5; + far_h(1, d) = (d == 0) ? -1.0 : 0.5; + far_h(2, d) = 0.5; + } + far_h(2, Dim - 1) = 5.0; + Kokkos::deep_copy(far_points, far_h); + const auto far_results = search.Locate(far_points); + const std::size_t far_routed = search.LastRoutedPointCount(); + LO n_far_located = 0; + for (const auto& r : far_results) { + if (r.status != ResultStatus::NOT_FOUND) { + ++n_far_located; + } + } + require(far_routed == 0, "a point outside the mesh bounding box was routed"); + require(mesh.comm()->allreduce(n_far_located, Omega_h_Op::OMEGA_H_SUM) == 0, + "a point outside the mesh bounding box was reported as located"); +} + +/// The global search grid's resolution can be chosen by the caller. Any +/// resolution must give the same answers -- it only changes how many candidate +/// ranks a point is sent to -- and a resolution that differs between ranks must +/// be rejected on every rank, since cell ids would then mean different things. +template +void check_grid_divisions(Omega_h::Mesh& mesh) +{ + using Search = pcms::DistributedPointSearch; + const auto gathered = distributed_test::GatherOwnedCentroids(mesh); + const auto points = distributed_test::ToPointView(gathered.coords); + const auto local_search = distributed_test::MakeLocalSearch( + mesh, pcms::DivisionsPerAxisForGlobalMesh(mesh)); + + Search by_default{mesh, local_search}; + const auto expected = by_default.Locate(points); + + LO n_differ = 0; + for (const LO divisions : {LO{3}, LO{37}}) { + Search custom{mesh, local_search, 0.0, divisions}; + for (int d = 0; d < Dim; ++d) { + require(custom.GetRoutingTable().GetGrid().divisions[d] == divisions, + "the search did not use the requested grid resolution"); + } + const auto got = custom.Locate(points); + for (size_t i = 0; i < got.size(); ++i) { + if (got[i].status != expected[i].status || + got[i].owner_rank != expected[i].owner_rank || + got[i].global_element_id != expected[i].global_element_id) { + ++n_differ; + } + } + } + require(mesh.comm()->allreduce(n_differ, Omega_h_Op::OMEGA_H_SUM) == 0, + "a custom grid resolution changed where a point was located"); + + bool threw = false; + try { + Search mismatched{mesh, local_search, 0.0, 10 + mesh.comm()->rank()}; + } catch (const std::invalid_argument&) { + threw = true; + } + require(threw, "a grid resolution that differs between ranks was accepted"); + + if (mesh.comm()->rank() == 0) { + std::printf("%dD: custom grid resolutions give the same answers; a " + "mismatched one is rejected\n", + Dim); + } +} + +template +void run_all(Omega_h::CommPtr world) +{ + auto mesh = distributed_test::BuildGhostedBox(world); + if (world->rank() == 0) { + std::printf("--- %dD, %d ranks ---\n", Dim, world->size()); + } + + const pcms::CandidateRankTable table(mesh); + check_shared_grid(mesh, table); + check_published_ranges(mesh, table, 0.0); + check_owned_points_route_to_owner(mesh, table, /*denom=*/Dim == 2 ? 4 : 3); + check_candidate_lists_agree_across_ranks(mesh, table); + check_tolerance_widening(mesh); + check_empty_ranks(mesh, table); + + check_local_classification(mesh); + check_distributed_locate(mesh); + check_grid_divisions(mesh); +} + +} // namespace + +int main(int argc, char** argv) +{ + MPI_Init(&argc, &argv); + { + Omega_h::Library library(&argc, &argv); + auto world = library.world(); + require(world->size() >= 2, "test requires at least 2 MPI ranks"); + + run_all<2>(world); + run_all<3>(world); + + if (world->rank() == 0) { + std::printf("distributed search test passed\n"); + } + } + MPI_Finalize(); + return 0; +} From 20c933e9da7efa8ebea0a3f3a7c03e44c735a1ce Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Tue, 27 Oct 2026 17:25:50 -0400 Subject: [PATCH 4/7] extend point search tests --- src/pcms/utility/uniform_grid.h | 19 +- test/test_point_search.cpp | 348 ++++++++++++++++++++++++++++++++ test/test_uniform_grid.cpp | 96 +++++++++ 3 files changed, 453 insertions(+), 10 deletions(-) diff --git a/src/pcms/utility/uniform_grid.h b/src/pcms/utility/uniform_grid.h index 81b7a612b..17ba35a8c 100644 --- a/src/pcms/utility/uniform_grid.h +++ b/src/pcms/utility/uniform_grid.h @@ -131,21 +131,20 @@ struct UniformGrid return true; } + /// Per-axis indices of a linear cell ID, in reversed order (z, y, x), the + /// order `GetCellIndex` takes. The exact inverse of `GetCellIndex`, which + /// numbers cells as `idx = x + Nx * (y + Ny * z)`. + /// + /// This is a numbering of cells, not a memory layout: it is the same on host + /// and device, whatever layout a Kokkos view uses. [[nodiscard]] KOKKOS_INLINE_FUNCTION Kokkos::Array GetDimensionedIndex(LO idx) const { - LO stride = 1; - for (std::size_t i = 0; i < dim - 1; ++i) { - stride *= divisions[i]; - } Kokkos::Array result; - - for (size_t i = 0; i < dim; ++i) { - result[i] = idx / stride; - idx -= result[i] * stride; - stride /= divisions[i]; + for (size_t axis = 0; axis < dim; ++axis) { + result[dim - 1 - axis] = idx % divisions[axis]; + idx /= divisions[axis]; } - return result; } diff --git a/test/test_point_search.cpp b/test/test_point_search.cpp index 6a5d6b996..482cc095d 100644 --- a/test/test_point_search.cpp +++ b/test/test_point_search.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -527,3 +528,350 @@ TEST_CASE("point search encodes outside points unambiguously", REQUIRE(owners_h(2) == 0); } } + +// Baseline for the reused local search: on a mesh with no ghost layer every +// element is owned and every element centroid must localize to its own element. +// The distributed search treats a local answer as authoritative only when this +// holds, so it is the contract the distributed layer builds on. +TEST_CASE("ghost-free search localizes every element centroid") +{ + using pcms::GridPointSearch2D; + auto lib = Omega_h::Library{}; + auto world = lib.world(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1, 1, 0, 4, 4, 0, false); + REQUIRE(mesh.nghost_layers() == 0); + + const auto n_elems = mesh.nelems(); + const auto owned = Omega_h::HostRead(mesh.owned(2)); + REQUIRE(static_cast(owned.size()) == n_elems); + for (Omega_h::LO i = 0; i < n_elems; ++i) { + REQUIRE(owned[i] != 0); + } + + const auto n = pcms::DivisionsPerAxisForGlobalMesh<2>(mesh); + GridPointSearch2D search{mesh, n, n}; + const auto centroids = + Omega_h::HostRead(pcms::get_entity_centroids(mesh, 2)); + Kokkos::View points("query points", n_elems); + auto points_h = Kokkos::create_mirror_view(points); + for (Omega_h::LO i = 0; i < n_elems; ++i) { + points_h(i, 0) = centroids[i * 2 + 0]; + points_h(i, 1) = centroids[i * 2 + 1]; + } + Kokkos::deep_copy(points, points_h); + + const auto owners_h = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, search.GetOwningElementIds(search(points))); + for (Omega_h::LO i = 0; i < n_elems; ++i) { + REQUIRE(owners_h(i) == i); + } +} + +// The distributed search decides what this rank may answer for by asking +// whether the element the local search reported is one it owns. On a mesh with +// no ghost layer every local element is owned, so every located point must be +// answerable here -- which is the baseline the multi-rank classification test +// builds on. +TEST_CASE("distributed search classifies local answers by ownership") +{ + using pcms::DistributedPointSearch; + using Status = DistributedPointSearch<2>::LocalStatus; + + auto lib = Omega_h::Library{}; + auto world = lib.world(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1, 1, 0, 4, 4, 0, false); + REQUIRE(mesh.nghost_layers() == 0); + + const auto n = pcms::DivisionsPerAxisForGlobalMesh<2>(mesh); + pcms::GridPointSearch2D local_search{mesh, n, n}; + const DistributedPointSearch<2> search{mesh, local_search}; + + REQUIRE_FALSE(search.HasGhostLayer()); + REQUIRE(search.GetMesh().nelems() == mesh.nelems()); + + // Query the centroid of every element plus one point outside the domain. + const auto n_elems = mesh.nelems(); + const auto centroids = + Omega_h::HostRead(pcms::get_entity_centroids(mesh, 2)); + Kokkos::View points("query points", n_elems + 1); + auto points_h = Kokkos::create_mirror_view(points); + for (Omega_h::LO i = 0; i < n_elems; ++i) { + points_h(i, 0) = centroids[i * 2 + 0]; + points_h(i, 1) = centroids[i * 2 + 1]; + } + // Outside the mesh and outside every grid cell that holds an element. + points_h(n_elems, 0) = 2.0; + points_h(n_elems, 1) = 2.0; + Kokkos::deep_copy(points, points_h); + + const auto classified = search.ClassifyLocally(points); + REQUIRE(classified.size() == static_cast(n_elems + 1)); + + for (Omega_h::LO i = 0; i < n_elems; ++i) { + const auto& c = classified[static_cast(i)]; + REQUIRE(c.status == Status::LOCATED_HERE); + // An element centroid lies strictly inside its element, so the reported + // element is that element and the barycentric coordinates are positive. + REQUIRE(c.element_id == i); + for (int d = 0; d < 3; ++d) { + REQUIRE(c.barycentric[d] > 0); + } + } + + // With no ghosts nothing can be "located elsewhere": either this rank answers + // or the point is outside what it holds. + for (const auto& c : classified) { + REQUIRE(c.status != Status::LOCATED_ELSEWHERE); + } + REQUIRE(classified[static_cast(n_elems)].status == + Status::NOT_FOUND_HERE); + REQUIRE(classified[static_cast(n_elems)].element_id == -1); +} +// A predicate chosen at the map-building call must govern *both* halves: the +// range of cells visited and the test applied to each. Checking the map against +// a brute-force build made with that same predicate is what catches the case +// where one half ignores the argument -- the fast map would then disagree. +// +// Only predicates instantiated in point_search.cpp can be named from here; the +// set is small and internal, which is why the map machinery no longer lives in +// the header. See the instantiations at the end of that file. +TEST_CASE("candidate map accepts a predicate other than the default") +{ + auto lib = Omega_h::Library{}; + auto world = lib.world(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1, 1, 0, 10, 10, 0, false); + + Kokkos::View grid_d("uniform grid"); + auto grid_h = Kokkos::create_mirror_view(grid_d); + grid_h(0) = Uniform2DGrid{ + .edge_length{1, 1}, .bot_left = {0, 0}, .divisions = {10, 10}}; + Kokkos::deep_copy(grid_d, grid_h); + const auto ncells = grid_h(0).GetNumCells(); + + // SimplexVsBbox is the 3D default; in 2D it is the conservative bounding-box + // test, which is a legitimate thing to want. No other caller asks for it in + // 2D. + auto conservative = + pcms::detail::construct_intersection_map<2, pcms::SimplexVsBbox>( + mesh, grid_d, ncells); + require_maps_identical( + conservative, + pcms::detail::construct_intersection_map_reference<2, pcms::SimplexVsBbox>( + mesh, grid_d, ncells)); + + // The conservative test accepts everything the refined default accepts and + // more, wherever an element's box reaches a cell its triangle does not. If + // the predicate were ignored, the two maps would be identical. + auto refined = + pcms::detail::construct_intersection_map_2d(mesh, grid_d, ncells); + REQUIRE(conservative.entries.size() > refined.entries.size()); +} + +// The search itself must accept a predicate the same way, chosen once when it +// is built. The predicate must be one instantiated in point_search.cpp. +TEST_CASE("point search accepts a predicate other than its default") +{ + using pcms::GridPointSearch2D; + auto lib = Omega_h::Library{}; + auto world = lib.world(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1, 1, 0, 10, 10, 0, false); + + auto tolerances = + GridPointSearch2D::PointSearchTolerances{"point search 2d tolerances"}; + auto tolerances_h = Kokkos::create_mirror_view(tolerances); + tolerances_h(0) = 0.01; + tolerances_h(1) = 0.01; + Kokkos::deep_copy(tolerances, tolerances_h); + + GridPointSearch2D default_search{mesh, 10, 10, tolerances}; + // The conservative 2D predicate instead of this search's default. + GridPointSearch2D custom_search{mesh, 10, 10, tolerances, + pcms::SimplexVsBbox{}}; + + Kokkos::View points("test_points", 4); + auto points_h = Kokkos::create_mirror_view(points); + points_h(0, 0) = 0.05; + points_h(0, 1) = 0.02; + points_h(1, 0) = 0.55; + points_h(1, 1) = 0.54; + points_h(2, 0) = 0.9; + points_h(2, 1) = 0.9; + points_h(3, 0) = 1.0; + points_h(3, 1) = 1.0; + Kokkos::deep_copy(points, points_h); + + auto default_results = default_search(points); + auto custom_results = custom_search(points); + auto default_h = Kokkos::create_mirror_view(default_results); + auto custom_h = Kokkos::create_mirror_view(custom_results); + Kokkos::deep_copy(default_h, default_results); + Kokkos::deep_copy(custom_h, custom_results); + + // The query tests point containment in each candidate, so the predicate only + // decides which elements are offered. For points inside the mesh both + // searches must therefore answer identically: the value-constructor path is a + // candidate-set choice, not a change of semantics. + for (size_t i = 0; i < 4; ++i) { + CAPTURE(i); + REQUIRE(custom_h(i).dimensionality == default_h(i).dimensionality); + REQUIRE(custom_h(i).element_id == default_h(i).element_id); + } +} + +// The 3D search takes a division count per axis. Unequal counts must locate +// points exactly as equal counts do: elements are binned by each cell's box, so +// a cell box computed for the wrong axis would drop candidates silently. +TEST_CASE("3D point search with unequal divisions per axis", + "[point_search][3d]") +{ + using pcms::GridPointSearch3D; + auto lib = Omega_h::Library{}; + auto world = lib.world(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1, 1, 1, 4, 4, 4, false); + GridPointSearch3D search{mesh, 3, 5, 7}; + + const auto n_elems = mesh.nelems(); + const auto centroids = + Omega_h::HostRead(pcms::get_entity_centroids(mesh, 3)); + Kokkos::View points("query points", n_elems); + auto points_h = Kokkos::create_mirror_view(points); + for (Omega_h::LO i = 0; i < n_elems; ++i) { + for (int d = 0; d < 3; ++d) { + points_h(i, d) = centroids[i * 3 + d]; + } + } + Kokkos::deep_copy(points, points_h); + + const auto owners_h = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace{}, search.GetOwningElementIds(search(points))); + Omega_h::LO n_wrong = 0; + for (Omega_h::LO i = 0; i < n_elems; ++i) { + n_wrong += owners_h(i) != i; + } + CAPTURE(n_wrong, n_elems); + REQUIRE(n_wrong == 0); +} + +// The grid and overlap layer under the point search is written for any +// dimension, though the search classes themselves are 2D and 3D only. These +// checks pin that down above 3D. Every grid has a different division count on +// each axis, and every probe uses different indices per axis, so an axis mix-up +// cannot hide behind a symmetric case. +namespace +{ + +/// Cells of `grid` whose id, per-axis indices and box disagree, counted on the +/// device: the same functions the search runs in its kernels. +template +pcms::LO CountCellRoundTripFailuresOnDevice(const pcms::UniformGrid& grid) +{ + pcms::LO failures = 0; + Kokkos::parallel_reduce( + "cell round trip", grid.GetNumCells(), + KOKKOS_LAMBDA(pcms::LO cell, pcms::LO & n) { + const auto axis = grid.AxisIndicesFromCellIndex(cell); + bool ok = grid.CellIndexFromAxisIndices(axis) == cell; + Omega_h::Vector center; + const auto box = grid.GetCellBBOX(cell); + for (unsigned d = 0; d < Dim; ++d) { + ok = ok && axis[d] >= 0 && axis[d] < grid.divisions[d]; + center[d] = box.center[d]; + } + ok = ok && grid.ClosestCellID(center) == cell; + n += ok ? 0 : 1; + }, + failures); + return failures; +} + +} // namespace + +TEST_CASE("grid and overlap layer above 3D", "[point_search][nd]") +{ + using pcms::LO; + + SECTION("4D grid: cell ids, per-axis indices and cell boxes agree") + { + // Cell size 1 on every axis, so a cell's box is easy to state. + const pcms::UniformGrid<4> grid{.edge_length = {2.0, 3.0, 4.0, 5.0}, + .bot_left = {0.0, 0.0, 0.0, 0.0}, + .divisions = {2, 3, 4, 5}}; + REQUIRE(grid.GetNumCells() == 120); + + // Ids number cells x-fastest: x + Nx * (y + Ny * (z + Nz * w)). + const Kokkos::Array last{1, 2, 3, 4}; + REQUIRE(grid.CellIndexFromAxisIndices(last) == 119); + const Kokkos::Array probe{1, 0, 2, 3}; + const LO probe_id = 1 + 2 * (0 + 3 * (2 + 4 * 3)); + REQUIRE(grid.CellIndexFromAxisIndices(probe) == probe_id); + const auto back = grid.AxisIndicesFromCellIndex(probe_id); + for (int d = 0; d < 4; ++d) { + REQUIRE(back[d] == probe[d]); + } + const auto box = grid.GetCellBBOX(probe_id); + for (int d = 0; d < 4; ++d) { + REQUIRE(box.center[d] == Catch::Approx(probe[d] + 0.5)); + REQUIRE(box.half_width[d] == Catch::Approx(0.5)); + } + + REQUIRE(CountCellRoundTripFailuresOnDevice(grid) == 0); + } + + SECTION("5D grid: every cell round-trips on the device") + { + const pcms::UniformGrid<5> grid{.edge_length = {1.0, 2.0, 3.0, 2.0, 1.0}, + .bot_left = {0.0, -1.0, 2.0, 0.5, 0.0}, + .divisions = {2, 3, 4, 3, 2}}; + REQUIRE(grid.GetNumCells() == 144); + REQUIRE(CountCellRoundTripFailuresOnDevice(grid) == 0); + } + + SECTION("4D candidate cells keep each axis separate") + { + const pcms::UniformGrid<4> grid{.edge_length = {2.0, 3.0, 4.0, 5.0}, + .bot_left = {0.0, 0.0, 0.0, 0.0}, + .divisions = {2, 3, 4, 5}}; + // Inside cell (0, 1, 2, 3); padded by one cell, clamped to the grid. + const AABBox<4> box{.center = {0.5, 1.5, 2.5, 3.5}, + .half_width = {0.25, 0.25, 0.25, 0.25}}; + const auto range = + pcms::candidate_cells<4, pcms::SimplexVsBbox>(box, grid); + const Kokkos::Array lo{0, 0, 1, 2}; + const Kokkos::Array hi{1, 2, 3, 4}; + for (int d = 0; d < 4; ++d) { + CAPTURE(d); + REQUIRE(range.lo[d] == lo[d]); + REQUIRE(range.hi[d] == hi[d]); + } + REQUIRE(range.Contains(Kokkos::Array{0, 1, 2, 3})); + REQUIRE_FALSE(range.Contains(Kokkos::Array{0, 1, 2, 1})); + } + + SECTION("4D simplex-vs-box overlap needs every axis to overlap") + { + const AABBox<4> box{.center = {1.0, 1.0, 1.0, 1.0}, + .half_width = {0.5, 0.5, 0.5, 0.5}}; + // A simplex spanning [0.8, 1.2] on axes 0-2. On axis 3 it either reaches + // into the box or stops short of it. + const auto simplex = [](pcms::Real axis3_lo, pcms::Real axis3_hi) { + Omega_h::Matrix<4, 5> coords{}; + for (int v = 0; v < 5; ++v) { + for (int d = 0; d < 3; ++d) { + coords(d, v) = (v % 2 == 0) ? 0.8 : 1.2; + } + coords(3, v) = (v == 0) ? axis3_lo : axis3_hi; + } + return coords; + }; + REQUIRE(pcms::Overlap<4, pcms::SimplexVsBbox>::apply(simplex(1.4, 2.0), + box)); + // Overlapping on axes 0-2 but entirely past the box on axis 3. + REQUIRE_FALSE(pcms::Overlap<4, pcms::SimplexVsBbox>::apply( + simplex(1.6, 2.0), box)); + } +} diff --git a/test/test_uniform_grid.cpp b/test/test_uniform_grid.cpp index d580c9f68..01efc6ea0 100644 --- a/test/test_uniform_grid.cpp +++ b/test/test_uniform_grid.cpp @@ -53,6 +53,40 @@ TEST_CASE("uniform grid") REQUIRE(0 == uniform_grid.GetCellIndex({0, 0})); REQUIRE(119 == uniform_grid.GetCellIndex({11, 9})); } + SECTION("AxisIndicesFromCellIndex") + { + // The inverse of CellIndexFromAxisIndices, in coordinate order. Cell 1 is + // the second cell along x, which the cell-BBOX test below also pins down + // (its center is (1.5, 0.5)): the grid stores its indices reversed relative + // to coordinates, so getting this wrong is silent. + { + const auto [i, j] = uniform_grid.AxisIndicesFromCellIndex(0); + REQUIRE(i == 0); + REQUIRE(j == 0); + } + { + const auto [i, j] = uniform_grid.AxisIndicesFromCellIndex(1); + REQUIRE(i == 1); + REQUIRE(j == 0); + } + { + const auto [i, j] = uniform_grid.AxisIndicesFromCellIndex(119); + REQUIRE(i == 9); + REQUIRE(j == 11); + } + + // Round-trip over every cell of the grid, which is the property a caller + // keyed by linear cell id depends on. + for (pcms::LO cell = 0; cell < uniform_grid.GetNumCells(); ++cell) { + const auto axis_indices = uniform_grid.AxisIndicesFromCellIndex(cell); + REQUIRE(uniform_grid.CellIndexFromAxisIndices(axis_indices) == cell); + // ...and it must agree with the binning rule used to find a point's cell. + for (int d = 0; d < 2; ++d) { + const auto center = uniform_grid.GetCellBBOX(cell).center[d]; + REQUIRE(uniform_grid.AxisCellIndex(d, center) == axis_indices[d]); + } + } + } } TEST_CASE("Create uniform grid from omega_h mesh") @@ -170,3 +204,65 @@ TEST_CASE("Create uniform grid from 3D mesh") REQUIRE(grid.divisions[2] == 12); } } + +// Grid resolution has to be a property of the mesh as a whole, not of one +// partition: search cells are named by a linear id, and an id identifies the +// same region of space on two ranks only if both resolved the grid the same +// way. A resolution derived from `mesh.nelems()` breaks that, because +// `nelems()` counts ghost copies and therefore changes with the partition. +TEST_CASE("global grid resolution") +{ + auto lib = Omega_h::Library{}; + auto world = lib.world(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, 0.0, 8, 8, 0, false); + + SECTION("global element count reduces to the local count when nothing is " + "ghosted") + { + REQUIRE(mesh.nghost_layers() == 0); + REQUIRE(pcms::GlobalElementCount(mesh) == mesh.nelems()); + } + + SECTION("the global resolution reproduces the local heuristic serially") + { + // Callers used to derive the resolution from mesh.nelems(). On an + // undecomposed mesh that is already the global count, so switching to the + // global helper must not change the grid they get. + REQUIRE(pcms::DivisionsPerAxisForGlobalMesh<2>(mesh) == + pcms::DivisionsPerAxisForMesh<2>(mesh.nelems())); + REQUIRE(pcms::DivisionsPerAxisForGlobalMesh<3>(mesh) == + pcms::DivisionsPerAxisForMesh<3>(mesh.nelems())); + } + + SECTION("resolution still targets roughly one element per cell") + { + const auto global_nelems = pcms::GlobalElementCount(mesh); + const auto n = pcms::DivisionsPerAxisForGlobalMesh<2>(mesh); + REQUIRE(n >= 10); // the min_divisions floor + REQUIRE(n * n >= global_nelems); + } +} + +// Linear cell ids and per-axis indices must be inverses on any grid, including +// one whose axes have different division counts. A grid with equal divisions +// on every axis hides an axis mix-up, so this one has three different counts. +TEST_CASE("cell index round trip on a grid with unequal divisions") +{ + const UniformGrid<3> grid{.edge_length = {3.0, 4.0, 5.0}, + .bot_left = {0.0, 0.0, 0.0}, + .divisions = {3, 4, 5}}; + for (pcms::LO cell = 0; cell < grid.GetNumCells(); ++cell) { + CAPTURE(cell); + const auto axis_indices = grid.AxisIndicesFromCellIndex(cell); + for (int d = 0; d < 3; ++d) { + REQUIRE(axis_indices[d] >= 0); + REQUIRE(axis_indices[d] < grid.divisions[d]); + } + REQUIRE(grid.CellIndexFromAxisIndices(axis_indices) == cell); + // The cell's own box must bin back to the same cell. + const auto box = grid.GetCellBBOX(cell); + REQUIRE(grid.ClosestCellID(Omega_h::Vector<3>{ + box.center[0], box.center[1], box.center[2]}) == cell); + } +} From e0f88b5dbfcab8b9f324990b351bd56977a6c904 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Tue, 27 Oct 2026 23:25:18 -0400 Subject: [PATCH 5/7] fix clang-tidy issues --- src/pcms/coupler/partition.h | 12 ---- src/pcms/localization/candidate_rank_table.h | 4 +- test/CMakeLists.txt | 10 ++- test/distributed_test_utils.h | 21 +++++++ test/test_distributed_coupling.cpp | 64 +++++++++++--------- test/test_distributed_evaluate.cpp | 56 +++++++++-------- test/test_distributed_field.cpp | 40 ++++++------ test/test_distributed_search.cpp | 37 ++++++----- 8 files changed, 138 insertions(+), 106 deletions(-) diff --git a/src/pcms/coupler/partition.h b/src/pcms/coupler/partition.h index 6a6461bf9..baba1751b 100644 --- a/src/pcms/coupler/partition.h +++ b/src/pcms/coupler/partition.h @@ -26,18 +26,6 @@ struct GetRank PCMS_FUNCTION_TIMER; return ptn.GetRank(coord_); } - // PointSearchPtn maps a point to a *list* of candidate ranks (a point may - // fall in the overlap region of several ranks' grid cells), unlike - // ClassPtn/RCBPtn which map to a single owning rank. GetDr() is a - // single-rank query, so this returns the first candidate as a - // representative owner; callers that need every candidate should use - // ptn.GetCandidateRanks() directly instead of GetDr(). - auto operator()(const redev::PointSearchPtn& ptn) - { - PCMS_FUNCTION_TIMER; - const auto& candidates = ptn.GetCandidateRanks(coord_); - return candidates.empty() ? static_cast(-1) : candidates.front(); - } LO id_; LO dim_; std::array coord_; diff --git a/src/pcms/localization/candidate_rank_table.h b/src/pcms/localization/candidate_rank_table.h index 6d7adf894..8a9eaaf13 100644 --- a/src/pcms/localization/candidate_rank_table.h +++ b/src/pcms/localization/candidate_rank_table.h @@ -127,10 +127,10 @@ class CandidateRankTable // of a fixed 2*Dim values per rank. `Omega_h::Comm::allgather` cannot be // used: it only exchanges with the ranks of the mesh's communication graph. // \see pcms/utility/mpi_utilities.h - Kokkos::Array local{}; + Kokkos::Array(Dim)> local{}; for (int d = 0; d < Dim; ++d) { local[static_cast(d)] = local_range_.lo[d]; - local[static_cast(Dim + d)] = local_range_.hi[d]; + local[static_cast(Dim) + d] = local_range_.hi[d]; } const auto gathered = AllGatherAllRanks(comm->get_impl(), local.data(), 2 * Dim); diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index a52dc4f80..f08c62598 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -76,10 +76,6 @@ if(PCMS_ENABLE_OMEGA_H) mpi_test(test_distributed_evaluate_2p 2 $) endif() - if(HOST_NPROC GREATER_EQUAL 3) - mpi_test(test_distributed_evaluate_3p 3 - $) - endif() if(HOST_NPROC GREATER_EQUAL 4) mpi_test(test_distributed_evaluate_4p 4 $) @@ -89,8 +85,10 @@ if(PCMS_ENABLE_OMEGA_H) mpi_test(test_distributed_search_2p 2 $) endif() - if(HOST_NPROC GREATER_EQUAL 3) - mpi_test(test_distributed_search_3p 3 + # Omega_h::build_box partitions by recursive bisection, which needs a + # power-of-two rank count, so the distributed tests skip 3 ranks. + if(HOST_NPROC GREATER_EQUAL 4) + mpi_test(test_distributed_search_4p 4 $) endif() set(d3d1p ${PCMS_TEST_DATA_DIR}/d3d/d3d-full_9k_sfc.osh/) diff --git a/test/distributed_test_utils.h b/test/distributed_test_utils.h index ad8387a17..e422e8edf 100644 --- a/test/distributed_test_utils.h +++ b/test/distributed_test_utils.h @@ -21,6 +21,8 @@ #include #include +#include +#include #include #include #include @@ -42,6 +44,17 @@ inline void require(bool condition, const char* message) } } +/// Report `what` and abort every rank. For exceptions escaping `main`: a rank +/// that just returned would leave the others blocked in their next collective. +[[noreturn]] inline void AbortAllRanks(const char* what) noexcept +{ + int rank = -1; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + std::fprintf(stderr, "[rank %d] uncaught exception: %s\n", rank, what); + MPI_Abort(MPI_COMM_WORLD, 1); + std::abort(); // MPI_Abort does not return, but is not declared noreturn +} + /// Abort unless every rank holds the same integer value. inline void require_same_on_all_ranks(Omega_h::CommPtr comm, LO value, const char* message) @@ -76,9 +89,17 @@ inline void require_same_real_on_all_ranks(Omega_h::CommPtr comm, Real value, /// The unit square (2D) or cube (3D) as a simplex mesh, partitioned over /// `world` with one ghost layer. The resolution grows with the rank count so /// every rank holds some elements. +/// +/// `world` must have a power-of-two size: `Omega_h::build_box` partitions by +/// recursive bisection. A debug Omega_h asserts on any other size; a release +/// one silently produces a broken partition. template Omega_h::Mesh BuildGhostedBox(Omega_h::CommPtr world) { + const int size = world->size(); + require(size > 0 && (size & (size - 1)) == 0, + "BuildGhostedBox needs a power-of-two rank count (Omega_h::build_box " + "partitions by recursive bisection)"); const LO n = (Dim == 2 ? 4 : 2) * world->size(); const Real z = Dim == 2 ? 0.0 : 1.0; const LO nz = Dim == 2 ? 0 : n; diff --git a/test/test_distributed_coupling.cpp b/test/test_distributed_coupling.cpp index 01ddd3242..2f2d9988a 100644 --- a/test/test_distributed_coupling.cpp +++ b/test/test_distributed_coupling.cpp @@ -423,35 +423,41 @@ void check_distributed_exchange(const DistributedCouplingContext& ctx) // (check_distributed_exchange). int main(int argc, char** argv) { - MPI_Init(&argc, &argv); - int result = 0; - { - 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"); - - const int n = 4 * nproc; - Omega_h::Mesh mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, - 0.0, /*nx=*/n, /*ny=*/n, /*nz=*/0, - /*symmetric=*/false); - mesh.set_parting(OMEGA_H_GHOSTED, 1, false); - - auto space = pcms::LagrangeFunctionSpace::FromMesh( - mesh, /*order=*/1, /*num_components=*/1, - pcms::csys::Cartesian::Deferred(), "global", - pcms::LagrangeFunctionSpace::Backend::OmegaH); - const DistributedCouplingContext ctx{mesh, space, rank, nproc}; - - check_serializer_round_trip(ctx); - - check_distributed_exchange(ctx); - - if (rank == 0) { - std::printf("distributed coupling test passed (nproc=%d)\n", nproc); + try { + MPI_Init(&argc, &argv); + int result = 0; + { + 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"); + + const int n = 4 * nproc; + Omega_h::Mesh mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, + 0.0, /*nx=*/n, /*ny=*/n, /*nz=*/0, + /*symmetric=*/false); + mesh.set_parting(OMEGA_H_GHOSTED, 1, false); + + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/1, /*num_components=*/1, + pcms::csys::Cartesian::Deferred(), "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + const DistributedCouplingContext ctx{mesh, space, rank, nproc}; + + check_serializer_round_trip(ctx); + + check_distributed_exchange(ctx); + + if (rank == 0) { + std::printf("distributed coupling test passed (nproc=%d)\n", nproc); + } } + MPI_Finalize(); + return result; + } catch (const std::exception& e) { + distributed_test::AbortAllRanks(e.what()); + } catch (...) { + distributed_test::AbortAllRanks("unknown exception"); } - MPI_Finalize(); - return result; } diff --git a/test/test_distributed_evaluate.cpp b/test/test_distributed_evaluate.cpp index 12ad3896a..07c95fc9a 100644 --- a/test/test_distributed_evaluate.cpp +++ b/test/test_distributed_evaluate.cpp @@ -741,7 +741,7 @@ void check_ghosts_must_be_synchronized(Omega_h::Mesh& mesh) std::printf("%dD ghost sync: %lld of %lld points wrong with stale ghosts, " "%lld after SynchronizeGhosts\n", Dim, static_cast(wrong_before_sync), - static_cast(n_points * mesh.comm()->size()), + static_cast(n_points) * mesh.comm()->size(), static_cast(wrong_after_sync)); } require(wrong_after_sync == 0, @@ -751,33 +751,39 @@ void check_ghosts_must_be_synchronized(Omega_h::Mesh& mesh) int main(int argc, char** argv) { - MPI_Init(&argc, &argv); - { - Omega_h::Library library(&argc, &argv); - auto world = library.world(); - require(world->size() >= 2, "test requires at least 2 MPI ranks"); - + try { + MPI_Init(&argc, &argv); { - auto mesh = distributed_test::BuildGhostedBox<2>(world); - check_agrees_with_serial<2>(mesh, /*order=*/1); - check_agrees_with_serial<2>(mesh, /*order=*/0); - check_boundary_arbitration<2>(mesh); - check_backend_without_distributed_code<2>(mesh); - check_out_of_bounds_policy<2>(mesh); - check_ghosts_must_be_synchronized<2>(mesh); - } + Omega_h::Library library(&argc, &argv); + auto world = library.world(); + require(world->size() >= 2, "test requires at least 2 MPI ranks"); + + { + auto mesh = distributed_test::BuildGhostedBox<2>(world); + check_agrees_with_serial<2>(mesh, /*order=*/1); + check_agrees_with_serial<2>(mesh, /*order=*/0); + check_boundary_arbitration<2>(mesh); + check_backend_without_distributed_code<2>(mesh); + check_out_of_bounds_policy<2>(mesh); + check_ghosts_must_be_synchronized<2>(mesh); + } - { - auto mesh = distributed_test::BuildGhostedBox<3>(world); - check_agrees_with_serial<3>(mesh, /*order=*/1); - check_ghosts_must_be_synchronized<3>(mesh); - check_out_of_bounds_policy<3>(mesh); - } + { + auto mesh = distributed_test::BuildGhostedBox<3>(world); + check_agrees_with_serial<3>(mesh, /*order=*/1); + check_ghosts_must_be_synchronized<3>(mesh); + check_out_of_bounds_policy<3>(mesh); + } - if (world->rank() == 0) { - std::printf("distributed evaluation test passed\n"); + if (world->rank() == 0) { + std::printf("distributed evaluation test passed\n"); + } } + MPI_Finalize(); + return 0; + } catch (const std::exception& e) { + distributed_test::AbortAllRanks(e.what()); + } catch (...) { + distributed_test::AbortAllRanks("unknown exception"); } - MPI_Finalize(); - return 0; } diff --git a/test/test_distributed_field.cpp b/test/test_distributed_field.cpp index ac1fd7d41..0cd35bb72 100644 --- a/test/test_distributed_field.cpp +++ b/test/test_distributed_field.cpp @@ -215,22 +215,28 @@ void check_distributed_evaluation(Omega_h::Mesh& mesh, int main(int argc, char** argv) { - MPI_Init(&argc, &argv); - { - Omega_h::Library library(&argc, &argv); - auto world = library.world(); - require(world->size() >= 2, "test requires at least 2 MPI ranks"); - - const int n = 4 * world->size(); - auto mesh = - Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, 0.0, n, n, 0, false); - mesh.set_parting(OMEGA_H_GHOSTED, 1, false); - auto space = pcms::LagrangeFunctionSpace::FromMesh( - mesh, 1, 1, pcms::csys::Cartesian::Deferred(), "global", - pcms::LagrangeFunctionSpace::Backend::OmegaH); - - check_distributed_evaluation(mesh, *space); + try { + MPI_Init(&argc, &argv); + { + Omega_h::Library library(&argc, &argv); + auto world = library.world(); + require(world->size() >= 2, "test requires at least 2 MPI ranks"); + + const int n = 4 * world->size(); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, 0.0, n, n, 0, false); + mesh.set_parting(OMEGA_H_GHOSTED, 1, false); + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, 1, 1, pcms::csys::Cartesian::Deferred(), "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + + check_distributed_evaluation(mesh, *space); + } + MPI_Finalize(); + return 0; + } catch (const std::exception& e) { + distributed_test::AbortAllRanks(e.what()); + } catch (...) { + distributed_test::AbortAllRanks("unknown exception"); } - MPI_Finalize(); - return 0; } diff --git a/test/test_distributed_search.cpp b/test/test_distributed_search.cpp index b0658a8fe..5d15c01f5 100644 --- a/test/test_distributed_search.cpp +++ b/test/test_distributed_search.cpp @@ -321,10 +321,10 @@ void check_published_ranges(Omega_h::Mesh& mesh, // Independent expectations, gathered and compared against the table: this is // the check that a rank-indexed gather went to all ranks rather than only to // the mesh's communication-graph neighbours. - Kokkos::Array local{}; + Kokkos::Array(Dim)> local{}; for (int d = 0; d < Dim; ++d) { local[static_cast(d)] = expected.lo[d]; - local[static_cast(Dim + d)] = expected.hi[d]; + local[static_cast(Dim) + d] = expected.hi[d]; } const auto all = pcms::AllGatherAllRanks(comm->get_impl(), local.data(), 2 * Dim); @@ -624,8 +624,9 @@ void check_local_classification(Omega_h::Mesh& mesh) static_cast(total_owned), static_cast(total_here), static_cast(total_elsewhere), static_cast(total_missing), - static_cast(total_owned_wrong + total_ghost_wrong), - static_cast(total_here + total_elsewhere + total_missing)); + static_cast(total_owned_wrong) + total_ghost_wrong, + static_cast(total_here) + total_elsewhere + + total_missing); } require( @@ -897,19 +898,25 @@ void run_all(Omega_h::CommPtr world) int main(int argc, char** argv) { - MPI_Init(&argc, &argv); - { - Omega_h::Library library(&argc, &argv); - auto world = library.world(); - require(world->size() >= 2, "test requires at least 2 MPI ranks"); + try { + MPI_Init(&argc, &argv); + { + Omega_h::Library library(&argc, &argv); + auto world = library.world(); + require(world->size() >= 2, "test requires at least 2 MPI ranks"); - run_all<2>(world); - run_all<3>(world); + run_all<2>(world); + run_all<3>(world); - if (world->rank() == 0) { - std::printf("distributed search test passed\n"); + if (world->rank() == 0) { + std::printf("distributed search test passed\n"); + } } + MPI_Finalize(); + return 0; + } catch (const std::exception& e) { + distributed_test::AbortAllRanks(e.what()); + } catch (...) { + distributed_test::AbortAllRanks("unknown exception"); } - MPI_Finalize(); - return 0; } From dd0b518a9503d5061c46d0cfe509510e005eccd5 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Wed, 28 Oct 2026 00:00:29 -0400 Subject: [PATCH 6/7] reformat --- src/pcms/field/data/mesh_fields.h | 7 +-- src/pcms/field/data/simple.h | 6 +-- .../evaluator/distributed_point_evaluator.h | 43 +++++++++--------- src/pcms/localization/overlap.h | 15 ++++--- src/pcms/localization/point_search.cpp | 45 ++++++++++--------- src/pcms/localization/point_search.h | 6 +-- src/pcms/utility/arrays.h | 1 - src/pcms/utility/mpi_utilities.h | 4 +- test/distributed_test_utils.h | 10 ++--- test/test_distributed_coupling.cpp | 8 ++-- test/test_distributed_evaluate.cpp | 14 +++--- test/test_distributed_field.cpp | 7 +-- test/test_distributed_search.cpp | 21 ++++----- test/test_point_search.cpp | 11 +++-- test/test_uniform_grid.cpp | 4 +- 15 files changed, 99 insertions(+), 103 deletions(-) diff --git a/src/pcms/field/data/mesh_fields.h b/src/pcms/field/data/mesh_fields.h index 225421a02..b37a4af3b 100644 --- a/src/pcms/field/data/mesh_fields.h +++ b/src/pcms/field/data/mesh_fields.h @@ -151,8 +151,9 @@ class MeshFieldsFieldData : public FieldData // subview of a LayoutLeft view is not representable as // LayoutLeft/LayoutRight (Kokkos falls back to LayoutStride), which // MakeConstRank2View/MakeRank2View do not detect. - Kokkos::View flat( - "ghost_sync_flat", static_cast(num_rows) * static_cast(nc)); + Kokkos::View flat("ghost_sync_flat", + static_cast(num_rows) * + static_cast(nc)); Kokkos::parallel_for( "SynchronizeGhostsFlatten", Kokkos::RangePolicy(0, num_rows), @@ -165,7 +166,7 @@ class MeshFieldsFieldData : public FieldData "SynchronizeGhostsUnflatten", Kokkos::RangePolicy(0, num_rows), UnflattenFunctor(flat, MakeRank2View(device_data_), row_off, - static_cast(nc))); + static_cast(nc))); row_offset += static_cast(num_rows); } diff --git a/src/pcms/field/data/simple.h b/src/pcms/field/data/simple.h index c6157968c..9af9ba45b 100644 --- a/src/pcms/field/data/simple.h +++ b/src/pcms/field/data/simple.h @@ -93,9 +93,9 @@ class SimpleFieldData : public FieldData // device_data_ may be component-major on device; flatten to dof-major // order for Omega_h's sync_array, then scatter the synchronized values // back. - Kokkos::View flat( - "ghost_sync_flat", - static_cast(num_rows) * static_cast(nc)); + Kokkos::View flat("ghost_sync_flat", + static_cast(num_rows) * + static_cast(nc)); CopyDeviceRank2ViewToDeviceView(flat, MakeConstRank2View(device_data_)); SynchronizeOmegaHBlock(oh->GetMesh(), dim, nc, flat); diff --git a/src/pcms/field/evaluator/distributed_point_evaluator.h b/src/pcms/field/evaluator/distributed_point_evaluator.h index 822391c2c..99a67b67a 100644 --- a/src/pcms/field/evaluator/distributed_point_evaluator.h +++ b/src/pcms/field/evaluator/distributed_point_evaluator.h @@ -313,12 +313,11 @@ void DistributedPointEvaluator::Setup( // ---- 2. send each remote point to its owner, once ---- // Only the remote points' coordinates leave the device, because MPI needs // them on the host. - const auto send = - CopyToHostRowMajor(GatherRows(points, remote_points_)); + const auto send = CopyToHostRowMajor(GatherRows(points, remote_points_)); std::vector incoming_scalar_counts; - const auto incoming = - ExchangeRecords(mesh_.comm()->get_impl(), detail::Scaled(remote_counts_, Dim), - send, incoming_scalar_counts); + const auto incoming = ExchangeRecords(mesh_.comm()->get_impl(), + detail::Scaled(remote_counts_, Dim), + send, incoming_scalar_counts); incoming_counts_.resize(incoming_scalar_counts.size()); for (size_t r = 0; r < incoming_counts_.size(); ++r) { incoming_counts_[r] = incoming_scalar_counts[r] / Dim; @@ -335,17 +334,16 @@ void DistributedPointEvaluator::Setup( // filled, once, so `Evaluate` never has to decode that from values. const auto n_incoming = static_cast(incoming_coords_.extent(0)); const auto incoming_filled = - incoming_evaluator_ ? detail::FilledFlags(incoming_evaluator_->FilledPoints(), - n_incoming) - : std::vector(static_cast(n_incoming), 0); - const auto remote_filled = ExchangeRecordsWithKnownCounts( - mesh_.comm()->get_impl(), incoming_counts_, incoming_filled, - remote_counts_); + incoming_evaluator_ + ? detail::FilledFlags(incoming_evaluator_->FilledPoints(), n_incoming) + : std::vector(static_cast(n_incoming), 0); + const auto remote_filled = + ExchangeRecordsWithKnownCounts(mesh_.comm()->get_impl(), incoming_counts_, + incoming_filled, remote_counts_); const auto own_filled = - own_evaluator_ - ? detail::FilledFlags(own_evaluator_->FilledPoints(), - static_cast(own.size())) - : std::vector{}; + own_evaluator_ ? detail::FilledFlags(own_evaluator_->FilledPoints(), + static_cast(own.size())) + : std::vector{}; std::vector answered(static_cast(n_points_), 0); for (size_t k = 0; k < own.size(); ++k) { @@ -389,12 +387,11 @@ Kokkos::View DistributedPointEvaluator::EvaluateLocally( const PointEvaluator* evaluator, LO n, const Field& field) const { - Kokkos::View values("local values", n, - n_components_); + Kokkos::View values("local values", n, n_components_); if (evaluator != nullptr) { - evaluator->Evaluate(field, ValueView( - field.GetData().GetValueBasis(), - MakeRank2View(values))); + evaluator->Evaluate( + field, ValueView(field.GetData().GetValueBasis(), + MakeRank2View(values))); } return values; } @@ -414,9 +411,9 @@ void DistributedPointEvaluator::Evaluate( // Owner side: evaluate this rank's own points and the points sent to it. const auto own_values = EvaluateLocally(own_evaluator_.get(), n_own, field); - const auto incoming_values = EvaluateLocally( - incoming_evaluator_.get(), static_cast(incoming_coords_.extent(0)), - field); + const auto incoming_values = + EvaluateLocally(incoming_evaluator_.get(), + static_cast(incoming_coords_.extent(0)), field); // Route the incoming values back to the ranks that asked: the only values // that cross to the host and back, because MPI takes host buffers. diff --git a/src/pcms/localization/overlap.h b/src/pcms/localization/overlap.h index 081463325..929c44817 100644 --- a/src/pcms/localization/overlap.h +++ b/src/pcms/localization/overlap.h @@ -324,13 +324,14 @@ struct Overlap const r3d::Few, Dim + 1>& b) { // The face normal below is a 2D perpendicular or a 3D cross product, so - // this predicate is 2D/3D only even though everything else about it reads as - // dimension-generic. Without this the higher-dimension case is silent: at - // Dim >= 4 `normal[Dim - 1]` is never written and is read as garbage in the - // loop below, and at Dim == 1 writing normal[1] and normal[2] runs off the - // end of the vector. - static_assert(Dim == 2 || Dim == 3, - "simplex-vs-simplex overlap is implemented for 2D and 3D only"); + // this predicate is 2D/3D only even though everything else about it reads + // as dimension-generic. Without this the higher-dimension case is silent: + // at Dim >= 4 `normal[Dim - 1]` is never written and is read as garbage in + // the loop below, and at Dim == 1 writing normal[1] and normal[2] runs off + // the end of the vector. + static_assert( + Dim == 2 || Dim == 3, + "simplex-vs-simplex overlap is implemented for 2D and 3D only"); // Each iteration treats one simplex as the half-space owner. Both are // needed: a face of `a` may separate the pair while no face of `b` does. diff --git a/src/pcms/localization/point_search.cpp b/src/pcms/localization/point_search.cpp index 14577c6b1..cd5a140d4 100644 --- a/src/pcms/localization/point_search.cpp +++ b/src/pcms/localization/point_search.cpp @@ -141,8 +141,8 @@ namespace detail * * struct MyElements { * static MyElements FromMesh(Omega_h::Mesh&); - * Operands operands(LO elem) const; // what Pred::apply takes - * AABBox bound(const Operands&) const; // a box containing it + * Operands operands(LO elem) const; // what Pred::apply + * takes AABBox bound(const Operands&) const; // a box containing it * }; * * `bound` is allowed to be conservative: the predicate rejects what it @@ -159,8 +159,8 @@ struct MeshSimplexElements return {mesh.ask_elem_verts(), mesh.coords()}; } - [[nodiscard]] KOKKOS_INLINE_FUNCTION Omega_h::Matrix - operands(LO elem) const + [[nodiscard]] KOKKOS_INLINE_FUNCTION Omega_h::Matrix operands( + LO elem) const { const auto verts = Omega_h::gather_verts(elems2verts, elem); return Omega_h::gather_vectors(coords, verts); @@ -182,9 +182,9 @@ struct MeshSimplexElements * own bounding box spans. That makes construction O(nelems * cells_per_elem) * instead of O(num_cells * nelems). * - * Operand-agnostic: `Elements` decides what an element *is* and how to bound it, - * `Pred` decides when it overlaps a cell. Neither is fixed here, so the same - * functor bins any shape a caller can supply a policy for. + * Operand-agnostic: `Elements` decides what an element *is* and how to bound + * it, `Pred` decides when it overlaps a cell. Neither is fixed here, so the + * same functor bins any shape a caller can supply a policy for. */ template > struct ElemCellOverlapFunctor @@ -250,8 +250,9 @@ construct_intersection_map(Omega_h::Mesh& mesh, int num_grid_cells) { if (mesh.dim() != Dim) { - std::cerr << "construct_intersection_map: mesh dimension does not match the " - "map dimension\n"; + std::cerr + << "construct_intersection_map: mesh dimension does not match the " + "map dimension\n"; std::terminate(); } @@ -259,9 +260,9 @@ construct_intersection_map(Omega_h::Mesh& mesh, const auto nelems = mesh.nelems(); CrsT elem2cells{}; - Kokkos::count_and_fill_crs( - elem2cells, nelems, - ElemCellOverlapFunctor{Elements::FromMesh(mesh), grid}); + Kokkos::count_and_fill_crs(elem2cells, nelems, + ElemCellOverlapFunctor{ + Elements::FromMesh(mesh), grid}); Kokkos::View counts("candidate_counts", num_grid_cells); Kokkos::parallel_for( @@ -356,10 +357,9 @@ construct_intersection_map_reference(Omega_h::Mesh& mesh, int num_grid_cells) { Kokkos::Crs intersection_map{}; - Kokkos::count_and_fill_crs( - intersection_map, num_grid_cells, - ReferenceCellMajorFunctor{Elements::FromMesh(mesh), grid, - mesh.nelems()}); + Kokkos::count_and_fill_crs(intersection_map, num_grid_cells, + ReferenceCellMajorFunctor{ + Elements::FromMesh(mesh), grid, mesh.nelems()}); return intersection_map; } Kokkos::Crs @@ -376,8 +376,8 @@ construct_intersection_map_reference_3d(Omega_h::Mesh& mesh, Kokkos::View grid, int num_grid_cells) { - return construct_intersection_map_reference<3, SimplexVsBbox>( - mesh, grid, num_grid_cells); + return construct_intersection_map_reference<3, SimplexVsBbox>(mesh, grid, + num_grid_cells); } // num_grid_cells should be result of grid.GetNumCells(), take as argument to @@ -412,13 +412,16 @@ construct_intersection_map_3d(Omega_h::Mesh& mesh, // callers outside this file need. template Kokkos::Crs construct_intersection_map<2, TriangleVsBbox>(Omega_h::Mesh&, - Kokkos::View[1]>, int); + Kokkos::View[1]>, + int); template Kokkos::Crs construct_intersection_map<2, SimplexVsBbox>(Omega_h::Mesh&, - Kokkos::View[1]>, int); + Kokkos::View[1]>, + int); template Kokkos::Crs construct_intersection_map<3, SimplexVsBbox>(Omega_h::Mesh&, - Kokkos::View[1]>, int); + Kokkos::View[1]>, + int); template Kokkos::Crs construct_intersection_map_reference<2, TriangleVsBbox>( diff --git a/src/pcms/localization/point_search.h b/src/pcms/localization/point_search.h index 74edd20f3..220d14020 100644 --- a/src/pcms/localization/point_search.h +++ b/src/pcms/localization/point_search.h @@ -86,9 +86,9 @@ construct_intersection_map_reference_3d(Omega_h::Mesh& mesh, /// How the candidate map learns an element's geometry, and a conservative bound /// on it. The default policy for an Omega_h simplex mesh; defined in the source -/// file alongside the rest of the map machinery. A different element shape means -/// a different policy with the same three members, which is the only thing that -/// has to change to bin that shape against the same grid. +/// file alongside the rest of the map machinery. A different element shape +/// means a different policy with the same three members, which is the only +/// thing that has to change to bin that shape against the same grid. template struct MeshSimplexElements; diff --git a/src/pcms/utility/arrays.h b/src/pcms/utility/arrays.h index 4928249c8..ce34baf73 100644 --- a/src/pcms/utility/arrays.h +++ b/src/pcms/utility/arrays.h @@ -527,7 +527,6 @@ void CopyRank1ViewToDevice(Rank1View dest, KOKKOS_LAMBDA(LO i) { dest(i) = src_tmp_d(i); }); } - // ---- std::vector <-> device copies ---- // // Each performs at most one host/device transfer. A row-major `std::vector` diff --git a/src/pcms/utility/mpi_utilities.h b/src/pcms/utility/mpi_utilities.h index 9004b5849..f77f75d76 100644 --- a/src/pcms/utility/mpi_utilities.h +++ b/src/pcms/utility/mpi_utilities.h @@ -106,8 +106,8 @@ template const auto per_rank = AllGatherAllRanks(comm, static_cast(values.size())); const auto displs = DisplacementsOf(per_rank); - std::vector gathered(static_cast( - std::accumulate(per_rank.begin(), per_rank.end(), 0))); + std::vector gathered( + static_cast(std::accumulate(per_rank.begin(), per_rank.end(), 0))); MPI_Allgatherv(values.data(), static_cast(values.size()), GetMPIType(T{}), gathered.data(), per_rank.data(), displs.data(), GetMPIType(T{}), comm); diff --git a/test/distributed_test_utils.h b/test/distributed_test_utils.h index e422e8edf..2d38c7a94 100644 --- a/test/distributed_test_utils.h +++ b/test/distributed_test_utils.h @@ -103,8 +103,8 @@ Omega_h::Mesh BuildGhostedBox(Omega_h::CommPtr world) const LO n = (Dim == 2 ? 4 : 2) * world->size(); const Real z = Dim == 2 ? 0.0 : 1.0; const LO nz = Dim == 2 ? 0 : n; - auto mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, z, n, n, nz, - false); + auto mesh = + Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, z, n, n, nz, false); mesh.set_parting(OMEGA_H_GHOSTED, 1, false); return mesh; } @@ -219,8 +219,7 @@ inline GatheredPoints GatherOwnedElementVertices(Omega_h::Mesh& mesh) /// Interleaved coordinates as a device point view, `[n][Dim]`. template Kokkos::View ToPointView(const std::vector& interleaved, - const char* label = "points") -{ + const char* label = "points") { const auto n = static_cast(interleaved.size() / Dim); Kokkos::View points(std::string(label), n); auto points_h = Kokkos::create_mirror_view(points); @@ -246,7 +245,8 @@ pcms::ValueView TagLike( template pcms::ValueView TagLike( - const pcms::Field& field, Kokkos::View values) + const pcms::Field& field, + Kokkos::View values) { return {field.GetData().GetValueBasis(), pcms::MakeRank2View(values)}; } diff --git a/test/test_distributed_coupling.cpp b/test/test_distributed_coupling.cpp index 2f2d9988a..4cf318ab2 100644 --- a/test/test_distributed_coupling.cpp +++ b/test/test_distributed_coupling.cpp @@ -144,8 +144,8 @@ void check_serializer_round_trip(const DistributedCouplingContext& ctx) auto target = ctx.space->CreateFunction(); serializer.Deserialize(target, pcms::make_const_array_view(buffer), pcms::make_const_array_view(permutation)); - auto target_data = pcms::FlattenToRank1View( - target.GetDOFHolderDataHost().GetValues()); + auto target_data = + pcms::FlattenToRank1View(target.GetDOFHolderDataHost().GetValues()); require(target_data.size() == static_cast(n_local), "target local size"); @@ -377,8 +377,8 @@ void check_distributed_exchange(const DistributedCouplingContext& ctx) exchange_serializer.Deserialize( echo_field, pcms::make_const_array_view(payload), pcms::make_const_array_view(recv_plan.permutation)); - auto echo_data = pcms::FlattenToRank1View( - echo_field.GetDOFHolderDataHost().GetValues()); + auto echo_data = + pcms::FlattenToRank1View(echo_field.GetDOFHolderDataHost().GetValues()); for (LO o = 0; o < n_owned; ++o) { const LO local = owned_to_local(o); if (recv_plan.permutation[static_cast(o)] >= 0) { diff --git a/test/test_distributed_evaluate.cpp b/test/test_distributed_evaluate.cpp index 07c95fc9a..22aa053ba 100644 --- a/test/test_distributed_evaluate.cpp +++ b/test/test_distributed_evaluate.cpp @@ -312,8 +312,7 @@ void check_boundary_arbitration(Omega_h::Mesh& mesh) auto field = space->CreateFunction(); SetLinearField(field.GetData(), *space->GetLayout()); - const auto gathered = - distributed_test::GatherOwnedElementVertices(mesh); + const auto gathered = distributed_test::GatherOwnedElementVertices(mesh); require(gathered.total > 0, "no owned element vertices were gathered"); auto device_coords = MakeQueryCoords(gathered.coords, dim); const LO n_points = static_cast(gathered.total); @@ -662,7 +661,6 @@ void check_out_of_bounds_policy(Omega_h::Mesh& mesh) require(threw, "the ERROR policy did not report a point outside the mesh"); } - /// The distributed evaluator reads ghost DOF values. /// /// A rank evaluates only points in elements it owns, but an owned element can @@ -688,8 +686,8 @@ void check_ghosts_must_be_synchronized(Omega_h::Mesh& mesh) // Owned DOFs hold the affine field; ghost DOFs hold a stale value. const auto dof_coords = layout.GetDOFHolderCoordinates().GetValues(); - Kokkos::View coords_d("dof coords", - n_local, Dim); + Kokkos::View coords_d("dof coords", n_local, + Dim); Kokkos::parallel_for( "copy_dof_coords", Kokkos::RangePolicy(0, n_local), @@ -703,9 +701,9 @@ void check_ghosts_must_be_synchronized(Omega_h::Mesh& mesh) const auto owned = layout.GetOwnedHost(); std::vector values(static_cast(n_local)); for (LO i = 0; i < n_local; ++i) { - values[static_cast(i)] = - owned[static_cast(i)] ? linear_f(coords(i, 0), coords(i, 1)) - : kStaleGhost; + values[static_cast(i)] = owned[static_cast(i)] + ? linear_f(coords(i, 0), coords(i, 1)) + : kStaleGhost; } auto field = space->CreateFunction(); field.SetDOFHolderDataHost(distributed_test::TagLike( diff --git a/test/test_distributed_field.cpp b/test/test_distributed_field.cpp index 0cd35bb72..11a0dbc6c 100644 --- a/test/test_distributed_field.cpp +++ b/test/test_distributed_field.cpp @@ -111,7 +111,8 @@ void check_distributed_evaluation(Omega_h::Mesh& mesh, n_local, 1))); field.SynchronizeGhosts(); - auto data = pcms::FlattenToRank1View(field.GetDOFHolderDataHost().GetValues()); + auto data = + pcms::FlattenToRank1View(field.GetDOFHolderDataHost().GetValues()); for (LO i = 0; i < n_local; ++i) { require(std::fabs(data[static_cast(i)] - linear_f(coords(i, 0), coords(i, 1))) < 1.0e-12, @@ -223,8 +224,8 @@ int main(int argc, char** argv) require(world->size() >= 2, "test requires at least 2 MPI ranks"); const int n = 4 * world->size(); - auto mesh = - Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, 0.0, n, n, 0, false); + auto mesh = Omega_h::build_box(world, OMEGA_H_SIMPLEX, 1.0, 1.0, 0.0, n, + n, 0, false); mesh.set_parting(OMEGA_H_GHOSTED, 1, false); auto space = pcms::LagrangeFunctionSpace::FromMesh( mesh, 1, 1, pcms::csys::Cartesian::Deferred(), "global", diff --git a/test/test_distributed_search.cpp b/test/test_distributed_search.cpp index 5d15c01f5..0481b09e9 100644 --- a/test/test_distributed_search.cpp +++ b/test/test_distributed_search.cpp @@ -553,8 +553,8 @@ void check_local_classification(Omega_h::Mesh& mesh) require(mesh.nghost_layers() > 0, "mesh is not ghosted"); const auto divisions = pcms::DivisionsPerAxisForGlobalMesh(mesh); - const auto local_search = distributed_test::MakeLocalSearch(mesh, - divisions); + const auto local_search = + distributed_test::MakeLocalSearch(mesh, divisions); const pcms::DistributedPointSearch search{mesh, local_search}; require(search.HasGhostLayer(), "distributed search missed the ghost layer"); @@ -625,8 +625,7 @@ void check_local_classification(Omega_h::Mesh& mesh) static_cast(total_elsewhere), static_cast(total_missing), static_cast(total_owned_wrong) + total_ghost_wrong, - static_cast(total_here) + total_elsewhere + - total_missing); + static_cast(total_here) + total_elsewhere + total_missing); } require( @@ -675,8 +674,8 @@ void check_distributed_locate(Omega_h::Mesh& mesh) gathered.coords, "global element centroids"); const auto divisions = pcms::DivisionsPerAxisForGlobalMesh(mesh); - const auto local_search = distributed_test::MakeLocalSearch(mesh, - divisions); + const auto local_search = + distributed_test::MakeLocalSearch(mesh, divisions); pcms::DistributedPointSearch search{mesh, local_search}; const auto results = search.Locate(points); @@ -716,10 +715,9 @@ void check_distributed_locate(Omega_h::Mesh& mesh) // A point no rank placed must not name an owner or an element, so a caller // cannot mistake it for a location. const auto& outside = results[static_cast(total)]; - const bool outside_not_found = outside.status == ResultStatus::NOT_FOUND && - outside.owner_rank == -1 && - outside.element_id == -1 && - outside.global_element_id == -1; + const bool outside_not_found = + outside.status == ResultStatus::NOT_FOUND && outside.owner_rank == -1 && + outside.element_id == -1 && outside.global_element_id == -1; const LO t_not_found = mesh.comm()->allreduce(n_not_found, Omega_h_Op::OMEGA_H_SUM); @@ -731,8 +729,7 @@ void check_distributed_locate(Omega_h::Mesh& mesh) mesh.comm()->allreduce(n_bad_bary, Omega_h_Op::OMEGA_H_SUM); const LO t_points = mesh.comm()->allreduce(n_points, Omega_h_Op::OMEGA_H_SUM); const LO t_routed_away = mesh.comm()->allreduce( - static_cast(total) - n_local, - Omega_h_Op::OMEGA_H_SUM); + static_cast(total) - n_local, Omega_h_Op::OMEGA_H_SUM); if (rank == 0) { std::printf( diff --git a/test/test_point_search.cpp b/test/test_point_search.cpp index d713abed8..2f90154f9 100644 --- a/test/test_point_search.cpp +++ b/test/test_point_search.cpp @@ -839,8 +839,7 @@ TEST_CASE("grid and overlap layer above 3D", "[point_search][nd]") // Inside cell (0, 1, 2, 3); padded by one cell, clamped to the grid. const AABBox<4> box{.center = {0.5, 1.5, 2.5, 3.5}, .half_width = {0.25, 0.25, 0.25, 0.25}}; - const auto range = - pcms::candidate_cells<4, pcms::SimplexVsBbox>(box, grid); + const auto range = pcms::candidate_cells<4, pcms::SimplexVsBbox>(box, grid); const Kokkos::Array lo{0, 0, 1, 2}; const Kokkos::Array hi{1, 2, 3, 4}; for (int d = 0; d < 4; ++d) { @@ -868,10 +867,10 @@ TEST_CASE("grid and overlap layer above 3D", "[point_search][nd]") } return coords; }; - REQUIRE(pcms::Overlap<4, pcms::SimplexVsBbox>::apply(simplex(1.4, 2.0), - box)); + REQUIRE( + pcms::Overlap<4, pcms::SimplexVsBbox>::apply(simplex(1.4, 2.0), box)); // Overlapping on axes 0-2 but entirely past the box on axis 3. - REQUIRE_FALSE(pcms::Overlap<4, pcms::SimplexVsBbox>::apply( - simplex(1.6, 2.0), box)); + REQUIRE_FALSE( + pcms::Overlap<4, pcms::SimplexVsBbox>::apply(simplex(1.6, 2.0), box)); } } diff --git a/test/test_uniform_grid.cpp b/test/test_uniform_grid.cpp index 01efc6ea0..55a99743b 100644 --- a/test/test_uniform_grid.cpp +++ b/test/test_uniform_grid.cpp @@ -262,7 +262,7 @@ TEST_CASE("cell index round trip on a grid with unequal divisions") REQUIRE(grid.CellIndexFromAxisIndices(axis_indices) == cell); // The cell's own box must bin back to the same cell. const auto box = grid.GetCellBBOX(cell); - REQUIRE(grid.ClosestCellID(Omega_h::Vector<3>{ - box.center[0], box.center[1], box.center[2]}) == cell); + REQUIRE(grid.ClosestCellID(Omega_h::Vector<3>{box.center[0], box.center[1], + box.center[2]}) == cell); } } From 4c5594ceb34558e93c36950aaf279acde9ee0c92 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Thu, 29 Oct 2026 21:05:44 -0400 Subject: [PATCH 7/7] refactor distributed evaluator --- .../evaluator/distributed_point_evaluator.h | 624 ++++++++++-------- src/pcms/localization/CMakeLists.txt | 1 + .../localization/distributed_point_search.cpp | 14 +- .../localization/distributed_point_search.h | 15 +- .../localization/distributed_routing_plan.h | 163 +++++ test/test_distributed_evaluate.cpp | 164 ++++- test/test_distributed_search.cpp | 26 +- 7 files changed, 674 insertions(+), 333 deletions(-) create mode 100644 src/pcms/localization/distributed_routing_plan.h diff --git a/src/pcms/field/evaluator/distributed_point_evaluator.h b/src/pcms/field/evaluator/distributed_point_evaluator.h index 99a67b67a..5123106b9 100644 --- a/src/pcms/field/evaluator/distributed_point_evaluator.h +++ b/src/pcms/field/evaluator/distributed_point_evaluator.h @@ -1,27 +1,32 @@ #ifndef PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H #define PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H +#include "pcms/field/coordinate_system.hpp" +#include "pcms/field/coordinate_view.hpp" #include "pcms/field/evaluation_request.h" +#include "pcms/field/evaluator/omega_h_lagrange.h" #include "pcms/field/field.h" #include "pcms/field/function_space.h" #include "pcms/field/out_of_bounds_policy.h" #include "pcms/field/point_evaluator.h" +#include "pcms/field/point_evaluator_factory.hpp" #include "pcms/field/value_view.hpp" #include "pcms/localization/distributed_point_search.h" +#include "pcms/localization/distributed_routing_plan.h" #include "pcms/localization/point_search.h" #include "pcms/utility/arrays.h" #include "pcms/utility/assert.h" #include "pcms/utility/memory_spaces.h" #include "pcms/utility/mpi_utilities.h" -#include "pcms/utility/uniform_grid.h" #include #include -#include #include +#include #include -#include +#include +#include #include namespace pcms @@ -30,37 +35,10 @@ namespace pcms namespace detail { -/// The concrete local search for a spatial dimension. -template -struct GridSearchOf; -template <> -struct GridSearchOf<2> -{ - using type = GridPointSearch2D; -}; -template <> -struct GridSearchOf<3> -{ - using type = GridPointSearch3D; -}; - -/// A search over the whole decomposition, so every rank builds the same grid. -/// -/// The local evaluator factories build their search from the local element -/// count, which differs per rank. Routing needs the opposite: a grid every rank -/// agrees on, because the rank a point is routed to is read off that grid. -template -typename GridSearchOf::type MakeGlobalGridSearch(Omega_h::Mesh& mesh) -{ - const auto n = DivisionsPerAxisForGlobalMesh(mesh); - if constexpr (Dim == 3) { - return typename GridSearchOf::type(mesh, n, n, n); - } else { - return typename GridSearchOf::type(mesh, n, n); - } -} - using DeviceExec = DeviceMemorySpace::execution_space; +using DeviceCoords = + Rank2View>; /// Set every entry of `values` to `fill`. template @@ -92,27 +70,84 @@ void ScatterRows(Values values, Kokkos::View rows, }); } -/// Whether each of the first `n` rows appears in `filled`, as host flags. -inline std::vector FilledFlags( - Kokkos::View filled, LO n) +/// Rows `rows(k)` of a coordinate view, as a new device matrix. +inline Kokkos::View GatherCoordinateRows( + DeviceCoords src, Kokkos::View rows) +{ + const int dim = static_cast(src.extent(1)); + Kokkos::View out("gathered coords", rows.extent(0), + dim); + Kokkos::parallel_for( + "DistributedPointEvaluator::GatherCoords", + Kokkos::RangePolicy(0, rows.extent(0)), KOKKOS_LAMBDA(LO k) { + for (int d = 0; d < dim; ++d) { + out(k, d) = src(rows(k), d); + } + }); + return out; +} + +/// A coordinate view copied into the fixed-dimension layout the search takes. +template +Kokkos::View ToSearchCoords(DeviceCoords src) { + const auto n = static_cast(src.extent(0)); + Kokkos::View out("search coords", n); + Kokkos::parallel_for( + "DistributedPointEvaluator::ToSearchCoords", + Kokkos::RangePolicy(0, n), KOKKOS_LAMBDA(LO i) { + for (int d = 0; d < Dim; ++d) { + out(i, d) = src(i, d); + } + }); + return out; +} + +/// Whether each of the first `n` rows appears in `rows`, as host flags. +inline std::vector FlagsOf( + Kokkos::View rows, LO n) { std::vector flags(static_cast(n), 0); - const auto filled_h = - Kokkos::create_mirror_view_and_copy(HostMemorySpace{}, filled); - for (size_t k = 0; k < filled_h.extent(0); ++k) { - flags[static_cast(filled_h(k))] = 1; + const auto rows_h = + Kokkos::create_mirror_view_and_copy(HostMemorySpace{}, rows); + for (size_t k = 0; k < rows_h.extent(0); ++k) { + flags[static_cast(rows_h(k))] = 1; } return flags; } -/// `counts` times `factor`, for sending `factor` scalars per point. -inline std::vector Scaled(const std::vector& counts, int factor) +/// `view` as a host vector. +inline std::vector ToHostVector( + Kokkos::View view) +{ + const auto h = Kokkos::create_mirror_view_and_copy(HostMemorySpace{}, view); + return std::vector(h.data(), h.data() + h.extent(0)); +} + +/// Throws on every rank if `error` is non-empty on any rank, so a failure on +/// one rank never leaves the others blocked in the next collective. +inline void ThrowIfAnyRankFailed(const Omega_h::CommPtr& comm, + const std::string& context, + const std::string& error) { - std::vector out(counts.size()); - for (size_t r = 0; r < counts.size(); ++r) { - out[r] = counts[r] * factor; + const int failed = + comm->allreduce(error.empty() ? 0 : 1, Omega_h_Op::OMEGA_H_MAX); + if (failed != 0) { + throw pcms_error( + context + ": " + + (error.empty() ? std::string("failed on another rank") : error)); + } +} + +inline void CheckSameCoordinateSystem( + const char* context, const CoordinateView& coords, + const std::shared_ptr& expected) +{ + if (!SameCoordinateSystem(coords.GetCoordinateSystem(), expected)) { + throw pcms_error(std::string(context) + ": query coordinate system '" + + std::string(coords.GetCoordinateSystem()->Kind()) + + "' is not the factory's coordinate system '" + + std::string(expected->Kind()) + "'"); } - return out; } } // namespace detail @@ -120,87 +155,61 @@ inline std::vector Scaled(const std::vector& counts, int factor) /** * \brief Evaluate a field at points that are distributed over ranks. * - * Any rank can evaluate the field at any point, including points outside its - * own partition, and every rank gets the same answer. + * Built by `DistributedPointEvaluatorFactory`, from a routing plan and an + * `EvaluationRequest`. Any rank can evaluate the field at any point, including + * points outside its own partition. * - * Construction (collective, runs once because the query points are fixed): - * 1. `DistributedPointSearch::Locate` finds the owning rank of every point: - * local search first, then the global search grid gives candidate ranks, - * which search the point locally; the lowest owning rank wins. - * 2. Each point owned by another rank is sent to its owner, once. - * 3. Each rank builds two ordinary local evaluators through - * `FunctionSpace::CreatePointEvaluator`: one for its own points, and one - * for the points other ranks sent it. + * Construction (collective): each rank builds two ordinary local evaluators + * through `FunctionSpace::CreatePointEvaluator`, one for the points it owns and + * one for the points other ranks sent it. Their requests are derived from the + * caller's: + * - The query layout is passed on only when this rank owns every one of its + * query points, because only then do the local points match the layout. + * Otherwise the local request carries coordinates only. + * - The local evaluators always use FILL. The caller's policy is applied here, + * after the exchange, so that ERROR throws on every rank together. * * `Evaluate` (collective): each rank evaluates both batches locally, and one * `MPI_Alltoallv` sends the values for the incoming batch back to the ranks - * that asked. Locations are never searched again; values are recomputed on - * every call, from the coefficients as they are now. - * - * Nothing here is specific to a field type: a backend gets distributed - * evaluation as soon as it can evaluate locally. + * that asked. Points are never searched again. * * Precondition: the field's ghost DOF values are synchronized * (`Field::SynchronizeGhosts`). A rank evaluates only points in elements it * owns, but an owned element can have vertices another rank owns, and their - * values on this rank are ghost copies. `FieldCommunicator::Receive` - * synchronizes after receiving; a field written some other way, with only its - * owned values set, must be synchronized before `Evaluate`. + * values on this rank are ghost copies. * * A point no rank can locate, or one its owner's local evaluator still cannot - * place (that evaluator's search disagreeing with the routing search), is an - * unanswered point. Which points those are is settled at construction: under - * FILL they get the fill value and are reported by `FilledPoints()`; under - * ERROR, `Evaluate` throws after its exchange, so no rank is left blocked. - * - * \tparam T Value type. Each backend accepts `CreatePointEvaluator` for - * whichever types it supports. - * \tparam Dim Spatial dimension of the mesh, 2 or 3. + * place, is an unanswered point. Which points those are is settled at + * construction: under FILL they get the fill value and are reported by + * `FilledPoints()`; under ERROR, `Evaluate` throws after its exchange. */ -template > class DistributedPointEvaluator : public PointEvaluator { public: - /** - * \param mesh The mesh points are located on, and whose ownership decides who - * answers for a point. Must outlive this object. - * \param space The field's function space. Every local evaluation goes - * through it. - * \param points Query coordinates, `[num_points][Dim]`, fixed for the life of - * the evaluator. - * \param policy What to do with a point no rank can locate. - * \param tolerance_reach How far outside an element the search may still - * accept a point. - * \param grid_divisions Cells per axis of the global search grid; 0 derives - * it from the mesh. Must be the same on every rank. - */ - DistributedPointEvaluator( - Omega_h::Mesh& mesh, std::shared_ptr space, - Kokkos::View points, - OutOfBoundsPolicy policy = {}, Real tolerance_reach = 1e-12, - LO grid_divisions = 0) - : mesh_(mesh), space_(std::move(space)), policy_(policy) + /// Collective. Prefer `DistributedPointEvaluatorFactory`, which builds the + /// plan and checks that it matches the request. + /// + /// \param space The field's function space. Every local evaluation goes + /// through it. + /// \param plan Routing for exactly the request's query points. + /// \param request Query points, policy and optional query layout. + DistributedPointEvaluator(std::shared_ptr space, + std::shared_ptr plan, + const EvaluationRequest& request) + : space_(std::move(space)), plan_(std::move(plan)), policy_(request.policy) { - if (!space_) { + if (!space_ || !plan_) { throw std::invalid_argument( - "DistributedPointEvaluator: function space is null"); - } - if (mesh_.dim() != Dim) { - throw std::invalid_argument( - "DistributedPointEvaluator: mesh dimension does not match"); + "DistributedPointEvaluator: function space and plan must not be null"); } n_components_ = space_->GetLayout()->GetNumComponents(); - n_points_ = static_cast(points.extent(0)); - Setup(points, tolerance_reach, grid_divisions); + Setup(request); } /// Evaluate the field at every query point. Collective. - /// - /// \param field Supplies the coefficients. Must be compatible with this - /// evaluator's function space. - /// \param values Output, `[num_query_points][num_components]`. void Evaluate( const Field& field, ValueView values) const override; @@ -219,49 +228,34 @@ class DistributedPointEvaluator : public PointEvaluator /// How many query points get their value from another rank. [[nodiscard]] std::size_t LastRoutedPointCount() const noexcept { - return remote_points_.extent(0); + return plan_->RoutedPointCount(); + } + + [[nodiscard]] const DistributedRoutingPlan& GetRoutingPlan() const noexcept + { + return *plan_; } - /// The space whose local evaluator produces the values. [[nodiscard]] const FunctionSpace& GetFunctionSpace() const noexcept { return *space_; } private: - void Setup(Kokkos::View points, - Real tolerance_reach, LO grid_divisions); - - /// A local evaluator over `coords`, or null when there are none. `coords` - /// must stay alive for as long as the evaluator. - [[nodiscard]] std::unique_ptr> MakeLocalEvaluator( - Kokkos::View coords) const; + void Setup(const EvaluationRequest& request); /// Run `evaluator` over its `n` points, on the device. [[nodiscard]] Kokkos::View EvaluateLocally( const PointEvaluator* evaluator, LO n, const Field& field) const; - Omega_h::Mesh& mesh_; std::shared_ptr space_; + std::shared_ptr plan_; OutOfBoundsPolicy policy_; int n_components_ = 0; - LO n_points_ = 0; - // Query points this rank owns (indices into the query set), their - // coordinates, and the evaluator over them. - Kokkos::View own_points_; + // Own points' coordinates, when they are a subset of the request's. Kokkos::View own_coords_; std::unique_ptr> own_evaluator_; - - // Query points another rank owns, grouped by owner rank in ascending order, - // and how many go to each rank. - Kokkos::View remote_points_; - std::vector remote_counts_; - - // Points other ranks sent this rank to evaluate, grouped by requesting rank, - // and how many came from each rank. - std::vector incoming_counts_; - Kokkos::View incoming_coords_; std::unique_ptr> incoming_evaluator_; // Query points nothing answers for, ascending: never located, or located @@ -269,83 +263,72 @@ class DistributedPointEvaluator : public PointEvaluator Kokkos::View unanswered_points_; }; -template -void DistributedPointEvaluator::Setup( - Kokkos::View points, - Real tolerance_reach, LO grid_divisions) +template +void DistributedPointEvaluator::Setup( + const EvaluationRequest& request) { - // ---- 1. find the owning rank of every point ---- - // The search is only needed here: once the owners are known, `Evaluate` - // never searches again. - const auto local_search = detail::MakeGlobalGridSearch(mesh_); - DistributedPointSearch search(mesh_, local_search, tolerance_reach, - grid_divisions); - const auto located = search.Locate(points); - - // Split the query indices by owner. The search result is host data, so this - // is host work on indices only; coordinates stay on the device. - const int rank = search.Rank(); - const int nranks = mesh_.comm()->size(); - std::vector> by_owner(static_cast(nranks)); - std::vector own; - for (LO i = 0; i < n_points_; ++i) { - const auto& r = located[static_cast(i)]; - if (r.status != DistributedPointSearch::ResultStatus::LOCATED) { - continue; // in neither list; `Evaluate` fills it + const auto& plan = *plan_; + const auto n_own = static_cast(plan.own_points.extent(0)); + const auto n_incoming = static_cast(plan.IncomingPointCount()); + + // Points this evaluator cannot place were already routed to their owner, so + // a local evaluator that still cannot place one must not throw or guess: + // FILL marks it, and its `FilledPoints()` reports it below. + OutOfBoundsPolicy fallback; + fallback.mode = OutOfBoundsMode::FILL; + fallback.fill_value = policy_.fill_value; + const auto& system = space_->GetCoordinateSystem(); + + // Local evaluator creation is local, but a failure must reach every rank + // before the exchange below. + std::string error; + try { + if (n_own > 0) { + if (plan.AllOwnedHere()) { + // The local points are the request's points, so the query layout still + // describes them and the backend may use it. + EvaluationRequest own_request = request; + own_request.policy = fallback; + own_evaluator_ = space_->CreatePointEvaluator(own_request); + } else { + own_coords_ = detail::GatherCoordinateRows(request.coords.GetValues(), + plan.own_points); + own_evaluator_ = + space_->CreatePointEvaluator(EvaluationRequest::FromCoordinates( + CoordinateView(system, + MakeConstRank2View(own_coords_)), + fallback)); + } } - if (r.owner_rank == rank) { - own.push_back(i); - } else { - by_owner[static_cast(r.owner_rank)].push_back(i); + if (n_incoming > 0) { + incoming_evaluator_ = + space_->CreatePointEvaluator(EvaluationRequest::FromCoordinates( + CoordinateView( + system, MakeConstRank2View(plan.incoming_coords)), + fallback)); } + } catch (const std::exception& e) { + error = e.what(); } - std::vector remote; - remote_counts_.assign(static_cast(nranks), 0); - for (int r = 0; r < nranks; ++r) { - const auto& points_of_r = by_owner[static_cast(r)]; - remote.insert(remote.end(), points_of_r.begin(), points_of_r.end()); - remote_counts_[static_cast(r)] = - static_cast(points_of_r.size()); - } - own_points_ = CopyToDevice(own); - remote_points_ = CopyToDevice(remote); - - // ---- 2. send each remote point to its owner, once ---- - // Only the remote points' coordinates leave the device, because MPI needs - // them on the host. - const auto send = CopyToHostRowMajor(GatherRows(points, remote_points_)); - std::vector incoming_scalar_counts; - const auto incoming = ExchangeRecords(mesh_.comm()->get_impl(), - detail::Scaled(remote_counts_, Dim), - send, incoming_scalar_counts); - incoming_counts_.resize(incoming_scalar_counts.size()); - for (size_t r = 0; r < incoming_counts_.size(); ++r) { - incoming_counts_[r] = incoming_scalar_counts[r] / Dim; - } + detail::ThrowIfAnyRankFailed(plan.comm, "DistributedPointEvaluator", error); - // ---- 3. one local evaluator per batch, built once ---- - own_coords_ = GatherRows(points, own_points_); // device only - incoming_coords_ = CopyRowMajorToDevice(incoming, Dim); - own_evaluator_ = MakeLocalEvaluator(own_coords_); - incoming_evaluator_ = MakeLocalEvaluator(incoming_coords_); - - // ---- 4. which points nothing answers for ---- // Each owner tells the requester which of its points its local evaluator // filled, once, so `Evaluate` never has to decode that from values. - const auto n_incoming = static_cast(incoming_coords_.extent(0)); const auto incoming_filled = incoming_evaluator_ - ? detail::FilledFlags(incoming_evaluator_->FilledPoints(), n_incoming) + ? detail::FlagsOf(incoming_evaluator_->FilledPoints(), n_incoming) : std::vector(static_cast(n_incoming), 0); const auto remote_filled = - ExchangeRecordsWithKnownCounts(mesh_.comm()->get_impl(), incoming_counts_, - incoming_filled, remote_counts_); + ExchangeRecordsWithKnownCounts(plan.comm->get_impl(), plan.incoming_counts, + incoming_filled, plan.remote_counts); const auto own_filled = - own_evaluator_ ? detail::FilledFlags(own_evaluator_->FilledPoints(), - static_cast(own.size())) + own_evaluator_ ? detail::FlagsOf(own_evaluator_->FilledPoints(), n_own) : std::vector{}; - std::vector answered(static_cast(n_points_), 0); + // Unlocated points stay unanswered. + std::vector answered(static_cast(plan.n_points), 0); + const auto own = detail::ToHostVector(plan.own_points); + const auto remote = detail::ToHostVector(plan.remote_points); for (size_t k = 0; k < own.size(); ++k) { answered[static_cast(own[k])] = own_filled[k] == 0; } @@ -353,7 +336,7 @@ void DistributedPointEvaluator::Setup( answered[static_cast(remote[k])] = remote_filled[k] == 0; } std::vector unanswered; - for (LO i = 0; i < n_points_; ++i) { + for (LO i = 0; i < plan.n_points; ++i) { if (answered[static_cast(i)] == 0) { unanswered.push_back(i); } @@ -361,30 +344,9 @@ void DistributedPointEvaluator::Setup( unanswered_points_ = CopyToDevice(unanswered); } -template -std::unique_ptr> -DistributedPointEvaluator::MakeLocalEvaluator( - Kokkos::View coords) const -{ - if (coords.extent(0) == 0) { - return nullptr; - } - // Only owned points reach this evaluator, so one it cannot place is a - // disagreement between its search and the routing search. It must not throw - // or guess: FILL marks it, and its `FilledPoints()` reports it at setup. - OutOfBoundsPolicy fallback; - fallback.mode = OutOfBoundsMode::FILL; - fallback.fill_value = policy_.fill_value; - - const CoordinateView view(space_->GetCoordinateSystem(), - MakeConstRank2View(coords)); - return space_->CreatePointEvaluator( - EvaluationRequest::FromCoordinates(view, fallback)); -} - -template +template Kokkos::View -DistributedPointEvaluator::EvaluateLocally( +DistributedPointEvaluator::EvaluateLocally( const PointEvaluator* evaluator, LO n, const Field& field) const { Kokkos::View values("local values", n, n_components_); @@ -396,37 +358,39 @@ DistributedPointEvaluator::EvaluateLocally( return values; } -template -void DistributedPointEvaluator::Evaluate( +template +void DistributedPointEvaluator::Evaluate( const Field& field, ValueView out) const { this->CheckEvaluateWriteTag("DistributedPointEvaluator::Evaluate", field, out); + const auto& plan = *plan_; const auto values = out.GetValues(); - PCMS_ALWAYS_ASSERT(values.extent(0) == static_cast(n_points_)); + PCMS_ALWAYS_ASSERT(values.extent(0) == static_cast(plan.n_points)); PCMS_ALWAYS_ASSERT(values.extent(1) == static_cast(n_components_)); const int nc = n_components_; - const auto n_own = static_cast(own_points_.extent(0)); // Owner side: evaluate this rank's own points and the points sent to it. - const auto own_values = EvaluateLocally(own_evaluator_.get(), n_own, field); + const auto own_values = EvaluateLocally( + own_evaluator_.get(), static_cast(plan.own_points.extent(0)), field); const auto incoming_values = EvaluateLocally(incoming_evaluator_.get(), - static_cast(incoming_coords_.extent(0)), field); + static_cast(plan.IncomingPointCount()), field); // Route the incoming values back to the ranks that asked: the only values // that cross to the host and back, because MPI takes host buffers. const auto remote_host = ExchangeRecordsWithKnownCounts( - mesh_.comm()->get_impl(), detail::Scaled(incoming_counts_, nc), - CopyToHostRowMajor(incoming_values), detail::Scaled(remote_counts_, nc)); + plan.comm->get_impl(), detail::Scaled(plan.incoming_counts, nc), + CopyToHostRowMajor(incoming_values), + detail::Scaled(plan.remote_counts, nc)); const auto remote_values = CopyRowMajorToDevice(remote_host, nc); // Assemble on the device. A point no rank located keeps the fill value; a // point an owner could not place already holds it, from the owner's FILL. detail::FillAll(values, static_cast(policy_.fill_value)); - detail::ScatterRows(values, own_values, own_points_); - detail::ScatterRows(values, remote_values, remote_points_); + detail::ScatterRows(values, own_values, plan.own_points); + detail::ScatterRows(values, remote_values, plan.remote_points); // Thrown after the exchange, so every rank has completed it. if (policy_.mode == OutOfBoundsMode::ERROR && @@ -438,83 +402,175 @@ void DistributedPointEvaluator::Evaluate( } /** - * \brief Make a distributed evaluator for `space`, dispatching on the mesh - * dimension. Collective. + * \brief Creates distributed point evaluators for a function space on a + * partitioned mesh. + * + * The distributed counterpart of a function space's own evaluator factory, with + * the same `CreatePointEvaluator(EvaluationRequest)` call. Where a serial + * factory turns a request into a localization hint, this one turns it into a + * `DistributedRoutingPlan`, then builds local evaluators for the points each + * rank owns. * - * \param points Query coordinates, `[num_points][mesh_dim]`; copied into the - * layout the evaluator uses. - * \param grid_divisions Cells per axis of the global search grid; 0 derives it - * from the mesh. Must be the same on every rank. + * Unlike a serial factory, **construction and every `CreatePointEvaluator` / + * `BuildRoutingPlan` call are collective**: every rank of the mesh's + * communicator must make the same call, with its own points (possibly none). * - * The returned evaluator is an ordinary `PointEvaluator`, so it can be used - * wherever a serial one can. + * The search is built once, here, and reused for every request. A plan can be + * built once and shared by several evaluators: it depends only on the mesh + * partition and the points, so evaluators for other function spaces on the + * same mesh can reuse it through `CreatePointEvaluator(plan, request)`. */ -template -[[nodiscard]] std::unique_ptr> MakeDistributedPointEvaluator( - Omega_h::Mesh& mesh, std::shared_ptr space, - Kokkos::View points, - OutOfBoundsPolicy policy = {}, Real tolerance_reach = 1e-12, - LO grid_divisions = 0) +class DistributedPointEvaluatorFactory : public PointEvaluatorFactory { - const int mesh_dim = mesh.dim(); - const auto n_points = static_cast(points.extent(0)); - if (mesh_dim != 2 && mesh_dim != 3) { - throw std::invalid_argument( - "MakeDistributedPointEvaluator: only 2D and 3D meshes are supported"); +public: + /** + * \param mesh The partitioned mesh. Its ownership decides who answers for a + * point. Must outlive this factory. + * \param space The function space evaluators are created for. Its layout + * must live on `mesh`. + * \param tolerance_reach How far outside an element the search may still + * accept a point. + * \param grid_divisions Cells per axis of the global routing grid; 0 derives + * it from the mesh. Must be the same on every rank. + */ + DistributedPointEvaluatorFactory(Omega_h::Mesh& mesh, + std::shared_ptr space, + Real tolerance_reach = 1e-12, + LO grid_divisions = 0) + : mesh_(mesh), space_(std::move(space)) + { + if (!space_) { + throw std::invalid_argument( + "DistributedPointEvaluatorFactory: function space is null"); + } + // The same local search the Omega_h Lagrange factory builds, so for that + // backend the routing and the owner's local evaluator place every point + // in the same element. + local_search_ = + std::make_unique(detail::MakeSearch(mesh_)); + if (mesh_.dim() == 2) { + search_ = std::make_unique>( + mesh_, std::get(*local_search_), tolerance_reach, + grid_divisions); + } else { + search_ = std::make_unique>( + mesh_, std::get(*local_search_), tolerance_reach, + grid_divisions); + } } - if (points.extent(1) != static_cast(mesh_dim)) { - throw std::invalid_argument( - "MakeDistributedPointEvaluator: query coordinates are not [num_points x " - "mesh_dim]"); + + [[nodiscard]] const std::shared_ptr& + GetCoordinateSystem() const override + { + return space_->GetCoordinateSystem(); } - if (mesh_dim == 2) { - Kokkos::View coords("distributed coords", - n_points); - Kokkos::parallel_for( - "copy_distributed_coords", - Kokkos::RangePolicy(0, n_points), - KOKKOS_LAMBDA(LO i) { - coords(i, 0) = points(i, 0); - coords(i, 1) = points(i, 1); - }); - return std::make_unique>( - mesh, std::move(space), coords, policy, tolerance_reach, grid_divisions); + /// Locate `coords` over the decomposition and route them. Collective. + [[nodiscard]] std::shared_ptr BuildRoutingPlan( + const CoordinateView& coords) const + { + detail::CheckSameCoordinateSystem( + "DistributedPointEvaluatorFactory::BuildRoutingPlan", coords, + GetCoordinateSystem()); + const auto values = coords.GetValues(); + if (values.extent(1) != static_cast(mesh_.dim())) { + throw std::invalid_argument( + "DistributedPointEvaluatorFactory: query coordinates are not " + "[num_points x mesh_dim]"); + } + return std::visit( + [&](auto& search) { + constexpr int Dim = std::decay_t::LocalSearch::DIM; + return std::make_shared( + pcms::BuildRoutingPlan(*search, + detail::ToSearchCoords(values))); + }, + search_); } - Kokkos::View coords("distributed coords", - n_points); - Kokkos::parallel_for( - "copy_distributed_coords", - Kokkos::RangePolicy(0, n_points), - KOKKOS_LAMBDA(LO i) { - coords(i, 0) = points(i, 0); - coords(i, 1) = points(i, 1); - coords(i, 2) = points(i, 2); - }); - return std::make_unique>( - mesh, std::move(space), coords, policy, tolerance_reach, grid_divisions); -} + + using PointEvaluatorFactory::CreatePointEvaluator; + + /// An evaluator for `request`'s points, routed by an existing `plan`. + /// Collective. `plan` must have been built from the same points. + template + [[nodiscard]] std::unique_ptr> CreatePointEvaluator( + std::shared_ptr plan, + const EvaluationRequest& request) const + { + detail::CheckSameCoordinateSystem( + "DistributedPointEvaluatorFactory::CreatePointEvaluator", request.coords, + GetCoordinateSystem()); + CheckRequest(request); + if (!plan) { + throw std::invalid_argument( + "DistributedPointEvaluatorFactory: routing plan is null"); + } + const auto values = request.coords.GetValues(); + if (plan->comm.get() != mesh_.comm().get() || + plan->n_points != static_cast(values.extent(0)) || + plan->dim != static_cast(values.extent(1))) { + throw std::invalid_argument( + "DistributedPointEvaluatorFactory: the routing plan was not built for " + "these query points on this mesh"); + } + return std::make_unique>( + space_, std::move(plan), request); + } + + [[nodiscard]] const FunctionSpace& GetFunctionSpace() const noexcept + { + return *space_; + } + + [[nodiscard]] Omega_h::Mesh& GetMesh() const noexcept { return mesh_; } + +protected: + PointEvaluatorVariant CreatePointEvaluatorImpl( + Type value_type, const EvaluationRequest& request) const override + { + CheckRequest(request); + auto plan = BuildRoutingPlan(request.coords); + switch (value_type) { + case Type::Real: return CreatePointEvaluator(plan, request); + case Type::LO: return CreatePointEvaluator(plan, request); + case Type::GO: return CreatePointEvaluator(plan, request); + case Type::Int8: return CreatePointEvaluator(plan, request); + case Type::Float: return CreatePointEvaluator(plan, request); + } + throw pcms_error( + "DistributedPointEvaluatorFactory: unsupported value type"); + } + +private: + static void CheckRequest(const EvaluationRequest& request) + { + if (request.policy.mode == OutOfBoundsMode::NEAREST_BOUNDARY) { + throw pcms_error("DistributedPointEvaluatorFactory: NEAREST_BOUNDARY is " + "not supported"); + } + } + + Omega_h::Mesh& mesh_; + std::shared_ptr space_; + // Declared before `search_`, which holds a reference to it. + std::unique_ptr local_search_; + std::variant>, + std::unique_ptr>> + search_; +}; /// How many query points of a distributed evaluator get their value from /// another rank, or 0 when `evaluator` is not one. -/// -/// Diagnostics for callers that hold only the `PointEvaluator` interface, which -/// is the interface a distributed evaluator is meant to be used through. template [[nodiscard]] std::size_t LastRoutedPointCount( const PointEvaluator& evaluator) { - if (auto* d2 = - dynamic_cast*>(&evaluator)) { - return d2->LastRoutedPointCount(); - } - if (auto* d3 = - dynamic_cast*>(&evaluator)) { - return d3->LastRoutedPointCount(); + if (auto* d = dynamic_cast*>(&evaluator)) { + return d->LastRoutedPointCount(); } return 0; } } // namespace pcms -#endif // PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H \ No newline at end of file +#endif // PCMS_FIELD_EVALUATOR_DISTRIBUTED_POINT_EVALUATOR_H diff --git a/src/pcms/localization/CMakeLists.txt b/src/pcms/localization/CMakeLists.txt index 4ec72a590..73dbe6097 100644 --- a/src/pcms/localization/CMakeLists.txt +++ b/src/pcms/localization/CMakeLists.txt @@ -4,6 +4,7 @@ set( overlap.h candidate_rank_table.h distributed_point_search.h + distributed_routing_plan.h queue_visited.hpp adj_search.hpp mls_support_helpers.h diff --git a/src/pcms/localization/distributed_point_search.cpp b/src/pcms/localization/distributed_point_search.cpp index fd93c9f48..1c2e532ae 100644 --- a/src/pcms/localization/distributed_point_search.cpp +++ b/src/pcms/localization/distributed_point_search.cpp @@ -81,7 +81,7 @@ DistributedPointSearch::DistributedPointSearch( } // Global ids identify an element across ranks. They are what lets the answer - // to "which element" be checked, and what a caller caches; a mesh without + // to "which element" be checked across ranks; a mesh without // them still works, the id is just unavailable. if (mesh.has_tag(mesh.dim(), "global")) { const auto globals = Omega_h::HostRead(mesh.globals(mesh.dim())); @@ -131,15 +131,10 @@ DistributedPointSearch::ClassifyLocally( template std::vector::Result> -DistributedPointSearch::Locate(Kokkos::View points, - std::vector* cache) +DistributedPointSearch::Locate(Kokkos::View points) { const auto n_points = static_cast(points.extent(0)); - std::vector storage; - std::vector& results = cache != nullptr ? *cache : storage; - // Resizing keeps entries already resolved, so a caller that passes the same - // cache for the same point set only pays for the points still missing. - results.resize(static_cast(n_points)); + std::vector results(static_cast(n_points)); const auto classified = ClassifyLocally(points); const auto points_h = @@ -153,9 +148,6 @@ DistributedPointSearch::Locate(Kokkos::View points, last_routed_ = 0; for (LO i = 0; i < n_points; ++i) { auto& r = results[static_cast(i)]; - if (r.status == ResultStatus::LOCATED) { - continue; // cached - } const auto& c = classified[static_cast(i)]; if (c.status == LocalStatus::LOCATED_HERE) { r.status = ResultStatus::LOCATED; diff --git a/src/pcms/localization/distributed_point_search.h b/src/pcms/localization/distributed_point_search.h index 223c31bdd..bcfe4865f 100644 --- a/src/pcms/localization/distributed_point_search.h +++ b/src/pcms/localization/distributed_point_search.h @@ -126,24 +126,19 @@ class DistributedPointSearch * from the **lowest rank id** wins. Any of them is geometrically right; the * rule just makes the outcome deterministic and identical on every rank. * + * To evaluate at the same points repeatedly, keep the result rather than + * locating again: `DistributedRoutingPlan` is built from it once and reused. + * * \param points Query coordinates, one per row. - * \param cache Optional, indexed by point. Entries that are already `LOCATED` - * are reused and not searched again, which is what makes repeated - * localization of the same point set cheap. The caller owns the cache - * and must clear it whenever the point coordinates or the mesh change: - * an entry is only valid for the point and mesh it was computed from. - * Points that could not be placed are not cached, so they are retried. * \return One result per row, in the same order. */ [[nodiscard]] std::vector Locate( - Kokkos::View points, - std::vector* cache = nullptr); + Kokkos::View points); /// How many points the most recent `Locate` sent to another rank. /// /// Diagnostics. Zero means every point was answered here, either because this - /// rank owns it, because it lies outside the mesh, or because the cache - /// already held its location. + /// rank owns it or because it lies outside the mesh. [[nodiscard]] std::size_t LastRoutedPointCount() const noexcept { return last_routed_; diff --git a/src/pcms/localization/distributed_routing_plan.h b/src/pcms/localization/distributed_routing_plan.h new file mode 100644 index 000000000..68c2696ad --- /dev/null +++ b/src/pcms/localization/distributed_routing_plan.h @@ -0,0 +1,163 @@ +#ifndef PCMS_LOCALIZATION_DISTRIBUTED_ROUTING_PLAN_H +#define PCMS_LOCALIZATION_DISTRIBUTED_ROUTING_PLAN_H + +#include "pcms/localization/distributed_point_search.h" +#include "pcms/utility/arrays.h" +#include "pcms/utility/memory_spaces.h" +#include "pcms/utility/mpi_utilities.h" +#include "pcms/utility/types.h" + +#include +#include + +#include +#include + +namespace pcms +{ + +/** + * \brief Which rank answers for each of this rank's query points, and which + * points this rank answers for the others. + * + * This is the distributed counterpart of a serial evaluator's localization + * hint. A serial hint records where each point is inside the local mesh, in a + * form only its own backend understands. A routing plan records only which + * rank owns each point. The owner then localizes the point with whatever + * backend it uses, so the plan does not depend on the field or the function + * space. One plan can be shared by evaluators of several function spaces on + * the same mesh partition. + * + * Built by `BuildRoutingPlan`, which is collective. The plan is valid for the + * query points and the mesh partition it was built from. + */ +struct DistributedRoutingPlan +{ + /// Communicator of the mesh the plan was built on. Every exchange that uses + /// the plan runs on it. + Omega_h::CommPtr comm; + /// Spatial dimension of the query points. + int dim = 0; + /// Number of query points on this rank. + LO n_points = 0; + + /// Query points this rank owns, as indices into the query set, ascending. + Kokkos::View own_points; + + /// Query points another rank owns, grouped by owner rank in ascending order. + Kokkos::View remote_points; + /// How many of `remote_points` go to each rank. + std::vector remote_counts; + + /// Coordinates other ranks sent this rank to evaluate, `[n][dim]`, grouped + /// by requesting rank. + Kokkos::View incoming_coords; + /// How many of `incoming_coords` came from each rank. + std::vector incoming_counts; + + /// Query points no rank could locate, ascending. + Kokkos::View unlocated_points; + + /// True when this rank owns every one of its query points: nothing is sent + /// away and nothing is unlocated. Other ranks may still send points here. + [[nodiscard]] bool AllOwnedHere() const noexcept + { + return static_cast(own_points.extent(0)) == n_points; + } + + /// How many query points get their value from another rank. + [[nodiscard]] std::size_t RoutedPointCount() const noexcept + { + return remote_points.extent(0); + } + + /// How many points other ranks asked this rank to evaluate. + [[nodiscard]] std::size_t IncomingPointCount() const noexcept + { + return incoming_coords.extent(0); + } +}; + +namespace detail +{ + +/// `counts` times `factor`, for sending `factor` scalars per point. +inline std::vector Scaled(const std::vector& counts, int factor) +{ + std::vector out(counts.size()); + for (size_t r = 0; r < counts.size(); ++r) { + out[r] = counts[r] * factor; + } + return out; +} + +} // namespace detail + +/** + * \brief Locate `points` over the whole decomposition and send each point + * another rank owns to that rank. Collective. + * + * Every rank must call it, with its own points (possibly none). + */ +template +[[nodiscard]] DistributedRoutingPlan BuildRoutingPlan( + DistributedPointSearch& search, + Kokkos::View points) +{ + using Search = DistributedPointSearch; + Omega_h::Mesh& mesh = search.GetMesh(); + + DistributedRoutingPlan plan; + plan.comm = mesh.comm(); + plan.dim = Dim; + plan.n_points = static_cast(points.extent(0)); + + const auto located = search.Locate(points); + + // Split the query indices by owner. The search result is host data, so this + // is host work on indices only; coordinates stay on the device. + const int rank = search.Rank(); + const int nranks = plan.comm->size(); + std::vector> by_owner(static_cast(nranks)); + std::vector own; + std::vector unlocated; + for (LO i = 0; i < plan.n_points; ++i) { + const auto& r = located[static_cast(i)]; + if (r.status != Search::ResultStatus::LOCATED) { + unlocated.push_back(i); + } else if (r.owner_rank == rank) { + own.push_back(i); + } else { + by_owner[static_cast(r.owner_rank)].push_back(i); + } + } + std::vector remote; + plan.remote_counts.assign(static_cast(nranks), 0); + for (int r = 0; r < nranks; ++r) { + const auto& points_of_r = by_owner[static_cast(r)]; + remote.insert(remote.end(), points_of_r.begin(), points_of_r.end()); + plan.remote_counts[static_cast(r)] = + static_cast(points_of_r.size()); + } + plan.own_points = CopyToDevice(own); + plan.remote_points = CopyToDevice(remote); + plan.unlocated_points = CopyToDevice(unlocated); + + // Send each remote point to its owner, once. Only the remote points' + // coordinates leave the device, because MPI takes host buffers. + const auto send = CopyToHostRowMajor(GatherRows(points, plan.remote_points)); + std::vector incoming_scalar_counts; + const auto incoming = ExchangeRecords(plan.comm->get_impl(), + detail::Scaled(plan.remote_counts, Dim), + send, incoming_scalar_counts); + plan.incoming_counts.resize(incoming_scalar_counts.size()); + for (size_t r = 0; r < plan.incoming_counts.size(); ++r) { + plan.incoming_counts[r] = incoming_scalar_counts[r] / Dim; + } + plan.incoming_coords = CopyRowMajorToDevice(incoming, Dim); + return plan; +} + +} // namespace pcms + +#endif // PCMS_LOCALIZATION_DISTRIBUTED_ROUTING_PLAN_H diff --git a/test/test_distributed_evaluate.cpp b/test/test_distributed_evaluate.cpp index 22aa053ba..6b9dee717 100644 --- a/test/test_distributed_evaluate.cpp +++ b/test/test_distributed_evaluate.cpp @@ -111,6 +111,18 @@ QueryCoords MakeQueryCoords(const std::vector& interleaved, int dim) pcms::csys::Cartesian::Deferred(), pcms::MakeRank2View(device))}; } +/// A distributed evaluator through the public factory API, as a caller with +/// one point set would build it. +std::unique_ptr> MakeDistributed( + Omega_h::Mesh& mesh, std::shared_ptr space, + const pcms::CoordinateView& coords, + pcms::OutOfBoundsPolicy policy) +{ + pcms::DistributedPointEvaluatorFactory factory(mesh, std::move(space)); + return factory.CreatePointEvaluator( + pcms::EvaluationRequest::FromCoordinates(coords, policy)); +} + /// Distributed evaluation must equal the serial evaluator wherever the serial /// evaluator is defined, and must equal the closed form everywhere. template @@ -148,8 +160,8 @@ void check_agrees_with_serial(Omega_h::Mesh& mesh, int order) // Distributed, from the same space and points. The space, not a layout, is // what is handed over: the distributed evaluator reaches a backend only // through the local evaluator the space creates. - auto distributed = pcms::MakeDistributedPointEvaluator( - mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + auto distributed = + MakeDistributed(mesh, space, device_coords.view, pcms::OutOfBoundsPolicy{}); Kokkos::View got("got", n_points, 1); distributed->Evaluate(field, distributed_test::TagLike(field, got)); @@ -330,8 +342,8 @@ void check_boundary_arbitration(Omega_h::Mesh& mesh) const auto reference_h = Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), reference); - auto distributed = pcms::MakeDistributedPointEvaluator( - mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + auto distributed = + MakeDistributed(mesh, space, device_coords.view, pcms::OutOfBoundsPolicy{}); Kokkos::View got("got", n_points, 1); distributed->Evaluate(field, distributed_test::TagLike(field, got)); const auto got_h = @@ -483,8 +495,8 @@ void check_backend_without_distributed_code(Omega_h::Mesh& mesh) Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace(), reference); // The same backend, distributed. - auto distributed = pcms::MakeDistributedPointEvaluator( - mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + auto distributed = + MakeDistributed(mesh, space, device_coords.view, pcms::OutOfBoundsPolicy{}); Kokkos::View got("got", n_points, 1); distributed->Evaluate(field, distributed_test::TagLike(field, got)); const std::size_t routed = pcms::LastRoutedPointCount(*distributed); @@ -598,8 +610,7 @@ void check_out_of_bounds_policy(Omega_h::Mesh& mesh) pcms::OutOfBoundsPolicy fill_policy; fill_policy.mode = pcms::OutOfBoundsMode::FILL; fill_policy.fill_value = fill_value; - auto filled = pcms::MakeDistributedPointEvaluator( - mesh, space, device_coords.storage, fill_policy); + auto filled = MakeDistributed(mesh, space, device_coords.view, fill_policy); Kokkos::View got("got", n_points, 1); filled->Evaluate(field, distributed_test::TagLike(field, got)); const auto got_h = @@ -650,8 +661,8 @@ void check_out_of_bounds_policy(Omega_h::Mesh& mesh) error_policy.mode = pcms::OutOfBoundsMode::ERROR; bool threw = false; try { - auto strict = pcms::MakeDistributedPointEvaluator( - mesh, space, device_coords.storage, error_policy); + auto strict = + MakeDistributed(mesh, space, device_coords.view, error_policy); Kokkos::View strict_values("strict", n_points, 1); strict->Evaluate(field, distributed_test::TagLike(field, strict_values)); @@ -713,8 +724,8 @@ void check_ghosts_must_be_synchronized(Omega_h::Mesh& mesh) const auto gathered = distributed_test::GatherOwnedCentroids(mesh); auto device_coords = MakeQueryCoords(gathered.coords, Dim); const LO n_points = static_cast(gathered.total); - auto distributed = pcms::MakeDistributedPointEvaluator( - mesh, space, device_coords.storage, pcms::OutOfBoundsPolicy{}); + auto distributed = + MakeDistributed(mesh, space, device_coords.view, pcms::OutOfBoundsPolicy{}); const auto count_wrong = [&]() { Kokkos::View got("got", n_points, 1); @@ -745,6 +756,131 @@ void check_ghosts_must_be_synchronized(Omega_h::Mesh& mesh) require(wrong_after_sync == 0, "distributed evaluation is wrong after ghosts are synchronized"); } +/// Values of `evaluator` at its `n` points, on the host. +auto EvaluateToHost(const pcms::PointEvaluator& evaluator, + const pcms::Field& field, LO n) +{ + Kokkos::View got("got", n, 1); + evaluator.Evaluate(field, distributed_test::TagLike(field, got)); + return Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace{}, got); +} + +/// One routing plan serves evaluators for two function spaces on the same +/// mesh, and gives exactly what a one-shot request gives. +template +void check_shared_routing_plan(Omega_h::Mesh& mesh) +{ + auto linear_space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/1, 1, pcms::csys::Cartesian::Deferred(), "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto constant_space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/0, 1, pcms::csys::Cartesian::Deferred(), "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto linear_field = linear_space->CreateFunction(); + auto constant_field = constant_space->CreateFunction(); + SetLinearField(linear_field.GetData(), *linear_space->GetLayout()); + SetLinearField(constant_field.GetData(), *constant_space->GetLayout()); + + const auto gathered = distributed_test::GatherOwnedCentroids(mesh); + auto device_coords = MakeQueryCoords(gathered.coords, Dim); + const LO n_points = static_cast(gathered.total); + const auto request = + pcms::EvaluationRequest::FromCoordinates(device_coords.view); + + pcms::DistributedPointEvaluatorFactory linear_factory(mesh, linear_space); + pcms::DistributedPointEvaluatorFactory constant_factory(mesh, constant_space); + + // One plan, two spaces. + const auto plan = linear_factory.BuildRoutingPlan(device_coords.view); + auto linear_shared = linear_factory.CreatePointEvaluator(plan, request); + auto constant_shared = + constant_factory.CreatePointEvaluator(plan, request); + + // The same requests, each building its own plan. + auto linear_own = linear_factory.CreatePointEvaluator(request); + auto constant_own = constant_factory.CreatePointEvaluator(request); + + require(pcms::LastRoutedPointCount(*linear_shared) == + plan->RoutedPointCount(), + "the evaluator does not report its plan's routed point count"); + + const auto a = EvaluateToHost(*linear_shared, linear_field, n_points); + const auto b = EvaluateToHost(*linear_own, linear_field, n_points); + const auto c = EvaluateToHost(*constant_shared, constant_field, n_points); + const auto d = EvaluateToHost(*constant_own, constant_field, n_points); + + LO n_bad = 0; + for (LO i = 0; i < n_points; ++i) { + const size_t base = static_cast(i) * Dim; + const Real truth = + linear_f(gathered.coords[base], gathered.coords[base + 1]); + n_bad += a(i, 0) != b(i, 0) ? 1 : 0; + n_bad += c(i, 0) != d(i, 0) ? 1 : 0; + n_bad += std::fabs(a(i, 0) - truth) > 1e-10 ? 1 : 0; + } + const LO t_bad = mesh.comm()->allreduce(n_bad, Omega_h_Op::OMEGA_H_SUM); + const auto t_routed = mesh.comm()->allreduce( + static_cast(plan->RoutedPointCount()), Omega_h_Op::OMEGA_H_SUM); + if (mesh.comm()->rank() == 0) { + std::printf("%dD shared plan: %lld points routed, %lld disagreements\n", + Dim, static_cast(t_routed), + static_cast(t_bad)); + } + require(t_routed > 0, "no point was routed, so the plan was not exercised"); + require( + t_bad == 0, + "an evaluator built from a shared plan disagrees with a one-shot one"); +} + +/// A request built from a layout: the query points are this rank's local +/// element centroids, ghosts included, so some are owned by another rank. +template +void check_layout_request(Omega_h::Mesh& mesh) +{ + auto space = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/1, 1, pcms::csys::Cartesian::Deferred(), "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto centroids = pcms::LagrangeFunctionSpace::FromMesh( + mesh, /*order=*/0, 1, pcms::csys::Cartesian::Deferred(), "global", + pcms::LagrangeFunctionSpace::Backend::OmegaH); + auto field = space->CreateFunction(); + SetLinearField(field.GetData(), *space->GetLayout()); + + pcms::DistributedPointEvaluatorFactory factory(mesh, space); + const auto request = pcms::EvaluationRequest::FromFunctionSpace(*centroids); + auto evaluator = factory.CreatePointEvaluator(request); + + const auto query = request.coords.GetValues(); + const LO n_points = static_cast(query.extent(0)); + Kokkos::View query_d("query", n_points, Dim); + Kokkos::parallel_for( + "copy_query", + Kokkos::RangePolicy(0, n_points), + KOKKOS_LAMBDA(LO i) { + for (int d = 0; d < Dim; ++d) { + query_d(i, d) = query(i, d); + } + }); + const auto query_h = + Kokkos::create_mirror_view_and_copy(pcms::HostMemorySpace{}, query_d); + const auto got = EvaluateToHost(*evaluator, field, n_points); + + LO n_bad = 0; + for (LO i = 0; i < n_points; ++i) { + const Real truth = linear_f(query_h(i, 0), query_h(i, 1)); + n_bad += std::fabs(got(i, 0) - truth) > 1e-10 ? 1 : 0; + } + const LO t_bad = mesh.comm()->allreduce(n_bad, Omega_h_Op::OMEGA_H_SUM); + const auto t_routed = mesh.comm()->allreduce( + static_cast(pcms::LastRoutedPointCount(*evaluator)), + Omega_h_Op::OMEGA_H_SUM); + if (mesh.comm()->rank() == 0) { + std::printf("%dD layout request: %lld ghost centroids routed, %lld wrong\n", + Dim, static_cast(t_routed), + static_cast(t_bad)); + } + require(t_bad == 0, "a layout request gave a wrong value"); +} } // namespace int main(int argc, char** argv) @@ -764,6 +900,8 @@ int main(int argc, char** argv) check_backend_without_distributed_code<2>(mesh); check_out_of_bounds_policy<2>(mesh); check_ghosts_must_be_synchronized<2>(mesh); + check_shared_routing_plan<2>(mesh); + check_layout_request<2>(mesh); } { @@ -771,6 +909,8 @@ int main(int argc, char** argv) check_agrees_with_serial<3>(mesh, /*order=*/1); check_ghosts_must_be_synchronized<3>(mesh); check_out_of_bounds_policy<3>(mesh); + check_shared_routing_plan<3>(mesh); + check_layout_request<3>(mesh); } if (world->rank() == 0) { diff --git a/test/test_distributed_search.cpp b/test/test_distributed_search.cpp index 0481b09e9..d344cfa04 100644 --- a/test/test_distributed_search.cpp +++ b/test/test_distributed_search.cpp @@ -6,8 +6,8 @@ // 2. The local pass (`ClassifyLocally`): a rank answers only for elements it // owns, never for ghosts. // 3. The full search (`Locate`): every point is attributed to its owning rank -// and element, including points that had to be routed, and a warm cache -// routes nothing. +// and element, including points that had to be routed, and repeating a +// batch gives the same answer. #include "distributed_test_utils.h" #include "pcms/localization/candidate_rank_table.h" @@ -755,13 +755,11 @@ void check_distributed_locate(Omega_h::Mesh& mesh) require(t_routed_away > 0, "no query left its rank, so routing was not exercised"); - // Determinism and caching: the same batch twice must give the same answer, - // and a second pass over a warm cache must not re-search anything it - // resolved. - std::vector::Result> cache; - const auto pass_one = search.Locate(points, &cache); + // Determinism: the same batch twice must give the same answer and the same + // routing work. + const auto pass_one = search.Locate(points); const auto routed_one = search.LastRoutedPointCount(); - const auto pass_two = search.Locate(points, &cache); + const auto pass_two = search.Locate(points); const auto routed_two = search.LastRoutedPointCount(); LO n_unresolved_results = 0; @@ -781,19 +779,15 @@ void check_distributed_locate(Omega_h::Mesh& mesh) } if (rank == 0) { - std::printf("routed %zu points on the first pass, %zu on the second " - "(only unresolved points are retried, and points outside the " - "mesh bounding box are never routed)\n", + std::printf("routed %zu points on the first pass, %zu on the second\n", routed_one, routed_two); } require(routed_one > 0, "the first pass routed nothing"); - // The only point still unresolved on a warm cache is the one outside the - // grid, and that one needs no message, so a warm pass routes nothing at all. + require(routed_two == routed_one, + "repeating a batch changed the routing work"); require(n_unresolved_results > 0, "the out-of-domain probe was located"); - require(routed_two == 0, "a cached pass still routed points"); - require(routed_two < routed_one, "caching did not reduce the routing work"); require(mesh.comm()->allreduce(n_unstable, Omega_h_Op::OMEGA_H_SUM) == 0, - "a cached pass gave a different answer"); + "repeating a batch gave a different answer"); // A batch of points outside the global bounding box needs no communication at // all: no element can contain them, so not a single query is routed.