Skip to content

Commit

Permalink
Use mpi_comm in functionspace::PointCloud
Browse files Browse the repository at this point in the history
  • Loading branch information
@wdeconinck
wdeconinck committed Sep 20, 2023
1 parent 9495fa8 commit 8d5b072
Show file tree
Hide file tree
Showing 3 changed files with 99 additions and 52 deletions.
75 changes: 45 additions & 30 deletions src/atlas/functionspace/PointCloud.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<PointXY>& points, const eckit::Configuration& config) {
mpi_comm_ = get_mpi_comm(config);
lonlat_ = Field("lonlat", array::make_datatype<double>(), array::make_shape(points.size(), 2));
auto lonlat = array::make_view<double, 2>(lonlat_);
for (idx_t j = 0, size = points.size(); j < size; ++j) {
Expand All @@ -57,6 +65,7 @@ PointCloud::PointCloud(const std::vector<PointXY>& points, const eckit::Configur

template <>
PointCloud::PointCloud(const std::vector<PointXYZ>& points, const eckit::Configuration& config) {
mpi_comm_ = get_mpi_comm(config);
lonlat_ = Field("lonlat", array::make_datatype<double>(), array::make_shape(points.size(), 2));
vertical_ = Field("vertical", array::make_datatype<double>(), array::make_shape(points.size()));
auto lonlat = array::make_view<double, 2>(lonlat_);
Expand All @@ -68,13 +77,17 @@ PointCloud::PointCloud(const std::vector<PointXYZ>& 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"];
}
Expand All @@ -93,13 +106,15 @@ 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"));
}
if( not partitioner.has("type") ) {
partitioner.set("type","equal_regions");
}
partitioner.set("mpi_comm",mpi_comm);
return grid::Partitioner(partitioner);
}

Expand All @@ -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<double>(), array::make_shape(size_halo, 2));
Expand All @@ -136,12 +153,12 @@ PointCloud::PointCloud(const Grid& grid, const grid::Partitioner& _partitioner,
auto ridx = array::make_indexview<idx_t,1>(remote_index_);
auto gidx = array::make_view<gidx_t,1>(global_index_);
array::make_view<int,1>(ghost_).assign(0);
array::make_view<int,1>(partition_).assign(part);
array::make_view<int,1>(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();
Expand All @@ -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);
}
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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<double, 2>(lonlat_);
// data structure containing a flag to identify the 'ghost points';
Expand Down Expand Up @@ -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<int, 1>(partition_).data(),
halo_exchange_->setup(mpi_comm_,
array::make_view<int, 1>(partition_).data(),
array::make_view<idx_t, 1>(remote_index_).data(),
REMOTE_IDX_BASE,
ghost_.size());

}

void PointCloud::adjointHaloExchange(const FieldSet& fieldset, bool on_device) const {
Expand Down Expand Up @@ -869,25 +884,25 @@ void PointCloud::adjointHaloExchange(const Field& field, bool) const {
PointCloud::PointCloud(const FunctionSpace& functionspace):
FunctionSpace(functionspace), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const Field& field):
FunctionSpace(new detail::PointCloud(field)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}
PointCloud::PointCloud(const Field& field, const eckit::Configuration& config):
FunctionSpace(new detail::PointCloud(field,config)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const Field& field1, const Field& field2):
FunctionSpace(new detail::PointCloud(field1, field2)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}
PointCloud::PointCloud(const Field& field1, const Field& field2, const eckit::Configuration& config):
FunctionSpace(new detail::PointCloud(field1, field2, config)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const FieldSet& fset):
FunctionSpace(new detail::PointCloud(fset)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}
PointCloud::PointCloud(const FieldSet& fset, const eckit::Configuration& config):
FunctionSpace(new detail::PointCloud(fset, config)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const std::vector<PointXY>& points):
FunctionSpace(new detail::PointCloud(points)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}
PointCloud::PointCloud(const std::vector<PointXY>& points, const eckit::Configuration& config):
FunctionSpace(new detail::PointCloud(points, config)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const std::vector<PointXYZ>& points):
FunctionSpace(new detail::PointCloud(points)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}
PointCloud::PointCloud(const std::vector<PointXYZ>& points, const eckit::Configuration& config):
FunctionSpace(new detail::PointCloud(points, config)), functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const std::initializer_list<std::initializer_list<double>>& points):
PointCloud::PointCloud(const std::initializer_list<std::initializer_list<double>>& points, const eckit::Configuration& config):
FunctionSpace((points.begin()->size() == 2
? new detail::PointCloud{std::vector<PointXY>(points.begin(), points.end())}
: new detail::PointCloud{std::vector<PointXYZ>(points.begin(), points.end())})),
? new detail::PointCloud{std::vector<PointXY>(points.begin(), points.end()), config}
: new detail::PointCloud{std::vector<PointXYZ>(points.begin(), points.end()), config})),
functionspace_(dynamic_cast<const detail::PointCloud*>(get())) {}

PointCloud::PointCloud(const Grid& grid, const eckit::Configuration& config):
Expand Down
18 changes: 11 additions & 7 deletions src/atlas/functionspace/PointCloud.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -166,7 +168,9 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
idx_t max_glb_idx_{0};
std::unique_ptr<parallel::HaloExchange> halo_exchange_;
idx_t levels_{0};
idx_t part_{0};
idx_t nb_partitions_{1};
std::string mpi_comm_;

void setupHaloExchange();

Expand All @@ -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<PointXY>&);
PointCloud(const std::vector<PointXYZ>&);
PointCloud(const std::initializer_list<std::initializer_list<double>>&);
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<PointXY>&, const eckit::Configuration& = util::NoConfig());
PointCloud(const std::vector<PointXYZ>&, const eckit::Configuration& = util::NoConfig());
PointCloud(const std::initializer_list<std::initializer_list<double>>&, const eckit::Configuration& = util::NoConfig());
PointCloud(const Grid&, const eckit::Configuration& = util::NoConfig());
PointCloud(const Grid&, const grid::Partitioner&, const eckit::Configuration& = util::NoConfig());

Expand Down
58 changes: 43 additions & 15 deletions src/tests/functionspace/test_pointcloud_haloexchange_2PE.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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<double>(option::name("f1") | option::levels(2));
Field f2 = fs2.createField<double>(option::name("f2") | option::levels(2));
auto f1v = array::make_view<double, 2>(f1);
Expand All @@ -158,7 +169,7 @@ CASE("test_halo_exchange_01") {
for (idx_t l = 0; l < f2v.shape(1); ++l) {
auto ghostv2 = array::make_view<int, 1>(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);
}
}
Expand All @@ -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 "
Expand All @@ -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");

}


Expand All @@ -202,7 +216,19 @@ CASE("test_halo_exchange_02") {
auto lonlatv = array::make_view<double, 2>(lonlat);
auto ghostv = array::make_view<int, 1>(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
Expand All @@ -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,
Expand All @@ -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<double>(option::name("f1") | option::levels(2));
Field f2 = pcfs.createField<double>(option::name("f2") | option::levels(2));
Expand All @@ -242,7 +268,7 @@ CASE("test_halo_exchange_02") {
for (idx_t l = 0; l < f2v.shape(1); ++l) {
auto ghostv2 = array::make_view<int, 1>(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);
}
}
Expand All @@ -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 "
Expand All @@ -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");

}


Expand Down

0 comments on commit 8d5b072

Please sign in to comment.