From c1ecb3cb12a0cca570b026efbcd8323efe018122 Mon Sep 17 00:00:00 2001 From: Marcel Koch Date: Wed, 14 Feb 2024 12:38:58 +0000 Subject: [PATCH] adds distributed index map tests --- core/test/CMakeLists.txt | 1 + core/test/distributed/CMakeLists.txt | 1 + core/test/distributed/index_map.cpp | 172 +++++++++ reference/test/distributed/CMakeLists.txt | 1 + .../test/distributed/index_map_kernels.cpp | 169 +++++++++ test/distributed/CMakeLists.txt | 1 + test/distributed/index_map_kernels.cpp | 325 ++++++++++++++++++ 7 files changed, 670 insertions(+) create mode 100644 core/test/distributed/CMakeLists.txt create mode 100644 core/test/distributed/index_map.cpp create mode 100644 reference/test/distributed/index_map_kernels.cpp create mode 100644 test/distributed/index_map_kernels.cpp diff --git a/core/test/CMakeLists.txt b/core/test/CMakeLists.txt index 69f7ddd749e..c1622fc584a 100644 --- a/core/test/CMakeLists.txt +++ b/core/test/CMakeLists.txt @@ -4,6 +4,7 @@ add_subdirectory(gtest) add_subdirectory(accessor) add_subdirectory(base) add_subdirectory(components) +add_subdirectory(distributed) if(GINKGO_BUILD_MPI) add_subdirectory(mpi) endif() diff --git a/core/test/distributed/CMakeLists.txt b/core/test/distributed/CMakeLists.txt new file mode 100644 index 00000000000..dcaac0a51c5 --- /dev/null +++ b/core/test/distributed/CMakeLists.txt @@ -0,0 +1 @@ +ginkgo_create_test(index_map) diff --git a/core/test/distributed/index_map.cpp b/core/test/distributed/index_map.cpp new file mode 100644 index 00000000000..d1000180aff --- /dev/null +++ b/core/test/distributed/index_map.cpp @@ -0,0 +1,172 @@ +// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include + + +#include + + +#include + + +#include "core/test/utils.hpp" + + +using gko::experimental::distributed::comm_index_type; + + +template +class IndexMap : public ::testing::Test { +public: + using local_index_type = + typename std::tuple_element<0, decltype(LocalGlobalIndexType())>::type; + using global_index_type = + typename std::tuple_element<1, decltype(LocalGlobalIndexType())>::type; + using part_type = + gko::experimental::distributed::Partition; + using map_type = + gko::experimental::distributed::index_map; + + std::shared_ptr exec = + gko::ReferenceExecutor::create(); + std::shared_ptr part = + part_type::build_from_global_size_uniform(exec, 3, 6); +}; + +TYPED_TEST_SUITE(IndexMap, gko::test::LocalGlobalIndexTypes, + PairTypenameNameGenerator); + + +template +void assert_default_constructed_array(const gko::array& a) +{ + ASSERT_EQ(a.get_executor(), nullptr); + ASSERT_EQ(a.get_size(), 0); +} + +template +void assert_collection_eq(const gko::collection::array& a, + const gko::collection::array& b) +{ + ASSERT_EQ(a.size(), b.size()); + for (int i = 0; i < a.size(); ++i) { + GKO_ASSERT_ARRAY_EQ(a[i], b[i]); + } +} + + +TYPED_TEST(IndexMap, CanDefaultConstruct) +{ + using map_type = typename TestFixture::map_type; + + auto imap = map_type(); + + ASSERT_EQ(imap.get_local_size(), 0); + ASSERT_EQ(imap.get_non_local_size(), 0); + assert_default_constructed_array(imap.get_remote_target_ids()); + ASSERT_EQ(imap.get_remote_global_idxs().size(), 0); + assert_default_constructed_array(imap.get_remote_global_idxs().get_flat()); + ASSERT_EQ(imap.get_remote_local_idxs().size(), 0); + assert_default_constructed_array(imap.get_remote_local_idxs().get_flat()); +} + + +TYPED_TEST(IndexMap, CanCopyConstruct) +{ + using map_type = typename TestFixture::map_type; + using global_index_type = typename TestFixture::global_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; + + GKO_ASSERT_ARRAY_EQ(copy.get_remote_target_ids(), + imap.get_remote_target_ids()); + assert_collection_eq(copy.get_remote_local_idxs(), + imap.get_remote_local_idxs()); + assert_collection_eq(copy.get_remote_global_idxs(), + imap.get_remote_global_idxs()); +} + +TYPED_TEST(IndexMap, CanMoveConstruct) +{ + using map_type = typename TestFixture::map_type; + using global_index_type = typename TestFixture::global_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; + auto imap_remote_global_it = + imap.get_remote_global_idxs().get_flat().get_const_data(); + auto imap_remote_local_it = + imap.get_remote_local_idxs().get_flat().get_const_data(); + auto imap_remote_target_it = imap.get_remote_target_ids().get_const_data(); + + auto move = std::move(imap); + + GKO_ASSERT_ARRAY_EQ(move.get_remote_target_ids(), + copy.get_remote_target_ids()); + assert_collection_eq(move.get_remote_local_idxs(), + copy.get_remote_local_idxs()); + assert_collection_eq(move.get_remote_global_idxs(), + copy.get_remote_global_idxs()); + ASSERT_EQ(move.get_remote_target_ids().get_const_data(), + imap_remote_target_it); + ASSERT_EQ(move.get_remote_local_idxs().get_flat().get_const_data(), + imap_remote_local_it); + ASSERT_EQ(move.get_remote_global_idxs().get_flat().get_const_data(), + imap_remote_global_it); +} + + +TYPED_TEST(IndexMap, CanCopyAssign) +{ + using map_type = typename TestFixture::map_type; + using global_index_type = typename TestFixture::global_index_type; + gko::array connections(this->exec, {4, 3, 3, 4, 2}); + auto imap = map_type(this->exec, this->part, 0, connections); + auto copy = map_type(); + + copy = imap; + + GKO_ASSERT_ARRAY_EQ(copy.get_remote_target_ids(), + imap.get_remote_target_ids()); + assert_collection_eq(copy.get_remote_local_idxs(), + imap.get_remote_local_idxs()); + assert_collection_eq(copy.get_remote_global_idxs(), + imap.get_remote_global_idxs()); +} + + +TYPED_TEST(IndexMap, CanMoveAssign) +{ + using map_type = typename TestFixture::map_type; + using global_index_type = typename TestFixture::global_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; + auto imap_remote_global_it = + imap.get_remote_global_idxs().get_flat().get_const_data(); + auto imap_remote_local_it = + imap.get_remote_local_idxs().get_flat().get_const_data(); + auto imap_remote_target_it = imap.get_remote_target_ids().get_const_data(); + auto move = map_type(); + + move = std::move(imap); + + GKO_ASSERT_ARRAY_EQ(move.get_remote_target_ids(), + copy.get_remote_target_ids()); + assert_collection_eq(move.get_remote_local_idxs(), + copy.get_remote_local_idxs()); + assert_collection_eq(move.get_remote_global_idxs(), + copy.get_remote_global_idxs()); + ASSERT_EQ(move.get_remote_target_ids().get_const_data(), + imap_remote_target_it); + ASSERT_EQ(move.get_remote_local_idxs().get_flat().get_const_data(), + imap_remote_local_it); + ASSERT_EQ(move.get_remote_global_idxs().get_flat().get_const_data(), + imap_remote_global_it); +} diff --git a/reference/test/distributed/CMakeLists.txt b/reference/test/distributed/CMakeLists.txt index 42ad2d7e1a2..c4619369d6d 100644 --- a/reference/test/distributed/CMakeLists.txt +++ b/reference/test/distributed/CMakeLists.txt @@ -1,3 +1,4 @@ +ginkgo_create_test(index_map_kernels) ginkgo_create_test(matrix_kernels) ginkgo_create_test(partition_helpers_kernels) ginkgo_create_test(partition_kernels) diff --git a/reference/test/distributed/index_map_kernels.cpp b/reference/test/distributed/index_map_kernels.cpp new file mode 100644 index 00000000000..c22389692ce --- /dev/null +++ b/reference/test/distributed/index_map_kernels.cpp @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include + + +#include +#include +#include + + +#include +#include + + +#include + + +#include "core/distributed/index_map_kernels.hpp" +#include "core/test/utils.hpp" + +namespace { + + +using comm_index_type = gko::experimental::distributed::comm_index_type; + + +class IndexMap : public ::testing::Test { +protected: + using local_index_type = gko::int32; + using global_index_type = gko::int64; + using part_type = + gko::experimental::distributed::Partition; + using map_type = + gko::experimental::distributed::index_map; + + IndexMap() : ref(gko::ReferenceExecutor::create()) {} + + std::shared_ptr ref; + std::shared_ptr part = + part_type::build_from_mapping(ref, {ref, {0, 0, 1, 1, 2, 2}}, 3); +}; + + +TEST_F(IndexMap, CanBuildMapping) +{ + gko::array target_ids(ref); + gko::collection::array remote_local_idxs(ref); + gko::collection::array remote_global_idxs(ref); + + gko::kernels::reference::index_map::build_mapping( + ref, part.get(), {ref, {2, 3, 3, 5, 5}}, target_ids, remote_local_idxs, + remote_global_idxs); + + auto expected_global = gko::array{ref, {2, 3, 5}}; + auto expected_local = gko::array{ref, {0, 1, 1}}; + auto expected_ids = gko::array{ref, {1, 2}}; + GKO_ASSERT_ARRAY_EQ(target_ids, expected_ids); + GKO_ASSERT_ARRAY_EQ(remote_global_idxs.get_flat(), expected_global); + GKO_ASSERT_ARRAY_EQ(remote_local_idxs.get_flat(), expected_local); +} + + +TEST_F(IndexMap, CanGetLocalWithNonLocalIS) +{ + gko::array global_ids(ref, {1, 1, 4, 0, 4}); + gko::array local_ids(ref); + gko::collection::array remote_global_idxs( + gko::array(ref, {0, 1, 4}), std::vector{2, 1}); + gko::array remote_target_ids(ref, {0, 2}); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), remote_target_ids, remote_global_idxs, 1, global_ids, + gko::experimental::distributed::index_space::non_local, local_ids); + + gko::array expected(ref, {1, 1, 2, 0, 2}); + GKO_ASSERT_ARRAY_EQ(local_ids, expected); +} + + +TEST_F(IndexMap, CanGetLocalWithNonLocalISWithInvalid) +{ + gko::array global_ids(ref, {1, 1, 4, 3, 0, 4}); + gko::array local_ids(ref); + gko::collection::array remote_global_idxs( + gko::array(ref, {0, 1, 4}), std::vector{2, 1}); + gko::array remote_target_ids(ref, {0, 2}); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), remote_target_ids, remote_global_idxs, 1, global_ids, + gko::experimental::distributed::index_space::non_local, local_ids); + + gko::array expected(ref, {1, 1, 2, -1, 0, 2}); + GKO_ASSERT_ARRAY_EQ(local_ids, expected); +} + + +TEST_F(IndexMap, CanGetLocalWithLocalIS) +{ + gko::array global_ids(ref, {2, 3, 3, 2}); + gko::array local_ids(ref); + gko::collection::array remote_global_idxs( + gko::array(ref, {0, 1, 4}), std::vector{2, 1}); + gko::array remote_target_ids(ref, {0, 2}); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), remote_target_ids, remote_global_idxs, 1, global_ids, + gko::experimental::distributed::index_space::local, local_ids); + + gko::array expected(ref, {0, 1, 1, 0}); + GKO_ASSERT_ARRAY_EQ(local_ids, expected); +} + + +TEST_F(IndexMap, CanGetLocalWithLocalISWithInvalid) +{ + gko::array global_ids(ref, {2, 4, 5, 3, 3, 2}); + gko::array local_ids(ref); + gko::collection::array remote_global_idxs( + gko::array(ref, {0, 1, 4}), std::vector{2, 1}); + gko::array remote_target_ids(ref, {0, 2}); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), remote_target_ids, remote_global_idxs, 1, global_ids, + gko::experimental::distributed::index_space::local, local_ids); + + gko::array expected(ref, {0, -1, -1, 1, 1, 0}); + GKO_ASSERT_ARRAY_EQ(local_ids, expected); +} + + +TEST_F(IndexMap, CanGetLocalWithCombinedIS) +{ + gko::array global_ids(ref, {0, 1, 2, 3, 0, 4, 3}); + gko::array local_ids(ref); + gko::collection::array remote_global_idxs( + gko::array(ref, {0, 1, 4}), std::vector{2, 1}); + gko::array remote_target_ids(ref, {0, 2}); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), remote_target_ids, remote_global_idxs, 1, global_ids, + gko::experimental::distributed::index_space::combined, local_ids); + + gko::array expected(ref, {2, 3, 0, 1, 2, 4, 1}); + GKO_ASSERT_ARRAY_EQ(local_ids, expected); +} + + +TEST_F(IndexMap, CanGetLocalWithCombinedISWithInvalid) +{ + gko::array global_ids(ref, {0, 1, 2, 3, 0, 4, 5, 3}); + gko::array local_ids(ref); + gko::collection::array remote_global_idxs( + gko::array(ref, {0, 1, 4}), std::vector{2, 1}); + gko::array remote_target_ids(ref, {0, 2}); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), remote_target_ids, remote_global_idxs, 1, global_ids, + gko::experimental::distributed::index_space::combined, local_ids); + + gko::array expected(ref, {2, 3, 0, 1, 2, 4, -1, 1}); + GKO_ASSERT_ARRAY_EQ(local_ids, expected); +} + + +} // namespace diff --git a/test/distributed/CMakeLists.txt b/test/distributed/CMakeLists.txt index 32b3810ea31..9e0c875de0e 100644 --- a/test/distributed/CMakeLists.txt +++ b/test/distributed/CMakeLists.txt @@ -1,3 +1,4 @@ +ginkgo_create_common_test(index_map_kernels DISABLE_EXECUTORS dpcpp) ginkgo_create_common_test(matrix_kernels DISABLE_EXECUTORS dpcpp) ginkgo_create_common_test(partition_kernels DISABLE_EXECUTORS dpcpp) ginkgo_create_common_test(vector_kernels DISABLE_EXECUTORS dpcpp) diff --git a/test/distributed/index_map_kernels.cpp b/test/distributed/index_map_kernels.cpp new file mode 100644 index 00000000000..09668340ae9 --- /dev/null +++ b/test/distributed/index_map_kernels.cpp @@ -0,0 +1,325 @@ +// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include "core/distributed/index_map_kernels.hpp" + + +#include +#include + + +#include +#include + + +#include +#include +#include +#include +#include + + +#include "core/distributed/partition_kernels.hpp" +#include "core/test/utils.hpp" +#include "test/utils/executor.hpp" + + +using comm_index_type = gko::experimental::distributed::comm_index_type; + + +class IndexMap : public CommonTestFixture { +protected: + using local_index_type = gko::int32; + using global_index_type = gko::int64; + using part_type = + gko::experimental::distributed::Partition; + using map_type = + gko::experimental::distributed::index_map; + + IndexMap() + { + auto connections = generate_connection_idxs(ref, this_rank, 11); + auto dconnections = gko::array(exec, connections); + gko::kernels::reference::index_map::build_mapping( + ref, part.get(), connections, target_ids, remote_local_idxs, + remote_global_idxs); + gko::kernels::EXEC_NAMESPACE::index_map::build_mapping( + exec, dpart.get(), dconnections, dtarget_ids, dremote_local_idxs, + dremote_global_idxs); + } + + 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, + gko::size_type num_queries) + { + auto host_connection_idxs = + gko::make_temporary_clone(ref, &connection_idxs); + // generate query from connection_index_space + std::uniform_int_distribution<> dist(0, connection_idxs.get_size() - 1); + gko::array query{ref, num_queries}; + std::generate_n(query.get_data(), query.get_size(), [&] { + return host_connection_idxs->get_const_data()[dist(engine)]; + }); + return {std::move(exec), std::move(query)}; + } + + gko::array generate_complement_idxs( + std::shared_ptr exec, + const gko::array& idxs) + { + auto host_idxs = gko::make_temporary_clone(ref, &idxs); + std::vector full_idxs(part->get_size()); + std::iota(full_idxs.begin(), full_idxs.end(), 0); + + std::set idxs_set( + host_idxs->get_const_data(), + host_idxs->get_const_data() + host_idxs->get_size()); + + auto end = std::remove_if( + full_idxs.begin(), full_idxs.end(), + [&](const auto v) { return idxs_set.find(v) != idxs_set.end(); }); + auto complement_size = std::distance(full_idxs.begin(), end); + return {std::move(exec), full_idxs.begin(), end}; + } + + + gko::array combine_arrays( + std::shared_ptr exec, + const gko::array& a, + const gko::array& b) + { + gko::array result(exec, a.get_size() + b.get_size()); + exec->copy_from(a.get_executor(), a.get_size(), a.get_const_data(), + result.get_data()); + exec->copy_from(b.get_executor(), b.get_size(), b.get_const_data(), + result.get_data() + a.get_size()); + return result; + } + + gko::array take_random( + const gko::array& a, gko::size_type n) + { + auto copy = gko::array(ref, a); + std::shuffle(copy.get_data(), copy.get_data() + copy.get_size(), + engine); + + return {a.get_executor(), copy.get_const_data(), + copy.get_const_data() + n}; + } + + gko::array target_ids{ref}; + gko::collection::array remote_local_idxs{ref}; + gko::collection::array remote_global_idxs{ref}; + gko::array dtarget_ids{exec}; + gko::collection::array dremote_local_idxs{exec}; + gko::collection::array dremote_global_idxs{exec}; + + 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); + + 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::collection::array remote_local_idxs{ref}; + gko::collection::array remote_global_idxs{ref}; + gko::array dtarget_ids{exec}; + gko::collection::array dremote_local_idxs{exec}; + gko::collection::array dremote_global_idxs{exec}; + + gko::kernels::reference::index_map::build_mapping( + ref, part.get(), query, target_ids, remote_local_idxs, + remote_global_idxs); + gko::kernels::EXEC_NAMESPACE::index_map::build_mapping( + exec, dpart.get(), dquery, dtarget_ids, dremote_local_idxs, + dremote_global_idxs); + + GKO_ASSERT_ARRAY_EQ(target_ids, dtarget_ids); + GKO_ASSERT_ARRAY_EQ(remote_local_idxs.get_flat(), + dremote_local_idxs.get_flat()); + GKO_ASSERT_ARRAY_EQ(remote_global_idxs.get_flat(), + dremote_global_idxs.get_flat()); +} + + +TEST_F(IndexMap, GetLocalWithLocalIndexSpaceSameAsRef) +{ + auto local_space = gko::array(ref, local_size); + std::iota(local_space.get_data(), local_space.get_data() + local_size, + this_rank * local_size); + auto query = generate_query(ref, local_space, 33); + auto dquery = gko::array(exec, query); + auto result = gko::array(ref); + auto dresult = gko::array(exec); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), target_ids, remote_global_idxs, this_rank, query, + gko::experimental::distributed::index_space::local, result); + gko::kernels::EXEC_NAMESPACE::index_map::get_local( + exec, dpart.get(), dtarget_ids, dremote_global_idxs, this_rank, dquery, + gko::experimental::distributed::index_space::local, dresult); + + GKO_ASSERT_ARRAY_EQ(result, dresult); +} + + +TEST_F(IndexMap, GetLocalWithLocalIndexSpaceWithInvalidIndexSameAsRef) +{ + auto local_space = gko::array(ref, local_size); + std::iota(local_space.get_data(), local_space.get_data() + local_size, + this_rank * local_size); + auto query = generate_query( + ref, + combine_arrays( + ref, local_space, + take_random(generate_complement_idxs(ref, local_space), 12)), + 33); + auto dquery = gko::array(exec, query); + auto result = gko::array(ref); + auto dresult = gko::array(exec); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), target_ids, remote_global_idxs, this_rank, query, + gko::experimental::distributed::index_space::local, result); + gko::kernels::EXEC_NAMESPACE::index_map::get_local( + exec, dpart.get(), dtarget_ids, dremote_global_idxs, this_rank, dquery, + gko::experimental::distributed::index_space::local, dresult); + + GKO_ASSERT_ARRAY_EQ(result, dresult); +} + + +TEST_F(IndexMap, GetLocalWithNonLocalIndexSpaceSameAsRef) +{ + auto query = generate_query(ref, remote_global_idxs.get_flat(), 33); + auto dquery = gko::array(exec, query); + auto result = gko::array(ref); + auto dresult = gko::array(exec); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), target_ids, remote_global_idxs, this_rank, query, + gko::experimental::distributed::index_space::non_local, result); + gko::kernels::EXEC_NAMESPACE::index_map::get_local( + exec, dpart.get(), dtarget_ids, dremote_global_idxs, this_rank, dquery, + gko::experimental::distributed::index_space::non_local, dresult); + + GKO_ASSERT_ARRAY_EQ(result, dresult); +} + + +TEST_F(IndexMap, GetLocalWithNonLocalIndexSpaceWithInvalidIndexSameAsRef) +{ + auto query = generate_query( + ref, + combine_arrays(ref, remote_global_idxs.get_flat(), + take_random(generate_complement_idxs( + ref, remote_global_idxs.get_flat()), + 12)), + 33); + auto dquery = gko::array(exec, query); + auto result = gko::array(ref); + auto dresult = gko::array(exec); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), target_ids, remote_global_idxs, this_rank, query, + gko::experimental::distributed::index_space::non_local, result); + gko::kernels::EXEC_NAMESPACE::index_map::get_local( + exec, dpart.get(), dtarget_ids, dremote_global_idxs, this_rank, dquery, + gko::experimental::distributed::index_space::non_local, dresult); + + GKO_ASSERT_ARRAY_EQ(result, dresult); +} + + +TEST_F(IndexMap, GetLocalWithCombinedIndexSpaceSameAsRef) +{ + auto local_space = gko::array(ref, local_size); + std::iota(local_space.get_data(), local_space.get_data() + local_size, + this_rank * local_size); + auto combined_space = + combine_arrays(ref, local_space, remote_global_idxs.get_flat()); + auto query = generate_query(ref, combined_space, 33); + auto dquery = gko::array(exec, query); + auto result = gko::array(ref); + auto dresult = gko::array(exec); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), target_ids, remote_global_idxs, this_rank, query, + gko::experimental::distributed::index_space::combined, result); + gko::kernels::EXEC_NAMESPACE::index_map::get_local( + exec, dpart.get(), dtarget_ids, dremote_global_idxs, this_rank, dquery, + gko::experimental::distributed::index_space::combined, dresult); + + GKO_ASSERT_ARRAY_EQ(result, dresult); +} + + +TEST_F(IndexMap, GetLocalWithCombinedIndexSpaceWithInvalidIndexSameAsRef) +{ + auto local_space = gko::array(ref, local_size); + std::iota(local_space.get_data(), local_space.get_data() + local_size, + this_rank * local_size); + auto combined_space = + combine_arrays(ref, local_space, remote_global_idxs.get_flat()); + auto query = generate_query( + ref, + combine_arrays( + ref, combined_space, + take_random(generate_complement_idxs(ref, combined_space), 12)), + 33); + auto dquery = gko::array(exec, query); + auto result = gko::array(ref); + auto dresult = gko::array(exec); + + gko::kernels::reference::index_map::get_local( + ref, part.get(), target_ids, remote_global_idxs, this_rank, query, + gko::experimental::distributed::index_space::non_local, result); + gko::kernels::EXEC_NAMESPACE::index_map::get_local( + exec, dpart.get(), dtarget_ids, dremote_global_idxs, this_rank, dquery, + gko::experimental::distributed::index_space::non_local, dresult); + + GKO_ASSERT_ARRAY_EQ(result, dresult); +}