diff --git a/common/unified/distributed/index_map_kernels.cpp b/common/unified/distributed/index_map_kernels.cpp index 0c53342eae8..84424976778 100644 --- a/common/unified/distributed/index_map_kernels.cpp +++ b/common/unified/distributed/index_map_kernels.cpp @@ -30,7 +30,7 @@ GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE( template -void get_local( +void map_to_local( std::shared_ptr exec, const experimental::distributed::Partition* partition, @@ -42,7 +42,7 @@ void get_local( array& local_ids) GKO_NOT_IMPLEMENTED; GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE( - GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY); + GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL); } // namespace index_map diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp index 4342e263583..2ee758eacea 100644 --- a/core/device_hooks/common_kernels.inc.cpp +++ b/core/device_hooks/common_kernels.inc.cpp @@ -256,7 +256,7 @@ namespace index_map { GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_INDEX_MAP_BUILD_MAPPING); -GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY); +GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL); } // namespace index_map diff --git a/core/distributed/index_map.cpp b/core/distributed/index_map.cpp index 2ac53ba26bc..5d2a1aebe18 100644 --- a/core/distributed/index_map.cpp +++ b/core/distributed/index_map.cpp @@ -13,7 +13,8 @@ namespace index_map_kernels { GKO_REGISTER_OPERATION(build_mapping, index_map::build_mapping); -GKO_REGISTER_OPERATION(get_local, index_map::get_local); +GKO_REGISTER_OPERATION(map_to_local, index_map::map_to_local); + } // namespace index_map_kernels @@ -76,14 +77,14 @@ size_type index_map::get_global_size() const template -array index_map::get_local( - const array& global_ids, index_space is) const +array index_map::map_to_local( + const array& global_ids, index_space index_space_v) const { array local_ids(exec_); - exec_->run(index_map_kernels::make_get_local( + exec_->run(index_map_kernels::make_map_to_local( partition_.get(), remote_target_ids_, to_device(remote_global_idxs_), - rank_, global_ids, is, local_ids)); + rank_, global_ids, index_space_v, local_ids)); return local_ids; } @@ -91,10 +92,11 @@ array index_map::get_local( template index_map::index_map( - std::shared_ptr exec, std::shared_ptr part, - comm_index_type rank, const array& recv_connections) + std::shared_ptr exec, + std::shared_ptr partition, comm_index_type rank, + const array& recv_connections) : exec_(std::move(exec)), - partition_(std::move(part)), + partition_(std::move(partition)), rank_(rank), remote_target_ids_(exec_), remote_local_idxs_(exec_), diff --git a/core/distributed/index_map_kernels.hpp b/core/distributed/index_map_kernels.hpp index d9d0bd55a57..c808a4b8d19 100644 --- a/core/distributed/index_map_kernels.hpp +++ b/core/distributed/index_map_kernels.hpp @@ -2,13 +2,15 @@ // // SPDX-License-Identifier: BSD-3-Clause -#ifndef INDEX_MAP_KERNELS_HPP -#define INDEX_MAP_KERNELS_HPP +#ifndef GKO_CORE_DISTRIBUTED_INDEX_MAP_KERNELS_HPP_ +#define GKO_CORE_DISTRIBUTED_INDEX_MAP_KERNELS_HPP_ + + +#include #include #include -#include #include @@ -20,6 +22,23 @@ namespace gko { namespace kernels { +/** + * This kernel creates an index map from a partition and global remote indices. + * + * The index map is defined by the output parameters remote_local_idxs, + * remote_global_idxs, and remote_sizes. After this functions: + * + * - remote_global_idxs contains the unique indices from recv_connections, + * sorted first by the owning part (as defined in the partition) and then by + * the global index + * - remote_local_idxs contains the indices of remote_global_idxs but mapped + * into the local index spaces of their owning parts + * - remote_sizes contains the number of remote indices (either in + * remote_global_idxs, or remote_local_indices) per owning part. + * + * The sizes array is used to create segmented arrays for both output index + * arrays. + */ #define GKO_DECLARE_INDEX_MAP_BUILD_MAPPING(_ltype, _gtype) \ void build_mapping( \ std::shared_ptr exec, \ @@ -30,12 +49,26 @@ namespace kernels { array& remote_sizes) -#define GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY(_ltype, _gtype) \ - void get_local( \ +/** + * This kernel maps global indices to local indices. + * + * The global indices in remote_global_idxs are mapped into the local index + * space defined by is. The resulting indices are stored in local_ids. + * The index map is defined by the input parameters: + * + * - partition: the global partition + * - remote_target_ids: the owning part ids of each segment of + * remote_global_idxs + * - remote_global_idxs: the remote global indices, segmented by the owning part + * ids + * - rank: the part id of this process + */ +#define GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL(_ltype, _gtype) \ + void map_to_local( \ std::shared_ptr exec, \ const experimental::distributed::Partition<_ltype, _gtype>* partition, \ const array& \ - remote_targed_ids, \ + remote_target_ids, \ device_segmented_array remote_global_idxs, \ experimental::distributed::comm_index_type rank, \ const array<_gtype>& global_ids, \ @@ -46,8 +79,7 @@ namespace kernels { template \ GKO_DECLARE_INDEX_MAP_BUILD_MAPPING(LocalIndexType, GlobalIndexType); \ template \ - GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY(LocalIndexType, \ - GlobalIndexType) + GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL(LocalIndexType, GlobalIndexType) GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(index_map, @@ -60,4 +92,4 @@ GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(index_map, } // namespace kernels } // namespace gko -#endif // INDEX_MAP_KERNELS_HPP +#endif // GKO_CORE_DISTRIBUTED_INDEX_MAP_KERNELS_HPP_ diff --git a/core/test/distributed/index_map.cpp b/core/test/distributed/index_map.cpp index d24710a2126..521a2e2d094 100644 --- a/core/test/distributed/index_map.cpp +++ b/core/test/distributed/index_map.cpp @@ -89,6 +89,7 @@ TYPED_TEST(IndexMap, CanMoveConstruct) { using map_type = typename TestFixture::map_type; using global_index_type = typename TestFixture::global_index_type; + using local_index_type = typename TestFixture::local_index_type; gko::array connections(this->exec, {4, 3, 3, 4, 2}); auto imap = map_type(this->exec, this->part, 0, connections); auto copy = imap; @@ -112,6 +113,13 @@ TYPED_TEST(IndexMap, CanMoveConstruct) imap_remote_local_it); ASSERT_EQ(move.get_remote_global_idxs().get_const_flat_data(), imap_remote_global_it); + ASSERT_EQ(imap.get_global_size(), 0); + ASSERT_EQ(imap.get_local_size(), 0); + GKO_ASSERT_ARRAY_EQ(imap.get_remote_target_ids(), {}); + assert_collection_eq(imap.get_remote_global_idxs(), + gko::segmented_array{this->exec}); + assert_collection_eq(imap.get_remote_local_idxs(), + gko::segmented_array{this->exec}); } @@ -138,6 +146,7 @@ TYPED_TEST(IndexMap, CanMoveAssign) { using map_type = typename TestFixture::map_type; using global_index_type = typename TestFixture::global_index_type; + using local_index_type = typename TestFixture::local_index_type; gko::array connections(this->exec, {4, 3, 3, 4, 2}); auto imap = map_type(this->exec, this->part, 0, connections); auto copy = imap; @@ -162,4 +171,11 @@ TYPED_TEST(IndexMap, CanMoveAssign) imap_remote_local_it); ASSERT_EQ(move.get_remote_global_idxs().get_const_flat_data(), imap_remote_global_it); + ASSERT_EQ(imap.get_global_size(), 0); + ASSERT_EQ(imap.get_local_size(), 0); + GKO_ASSERT_ARRAY_EQ(imap.get_remote_target_ids(), {}); + assert_collection_eq(imap.get_remote_global_idxs(), + gko::segmented_array{this->exec}); + assert_collection_eq(imap.get_remote_local_idxs(), + gko::segmented_array{this->exec}); } diff --git a/include/ginkgo/core/distributed/index_map.hpp b/include/ginkgo/core/distributed/index_map.hpp index 01e53b37092..35d072ccc7e 100644 --- a/include/ginkgo/core/distributed/index_map.hpp +++ b/include/ginkgo/core/distributed/index_map.hpp @@ -17,6 +17,8 @@ namespace distributed { /** * \brief Index space classification for the locally stored indices. + * + * The definitions of the enum values is clarified in @ref index_map. */ enum class index_space { local, //!< indices that are locally owned @@ -42,17 +44,20 @@ enum class index_space { * (index_space::non_local). The locally owned indices are defined as * $[0, \dots, |I_k|)$, and the non-locally owned as $[0, \dots, |R_k|)$. * With these index sets, the following mappings are defined: + * * - $c_k : \hat{I}_k \mapsto [0, \dots, |\hat{I}_k|)$ which maps global indices - * into the combined/full local index space, + * into the combined/full local index space (denoted as + * index_space::combined), * - $l_k: I_k \mapsto [0, \dots, |I_k|)$ which maps global indices into the - * locally owned index space, + * locally owned index space (denoted as index_space::local), * - $r_k: R_k \mapsto [0, \dots, |R_k|)$ which maps global indices into the - * non-locally owned index space. + * non-locally owned index space (denoted as index_space::non_local). + * * The required map can be selected by passing the appropriate type of an * index_space. * * The index map for $I_k$ has no knowledge about any other index maps for - * $I_l, l \neq k$. In particular, any global index passed to the `get_local` + * $I_l, l \neq k$. In particular, any global index passed to the `map_to_local` * map that is not part of the specified index space, will be mapped to an * invalid_index. * @@ -61,18 +66,19 @@ enum class index_space { */ template struct index_map { - using part_type = Partition; + using partition_type = Partition; /** * \brief Maps global indices to local indices * - * \param global_ids the global indices to map - * \param is the index space in which the returned local indices are defined + * \param global_ids the global indices to map + * \param index_space_v the index space in which the returned local indices + * are defined * - * \return the mapped local indices. Any global index that is not in the - * specified index space is mapped to invalid_index. + * \return the mapped local indices. Any global index that is not in the + * specified index space is mapped to invalid_index. */ - array get_local( + array map_to_local( const array& global_ids, index_space index_space_v = index_space::combined) const; @@ -98,13 +104,14 @@ struct index_map { * filtered out. * * \param exec the executor - * \param part the partition of the global index set + * \param partition the partition of the global index set * \param rank the id of the global index space subset * \param recv_connections the global indices that are not owned by this * rank, but accessed by it */ index_map(std::shared_ptr exec, - std::shared_ptr part, comm_index_type rank, + std::shared_ptr partition, + comm_index_type rank, const array& recv_connections); /** @@ -123,7 +130,8 @@ struct index_map { * \brief get the index set $R_k$, but mapped to their respective local * index space. * - * The indices are ordered by their owning rank and global index. + * The indices are grouped by their owning rank and sorted according to + * their global index within each group. * * The set $R_k = \hat{I}_k \setminus I_k$ can also be written as the union * of the intersection of $\hat{I}_k$ with other disjoint sets @@ -157,7 +165,7 @@ struct index_map { private: std::shared_ptr exec_; - std::shared_ptr partition_; + std::shared_ptr partition_; comm_index_type rank_; array remote_target_ids_; diff --git a/reference/distributed/index_map_kernels.cpp b/reference/distributed/index_map_kernels.cpp index 7e9e9a79313..5f13581eee0 100644 --- a/reference/distributed/index_map_kernels.cpp +++ b/reference/distributed/index_map_kernels.cpp @@ -8,6 +8,7 @@ #include +#include "core/base/allocator.hpp" #include "core/base/segmented_array.hpp" #include "reference/distributed/partition_helpers.hpp" @@ -29,35 +30,12 @@ void build_mapping( array& remote_global_idxs, array& remote_sizes) { using experimental::distributed::comm_index_type; - using partition_type = - experimental::distributed::Partition; auto part_ids = part->get_part_ids(); - std::vector unique_indices(recv_connections.get_size()); + vector unique_indices(recv_connections.get_size(), {exec}); std::copy_n(recv_connections.get_const_data(), recv_connections.get_size(), unique_indices.begin()); - auto find_range = [](GlobalIndexType idx, const partition_type* partition, - size_type hint) { - auto range_bounds = partition->get_range_bounds(); - auto num_ranges = partition->get_num_ranges(); - if (range_bounds[hint] <= idx && idx < range_bounds[hint + 1]) { - return hint; - } else { - auto it = std::upper_bound(range_bounds + 1, - range_bounds + num_ranges + 1, idx); - return static_cast(std::distance(range_bounds + 1, it)); - } - }; - - auto map_to_local = [](GlobalIndexType idx, const partition_type* partition, - size_type range_id) { - auto range_bounds = partition->get_range_bounds(); - auto range_starting_indices = partition->get_range_starting_indices(); - return static_cast(idx - range_bounds[range_id]) + - range_starting_indices[range_id]; - }; - auto find_part = [&](GlobalIndexType idx) { auto range_id = find_range(idx, part, 0); return part_ids[range_id]; @@ -97,12 +75,12 @@ void build_mapping( }); // get part-ids - std::vector full_part_ids(unique_size); + vector full_part_ids(unique_size, exec); std::transform(unique_indices.begin(), unique_indices_end, full_part_ids.begin(), [&](const auto idx) { return find_part(idx); }); - std::vector unique_part_ids(full_part_ids); + vector unique_part_ids(full_part_ids, exec); auto unique_part_ids_end = std::unique(unique_part_ids.begin(), unique_part_ids.end()); @@ -113,7 +91,7 @@ void build_mapping( remote_part_ids.get_data()); // get recv size per part - std::vector full_remote_sizes(part->get_num_parts()); + vector full_remote_sizes(part->get_num_parts(), exec); for (size_type i = 0; i < full_part_ids.size(); ++i) { full_remote_sizes[full_part_ids[i]]++; } @@ -127,11 +105,11 @@ GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE( template -void get_local( +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,14 +140,14 @@ void get_local( // global index. As a result, the array is not sorted wrt. // 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(), - part_id)); - - if (set_id == remote_targed_ids.get_size()) { + auto set_id = + 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_target_ids.get_size()) { return invalid_index(); } @@ -220,7 +198,8 @@ void get_local( } GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE( - GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY); + GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL); + } // namespace index_map } // namespace reference diff --git a/reference/test/distributed/index_map_kernels.cpp b/reference/test/distributed/index_map_kernels.cpp index 7bff4fcc5b3..972db10654f 100644 --- a/reference/test/distributed/index_map_kernels.cpp +++ b/reference/test/distributed/index_map_kernels.cpp @@ -20,8 +20,6 @@ #include "core/distributed/index_map_kernels.hpp" #include "core/test/utils.hpp" -namespace { - using comm_index_type = gko::experimental::distributed::comm_index_type; @@ -98,7 +96,7 @@ TEST_F(IndexMap, CanGetLocalWithNonLocalIS) {ref, {0, 1, 4}}, {ref, {2, 1}}); gko::array remote_target_ids(ref, {0, 2}); - gko::kernels::reference::index_map::get_local( + gko::kernels::reference::index_map::map_to_local( ref, part.get(), remote_target_ids, to_device_const(remote_global_idxs), 1, global_ids, gko::experimental::distributed::index_space::non_local, local_ids); @@ -117,7 +115,7 @@ TEST_F(IndexMap, CanGetLocalWithNonLocalISWithInvalid) {ref, {0, 1, 4}}, {ref, {2, 1}}); gko::array remote_target_ids(ref, {0, 2}); - gko::kernels::reference::index_map::get_local( + gko::kernels::reference::index_map::map_to_local( ref, part.get(), remote_target_ids, to_device_const(remote_global_idxs), 1, global_ids, gko::experimental::distributed::index_space::non_local, local_ids); @@ -136,7 +134,7 @@ TEST_F(IndexMap, CanGetLocalWithLocalIS) {ref, {0, 1, 4}}, {ref, {2, 1}}); gko::array remote_target_ids(ref, {0, 2}); - gko::kernels::reference::index_map::get_local( + gko::kernels::reference::index_map::map_to_local( ref, part.get(), remote_target_ids, to_device_const(remote_global_idxs), 1, global_ids, gko::experimental::distributed::index_space::local, local_ids); @@ -155,7 +153,7 @@ TEST_F(IndexMap, CanGetLocalWithLocalISWithInvalid) {ref, {0, 1, 4}}, {ref, {2, 1}}); gko::array remote_target_ids(ref, {0, 2}); - gko::kernels::reference::index_map::get_local( + gko::kernels::reference::index_map::map_to_local( ref, part.get(), remote_target_ids, to_device_const(remote_global_idxs), 1, global_ids, gko::experimental::distributed::index_space::local, local_ids); @@ -174,7 +172,7 @@ TEST_F(IndexMap, CanGetLocalWithCombinedIS) {ref, {0, 1, 4}}, {ref, {2, 1}}); gko::array remote_target_ids(ref, {0, 2}); - gko::kernels::reference::index_map::get_local( + gko::kernels::reference::index_map::map_to_local( ref, part.get(), remote_target_ids, to_device_const(remote_global_idxs), 1, global_ids, gko::experimental::distributed::index_space::combined, local_ids); @@ -193,7 +191,7 @@ TEST_F(IndexMap, CanGetLocalWithCombinedISWithInvalid) {ref, {0, 1, 4}}, {ref, {2, 1}}); gko::array remote_target_ids(ref, {0, 2}); - gko::kernels::reference::index_map::get_local( + gko::kernels::reference::index_map::map_to_local( ref, part.get(), remote_target_ids, to_device_const(remote_global_idxs), 1, global_ids, gko::experimental::distributed::index_space::combined, local_ids); @@ -201,6 +199,3 @@ TEST_F(IndexMap, CanGetLocalWithCombinedISWithInvalid) gko::array expected(ref, {2, 3, 0, 1, 2, 4, -1, 1}); GKO_ASSERT_ARRAY_EQ(local_ids, expected); } - - -} // namespace