From cdb0447868f57aeb90360418f5c74a684dee069a Mon Sep 17 00:00:00 2001 From: Marcel Koch Date: Tue, 2 Apr 2024 07:25:56 +0000 Subject: [PATCH] review updates: - renaming - separate tests Co-authored-by: Pratik Nayak --- .../distributed/index_map_kernels.hpp.inc | 7 +- omp/distributed/index_map_kernels.cpp | 12 +- test/distributed/index_map_kernels.cpp | 139 +++++++++++------- 3 files changed, 92 insertions(+), 66 deletions(-) diff --git a/common/cuda_hip/distributed/index_map_kernels.hpp.inc b/common/cuda_hip/distributed/index_map_kernels.hpp.inc index f60fb2f1963..9d312cc43aa 100644 --- a/common/cuda_hip/distributed/index_map_kernels.hpp.inc +++ b/common/cuda_hip/distributed/index_map_kernels.hpp.inc @@ -55,6 +55,7 @@ array compute_range_ids( return range_ids; } + template void build_mapping( std::shared_ptr exec, @@ -161,7 +162,7 @@ void map_to_local( std::shared_ptr exec, const experimental::distributed::Partition* partition, - const array& remote_targed_ids, + const array& remote_target_ids, device_segmented_array remote_global_idxs, experimental::distributed::comm_index_type rank, const array& global_ids, @@ -197,8 +198,8 @@ void map_to_local( : invalid_index(); }; - auto remote_target_ids_ptr = remote_targed_ids.get_const_data(); - auto num_target_ids = remote_targed_ids.get_size(); + auto remote_target_ids_ptr = remote_target_ids.get_const_data(); + auto num_target_ids = remote_target_ids.get_size(); auto remote_global_idxs_ptr = remote_global_idxs.flat_begin; auto offsets_ptr = remote_global_idxs.offsets_begin; auto map_non_local = diff --git a/omp/distributed/index_map_kernels.cpp b/omp/distributed/index_map_kernels.cpp index 70efe56b878..27c98593fab 100644 --- a/omp/distributed/index_map_kernels.cpp +++ b/omp/distributed/index_map_kernels.cpp @@ -128,7 +128,7 @@ void map_to_local( std::shared_ptr exec, const experimental::distributed::Partition* partition, - const array& remote_targed_ids, + const array& remote_target_ids, device_segmented_array remote_global_idxs, experimental::distributed::comm_index_type rank, const array& global_ids, @@ -162,13 +162,13 @@ void map_to_local( // the global indexing. So find the part-id that corresponds // to the global index first auto set_id = - std::distance(remote_targed_ids.get_const_data(), - std::lower_bound(remote_targed_ids.get_const_data(), - remote_targed_ids.get_const_data() + - remote_targed_ids.get_size(), + std::distance(remote_target_ids.get_const_data(), + std::lower_bound(remote_target_ids.get_const_data(), + remote_target_ids.get_const_data() + + remote_target_ids.get_size(), part_id)); - if (set_id == remote_targed_ids.get_size()) { + if (set_id == remote_target_ids.get_size()) { return std::make_pair(invalid_index(), range_id); } diff --git a/test/distributed/index_map_kernels.cpp b/test/distributed/index_map_kernels.cpp index 8ca93c2bcdf..458ca594a56 100644 --- a/test/distributed/index_map_kernels.cpp +++ b/test/distributed/index_map_kernels.cpp @@ -28,6 +28,86 @@ using comm_index_type = gko::experimental::distributed::comm_index_type; +template +gko::array generate_connection_idxs( + const std::shared_ptr& exec, comm_index_type rank, + std::shared_ptr> + partition, + std::default_random_engine engine, gko::size_type num_connections) +{ + auto ref = exec->get_master(); + auto num_parts = partition->get_num_parts(); + auto local_size = + static_cast(partition->get_part_size(rank)); + // create vector with [0, ..., num_parts) excluding excluded_pid + std::vector part_ids(num_parts - 1); + std::iota(part_ids.begin(), part_ids.end(), rank + 1); + std::transform(part_ids.begin(), part_ids.end(), part_ids.begin(), + [&](const auto pid) { return pid % num_parts; }); + // get random connections + std::shuffle(part_ids.begin(), part_ids.end(), engine); + std::vector connected_ids( + part_ids.begin(), part_ids.begin() + num_connections); + // create global index space of connections + std::vector connections_index_space; + for (auto pid : connected_ids) { + for (GlobalIndexType i = 0; i < local_size; ++i) { + connections_index_space.push_back( + i + static_cast(pid * local_size)); + } + } + // generate query from connection_index_space + std::uniform_int_distribution<> dist(0, connections_index_space.size() - 1); + gko::array connection_idxs{ref, 11}; + std::generate_n(connection_idxs.get_data(), connection_idxs.get_size(), + [&] { return connections_index_space[dist(engine)]; }); + return {exec, std::move(connection_idxs)}; +} + + +class IndexMapBuildMapping : public CommonTestFixture {}; + + +TEST_F(IndexMapBuildMapping, BuildMappingSameAsRef) +{ + using local_index_type = gko::int32; + using global_index_type = gko::int64; + using part_type = + gko::experimental::distributed::Partition; + std::default_random_engine engine; + comm_index_type num_parts = 13; + global_index_type local_size = 41; + comm_index_type this_rank = 5; + std::shared_ptr part = part_type::build_from_global_size_uniform( + ref, num_parts, num_parts * local_size); + std::shared_ptr dpart = gko::clone(exec, part); + auto query = generate_connection_idxs(ref, this_rank, part, engine, 11); + auto dquery = gko::array(exec, query); + gko::array target_ids{ref}; + gko::array remote_local_idxs{ref}; + gko::array remote_global_idxs{ref}; + gko::array remote_sizes{ref}; + gko::array dtarget_ids{exec}; + gko::array dremote_local_idxs{exec}; + gko::array dremote_global_idxs{exec}; + gko::array dremote_sizes{exec}; + + gko::kernels::reference::index_map::build_mapping( + ref, part.get(), query, target_ids, remote_local_idxs, + remote_global_idxs, remote_sizes); + gko::kernels::EXEC_NAMESPACE::index_map::build_mapping( + exec, dpart.get(), dquery, dtarget_ids, dremote_local_idxs, + dremote_global_idxs, dremote_sizes); + + GKO_ASSERT_ARRAY_EQ(remote_sizes, dremote_sizes); + GKO_ASSERT_ARRAY_EQ(target_ids, dtarget_ids); + GKO_ASSERT_ARRAY_EQ(remote_local_idxs, dremote_local_idxs); + GKO_ASSERT_ARRAY_EQ(remote_global_idxs, dremote_global_idxs); +} + + class IndexMap : public CommonTestFixture { protected: using local_index_type = gko::int32; @@ -41,7 +121,8 @@ class IndexMap : public CommonTestFixture { IndexMap() { - auto connections = generate_connection_idxs(ref, this_rank, 11); + auto connections = + generate_connection_idxs(ref, this_rank, part, engine, 11); auto dconnections = gko::array(exec, connections); auto flat_remote_local_idxs = gko::array(ref); @@ -73,36 +154,6 @@ class IndexMap : public CommonTestFixture { std::move(dflat_remote_global_idxs), dremote_sizes); } - gko::array generate_connection_idxs( - std::shared_ptr exec, comm_index_type excluded_pid, - gko::size_type num_connections) - { - // create vector with [0, ..., num_parts) excluding excluded_pid - std::vector part_ids(num_parts - 1); - std::iota(part_ids.begin(), part_ids.end(), excluded_pid + 1); - std::transform(part_ids.begin(), part_ids.end(), part_ids.begin(), - [&](const auto pid) { return pid % num_parts; }); - // get random connections - std::shuffle(part_ids.begin(), part_ids.end(), engine); - std::vector connected_ids( - part_ids.begin(), part_ids.begin() + num_connections); - // create global index space of connections - std::vector connections_index_space; - for (auto pid : connected_ids) { - for (global_index_type i = 0; i < local_size; ++i) { - connections_index_space.push_back( - i + static_cast(pid * local_size)); - } - } - // generate query from connection_index_space - std::uniform_int_distribution<> dist( - 0, connections_index_space.size() - 1); - gko::array connection_idxs{ref, 11}; - std::generate_n(connection_idxs.get_data(), connection_idxs.get_size(), - [&] { return connections_index_space[dist(engine)]; }); - return {std::move(exec), std::move(connection_idxs)}; - } - gko::array generate_query( std::shared_ptr exec, const gko::array& connection_idxs, @@ -181,32 +232,6 @@ class IndexMap : public CommonTestFixture { std::default_random_engine engine; }; -TEST_F(IndexMap, BuildMappingSameAsRef) -{ - auto query = generate_connection_idxs(ref, this_rank, 11); - auto dquery = gko::array(exec, query); - gko::array target_ids{ref}; - gko::array remote_local_idxs{ref}; - gko::array remote_global_idxs{ref}; - gko::array remote_sizes{ref}; - gko::array dtarget_ids{exec}; - gko::array dremote_local_idxs{exec}; - gko::array dremote_global_idxs{exec}; - gko::array dremote_sizes{exec}; - - gko::kernels::reference::index_map::build_mapping( - ref, part.get(), query, target_ids, remote_local_idxs, - remote_global_idxs, remote_sizes); - gko::kernels::EXEC_NAMESPACE::index_map::build_mapping( - exec, dpart.get(), dquery, dtarget_ids, dremote_local_idxs, - dremote_global_idxs, dremote_sizes); - - GKO_ASSERT_ARRAY_EQ(remote_sizes, dremote_sizes); - GKO_ASSERT_ARRAY_EQ(target_ids, dtarget_ids); - GKO_ASSERT_ARRAY_EQ(remote_local_idxs, dremote_local_idxs); - GKO_ASSERT_ARRAY_EQ(remote_global_idxs, dremote_global_idxs); -} - TEST_F(IndexMap, GetLocalWithLocalIndexSpaceSameAsRef) {