From 8d5b072e60d79536e01af8d6e9964c89613609ec Mon Sep 17 00:00:00 2001 From: Willem Deconinck Date: Wed, 13 Sep 2023 09:55:34 +0200 Subject: [PATCH] Use mpi_comm in functionspace::PointCloud --- src/atlas/functionspace/PointCloud.cc | 75 +++++++++++-------- src/atlas/functionspace/PointCloud.h | 18 +++-- .../test_pointcloud_haloexchange_2PE.cc | 58 ++++++++++---- 3 files changed, 99 insertions(+), 52 deletions(-) diff --git a/src/atlas/functionspace/PointCloud.cc b/src/atlas/functionspace/PointCloud.cc index ab92e1125..a62a53233 100644 --- a/src/atlas/functionspace/PointCloud.cc +++ b/src/atlas/functionspace/PointCloud.cc @@ -45,8 +45,16 @@ namespace functionspace { namespace detail { +static std::string get_mpi_comm(const eckit::Configuration& config) { + if(config.has("mpi_comm")) { + return config.getString("mpi_comm"); + } + return mpi::comm().name(); +} + template <> PointCloud::PointCloud(const std::vector& points, const eckit::Configuration& config) { + mpi_comm_ = get_mpi_comm(config); lonlat_ = Field("lonlat", array::make_datatype(), array::make_shape(points.size(), 2)); auto lonlat = array::make_view(lonlat_); for (idx_t j = 0, size = points.size(); j < size; ++j) { @@ -57,6 +65,7 @@ PointCloud::PointCloud(const std::vector& points, const eckit::Configur template <> PointCloud::PointCloud(const std::vector& points, const eckit::Configuration& config) { + mpi_comm_ = get_mpi_comm(config); lonlat_ = Field("lonlat", array::make_datatype(), array::make_shape(points.size(), 2)); vertical_ = Field("vertical", array::make_datatype(), array::make_shape(points.size())); auto lonlat = array::make_view(lonlat_); @@ -68,13 +77,17 @@ PointCloud::PointCloud(const std::vector& points, const eckit::Configu } } -PointCloud::PointCloud(const Field& lonlat, const eckit::Configuration& config): lonlat_(lonlat) {} +PointCloud::PointCloud(const Field& lonlat, const eckit::Configuration& config): lonlat_(lonlat) { + mpi_comm_ = get_mpi_comm(config); +} PointCloud::PointCloud(const Field& lonlat, const Field& ghost, const eckit::Configuration& config): lonlat_(lonlat), ghost_(ghost) { + mpi_comm_ = get_mpi_comm(config); setupHaloExchange(); } PointCloud::PointCloud(const FieldSet& flds, const eckit::Configuration& config): lonlat_(flds["lonlat"]) { + mpi_comm_ = get_mpi_comm(config); if (flds.has("ghost")) { ghost_ = flds["ghost"]; } @@ -93,6 +106,7 @@ PointCloud::PointCloud(const FieldSet& flds, const eckit::Configuration& config) } grid::Partitioner make_partitioner(const Grid& grid, const eckit::Configuration& config) { + auto mpi_comm = get_mpi_comm(config); auto partitioner = grid.partitioner(); if( config.has("partitioner") ) { partitioner.set("type",config.getString("partitioner")); @@ -100,6 +114,7 @@ grid::Partitioner make_partitioner(const Grid& grid, const eckit::Configuration& if( not partitioner.has("type") ) { partitioner.set("type","equal_regions"); } + partitioner.set("mpi_comm",mpi_comm); return grid::Partitioner(partitioner); } @@ -109,21 +124,23 @@ PointCloud::PointCloud(const Grid& grid, const eckit::Configuration& config) : PointCloud::PointCloud(const Grid& grid, const grid::Partitioner& _partitioner, const eckit::Configuration& config) { ATLAS_TRACE("PointCloud(grid,partitioner,config)"); - int part = mpi::rank(); + mpi_comm_ = get_mpi_comm(config); + auto& comm = mpi::comm(mpi_comm_); double halo_radius; config.get("halo_radius", halo_radius = 0.); grid::Partitioner partitioner(_partitioner); if ( not partitioner ) { - partitioner = grid::Partitioner("equal_regions"); + partitioner = grid::Partitioner("equal_regions", util::Config("mpi_comm",mpi_comm_)); } + part_ = comm.rank(); nb_partitions_ = partitioner.nb_partitions(); auto distribution = partitioner.partition(grid); - auto size_owned = distribution.nb_pts()[part]; + auto size_owned = distribution.nb_pts()[part_]; size_owned_ = size_owned; - if (halo_radius == 0. || mpi::size() == 1) { + if (halo_radius == 0. || nb_partitions_ == 1) { idx_t size_halo = size_owned; ATLAS_ASSERT(size_owned > 0); lonlat_ = Field("lonlat", array::make_datatype(), array::make_shape(size_halo, 2)); @@ -136,12 +153,12 @@ PointCloud::PointCloud(const Grid& grid, const grid::Partitioner& _partitioner, auto ridx = array::make_indexview(remote_index_); auto gidx = array::make_view(global_index_); array::make_view(ghost_).assign(0); - array::make_view(partition_).assign(part); + array::make_view(partition_).assign(part_); idx_t j{0}; gidx_t g{0}; for (auto p : grid.lonlat()) { - if( distribution.partition(g++) == part ) { + if( distribution.partition(g++) == part_ ) { gidx(j) = g+1; ridx(j) = j; lonlat(j, 0) = p.lon(); @@ -167,7 +184,7 @@ PointCloud::PointCloud(const Grid& grid, const grid::Partitioner& _partitioner, kdtree.reserve(grid.size()); idx_t j{0}; for (auto p : grid.lonlat()) { - if( distribution.partition(j) == part ) { + if( distribution.partition(j) == part_ ) { owned_lonlat.emplace_back(p); owned_grid_idx.emplace_back(j); } @@ -213,7 +230,7 @@ PointCloud::PointCloud(const Grid& grid, const grid::Partitioner& _partitioner, lonlat(j, 0) = p.lon(); lonlat(j, 1) = p.lat(); partition(j) = distribution.partition(g); - ghost(j) = partition(j) != part; + ghost(j) = partition(j) != part_; glb_idx(j) = g+1; ++j; } @@ -446,7 +463,7 @@ void PointCloud::haloExchange(const Field& field, bool on_device) const { void PointCloud::create_remote_index() const { ATLAS_TRACE(); - const eckit::mpi::Comm& comm = atlas::mpi::comm(); + const auto& comm = mpi::comm(mpi_comm_); const int mpi_rank = comm.rank(); const int mpi_size = comm.size(); auto size_halo = lonlat_.shape(0); @@ -690,16 +707,14 @@ void PointCloud::create_remote_index() const { void PointCloud::setupHaloExchange() { ATLAS_TRACE(); - const eckit::mpi::Comm& comm = atlas::mpi::comm(); - const int mpi_rank = comm.rank(); - const int mpi_size = comm.size(); - - if (ghost_ and partition_ and global_index_ and not remote_index_) { create_remote_index(); } else if (not partition_ or not remote_index_) { ATLAS_TRACE("do setup"); + const auto& comm = mpi::comm(mpi_comm_); + const int mpi_rank = comm.rank(); + const int mpi_size = comm.size(); auto lonlat_v = array::make_view(lonlat_); // data structure containing a flag to identify the 'ghost points'; @@ -826,11 +841,11 @@ void PointCloud::setupHaloExchange() { ATLAS_ASSERT(ghost_.size() == partition_.size()); halo_exchange_.reset(new parallel::HaloExchange()); - halo_exchange_->setup(array::make_view(partition_).data(), + halo_exchange_->setup(mpi_comm_, + array::make_view(partition_).data(), array::make_view(remote_index_).data(), REMOTE_IDX_BASE, ghost_.size()); - } void PointCloud::adjointHaloExchange(const FieldSet& fieldset, bool on_device) const { @@ -869,25 +884,25 @@ void PointCloud::adjointHaloExchange(const Field& field, bool) const { PointCloud::PointCloud(const FunctionSpace& functionspace): FunctionSpace(functionspace), functionspace_(dynamic_cast(get())) {} -PointCloud::PointCloud(const Field& field): - FunctionSpace(new detail::PointCloud(field)), functionspace_(dynamic_cast(get())) {} +PointCloud::PointCloud(const Field& field, const eckit::Configuration& config): + FunctionSpace(new detail::PointCloud(field,config)), functionspace_(dynamic_cast(get())) {} -PointCloud::PointCloud(const Field& field1, const Field& field2): - FunctionSpace(new detail::PointCloud(field1, field2)), functionspace_(dynamic_cast(get())) {} +PointCloud::PointCloud(const Field& field1, const Field& field2, const eckit::Configuration& config): + FunctionSpace(new detail::PointCloud(field1, field2, config)), functionspace_(dynamic_cast(get())) {} -PointCloud::PointCloud(const FieldSet& fset): - FunctionSpace(new detail::PointCloud(fset)), functionspace_(dynamic_cast(get())) {} +PointCloud::PointCloud(const FieldSet& fset, const eckit::Configuration& config): + FunctionSpace(new detail::PointCloud(fset, config)), functionspace_(dynamic_cast(get())) {} -PointCloud::PointCloud(const std::vector& points): - FunctionSpace(new detail::PointCloud(points)), functionspace_(dynamic_cast(get())) {} +PointCloud::PointCloud(const std::vector& points, const eckit::Configuration& config): + FunctionSpace(new detail::PointCloud(points, config)), functionspace_(dynamic_cast(get())) {} -PointCloud::PointCloud(const std::vector& points): - FunctionSpace(new detail::PointCloud(points)), functionspace_(dynamic_cast(get())) {} +PointCloud::PointCloud(const std::vector& points, const eckit::Configuration& config): + FunctionSpace(new detail::PointCloud(points, config)), functionspace_(dynamic_cast(get())) {} -PointCloud::PointCloud(const std::initializer_list>& points): +PointCloud::PointCloud(const std::initializer_list>& points, const eckit::Configuration& config): FunctionSpace((points.begin()->size() == 2 - ? new detail::PointCloud{std::vector(points.begin(), points.end())} - : new detail::PointCloud{std::vector(points.begin(), points.end())})), + ? new detail::PointCloud{std::vector(points.begin(), points.end()), config} + : new detail::PointCloud{std::vector(points.begin(), points.end()), config})), functionspace_(dynamic_cast(get())) {} PointCloud::PointCloud(const Grid& grid, const eckit::Configuration& config): diff --git a/src/atlas/functionspace/PointCloud.h b/src/atlas/functionspace/PointCloud.h index b7c1fea81..f603a3146 100644 --- a/src/atlas/functionspace/PointCloud.h +++ b/src/atlas/functionspace/PointCloud.h @@ -61,7 +61,9 @@ class PointCloud : public functionspace::FunctionSpaceImpl { Field global_index() const override { return global_index_; } Field partition() const override { return partition_; } idx_t size() const override { return lonlat_.shape(0); } - idx_t nb_partitions() const override { return nb_partitions_; } + idx_t part() const override { return part_; } + idx_t nb_parts() const override { return nb_partitions_; } + std::string mpi_comm() const override { return mpi_comm_; } using FunctionSpaceImpl::createField; Field createField(const eckit::Configuration&) const override; @@ -166,7 +168,9 @@ class PointCloud : public functionspace::FunctionSpaceImpl { idx_t max_glb_idx_{0}; std::unique_ptr halo_exchange_; idx_t levels_{0}; + idx_t part_{0}; idx_t nb_partitions_{1}; + std::string mpi_comm_; void setupHaloExchange(); @@ -181,12 +185,12 @@ class PointCloud : public functionspace::FunctionSpaceImpl { class PointCloud : public FunctionSpace { public: PointCloud(const FunctionSpace&); - PointCloud(const Field& points); - PointCloud(const Field&, const Field&); - PointCloud(const FieldSet& flds); - PointCloud(const std::vector&); - PointCloud(const std::vector&); - PointCloud(const std::initializer_list>&); + PointCloud(const Field& points, const eckit::Configuration& = util::NoConfig()); + PointCloud(const Field&, const Field&, const eckit::Configuration& = util::NoConfig()); + PointCloud(const FieldSet& flds, const eckit::Configuration& = util::NoConfig()); + PointCloud(const std::vector&, const eckit::Configuration& = util::NoConfig()); + PointCloud(const std::vector&, const eckit::Configuration& = util::NoConfig()); + PointCloud(const std::initializer_list>&, const eckit::Configuration& = util::NoConfig()); PointCloud(const Grid&, const eckit::Configuration& = util::NoConfig()); PointCloud(const Grid&, const grid::Partitioner&, const eckit::Configuration& = util::NoConfig()); diff --git a/src/tests/functionspace/test_pointcloud_haloexchange_2PE.cc b/src/tests/functionspace/test_pointcloud_haloexchange_2PE.cc index 35a93b268..d20c76054 100644 --- a/src/tests/functionspace/test_pointcloud_haloexchange_2PE.cc +++ b/src/tests/functionspace/test_pointcloud_haloexchange_2PE.cc @@ -35,8 +35,8 @@ double innerproductwithhalo(const Field& f1, const Field& f2) { sum += view1(jn, jl) * view2(jn, jl); } } - - atlas::mpi::comm().allReduceInPlace(sum, eckit::mpi::sum()); + auto& comm = atlas::mpi::comm(f1.functionspace().mpi_comm()); + comm.allReduceInPlace(sum, eckit::mpi::sum()); return sum; } @@ -111,8 +111,19 @@ CASE("test_halo_exchange_01") { remote_indexv.assign({0, 1, 2, 3, 1, 0, 0, 2, 2, 3, 3, 1}); - - if (atlas::mpi::rank() == 0) { + int color = 0; + if( mpi::rank() >= mpi::size() - 2) { + color = 1; + } + auto& comm = mpi::comm().split(color,"split"); + if( color == 0 ) { + eckit::mpi::deleteComm("split"); + return; + } + util::Config mpi_comm("mpi_comm","split"); + int mpi_rank = comm.rank(); + + if (mpi_rank == 0) { // center followed by clockwise halo starting from top left lonlatv.assign({-45.0, 45.0, 45.0, 45.0, // center, first row -45.0, -45.0, 45.0, -45.0, // center, second row @@ -126,7 +137,7 @@ CASE("test_halo_exchange_01") { - } else if (atlas::mpi::rank() == 1) { + } else if (mpi_rank == 1) { // center followed by clockwise halo starting from top left lonlatv.assign({135.0, 45.0, 225.0, 45.0, 135.0, -45.0, 225.0, -45.0, @@ -146,7 +157,7 @@ CASE("test_halo_exchange_01") { fset.add(remote_index); fset.add(partition); - auto fs2 = functionspace::PointCloud(fset); + auto fs2 = functionspace::PointCloud(fset, mpi_comm); Field f1 = fs2.createField(option::name("f1") | option::levels(2)); Field f2 = fs2.createField(option::name("f2") | option::levels(2)); auto f1v = array::make_view(f1); @@ -158,7 +169,7 @@ CASE("test_halo_exchange_01") { for (idx_t l = 0; l < f2v.shape(1); ++l) { auto ghostv2 = array::make_view(fs2.ghost()); if (ghostv2(i) == 0) { - f1v(i, l) = (atlas::mpi::rank() +1) * 10.0 + i; + f1v(i, l) = (mpi_rank +1) * 10.0 + i; f2v(i, l) = f1v(i, l); } } @@ -173,8 +184,8 @@ CASE("test_halo_exchange_01") { double sum2 = innerproductwithhalo(f1, f2); - atlas::mpi::comm().allReduceInPlace(sum1, eckit::mpi::sum()); - atlas::mpi::comm().allReduceInPlace(sum2, eckit::mpi::sum()); + comm.allReduceInPlace(sum1, eckit::mpi::sum()); + comm.allReduceInPlace(sum2, eckit::mpi::sum()); EXPECT(std::abs(sum1 - sum2)/ std::abs(sum1) < tolerance); atlas::Log::info() << "adjoint test passed :: " << "sum1 " << sum1 << " sum2 " << sum2 << " normalised difference " @@ -189,6 +200,9 @@ CASE("test_halo_exchange_01") { } atlas::Log::info() << "values from halo followed by halo adjoint are as expected " << std::endl; + + eckit::mpi::deleteComm("split"); + } @@ -202,7 +216,19 @@ CASE("test_halo_exchange_02") { auto lonlatv = array::make_view(lonlat); auto ghostv = array::make_view(ghost); - if (atlas::mpi::rank() == 0) { + int color = 0; + if( mpi::rank() >= mpi::size() - 2) { + color = 1; + } + auto& comm = mpi::comm().split(color,"split"); + if( color == 0 ) { + eckit::mpi::deleteComm("split"); + return; + } + util::Config mpi_comm("mpi_comm","split"); + int mpi_rank = comm.rank(); + + if (mpi_rank == 0) { // center followed by clockwise halo starting from top left lonlatv.assign({-45.0, 45.0, 45.0, 45.0, // center, first row -45.0, -45.0, 45.0, -45.0, // center, second row @@ -214,7 +240,7 @@ CASE("test_halo_exchange_02") { ghostv.assign({ 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1}); - } else if (atlas::mpi::rank() == 1) { + } else if (mpi_rank == 1) { // center followed by clockwise halo starting from top left lonlatv.assign({135.0, 45.0, 225.0, 45.0, 135.0, -45.0, 225.0, -45.0, @@ -228,7 +254,7 @@ CASE("test_halo_exchange_02") { } - auto pcfs = functionspace::PointCloud(lonlat, ghost); + auto pcfs = functionspace::PointCloud(lonlat, ghost, mpi_comm); Field f1 = pcfs.createField(option::name("f1") | option::levels(2)); Field f2 = pcfs.createField(option::name("f2") | option::levels(2)); @@ -242,7 +268,7 @@ CASE("test_halo_exchange_02") { for (idx_t l = 0; l < f2v.shape(1); ++l) { auto ghostv2 = array::make_view(pcfs.ghost()); if (ghostv2(i) == 0) { - f1v(i, l) = (atlas::mpi::rank() +1) * 10.0 + i; + f1v(i, l) = (mpi_rank +1) * 10.0 + i; f2v(i, l) = f1v(i, l); } } @@ -257,8 +283,8 @@ CASE("test_halo_exchange_02") { double sum2 = innerproductwithhalo(f1, f2); - atlas::mpi::comm().allReduceInPlace(sum1, eckit::mpi::sum()); - atlas::mpi::comm().allReduceInPlace(sum2, eckit::mpi::sum()); + comm.allReduceInPlace(sum1, eckit::mpi::sum()); + comm.allReduceInPlace(sum2, eckit::mpi::sum()); EXPECT(std::abs(sum1 - sum2)/ std::abs(sum1) < tolerance); atlas::Log::info() << "adjoint test passed :: " << "sum1 " << sum1 << " sum2 " << sum2 << " normalised difference " @@ -274,6 +300,8 @@ CASE("test_halo_exchange_02") { atlas::Log::info() << "values from halo followed by halo adjoint are as expected " << std::endl; + eckit::mpi::deleteComm("split"); + }