adds distributed index map core + reference tests

core/test/CMakeLists.txt

@@ -5,6 +5,7 @@ add_subdirectory(accessor)
 add_subdirectory(base)
 add_subdirectory(components)
 add_subdirectory(config)
+add_subdirectory(distributed)
 if(GINKGO_BUILD_MPI)
  add_subdirectory(mpi)
 endif()

core/test/distributed/CMakeLists.txt

@@ -0,0 +1 @@
+ginkgo_create_test(index_map)

core/test/distributed/index_map.cpp

@@ -0,0 +1,165 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/distributed/index_map.hpp>
+
+
+#include <gtest/gtest.h>
+
+
+#include <ginkgo/core/base/executor.hpp>
+
+
+#include "core/test/utils.hpp"
+
+
+using gko::experimental::distributed::comm_index_type;
+
+
+template <typename LocalGlobalIndexType>
+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<local_index_type,
+ global_index_type>;
+ using map_type =
+ gko::experimental::distributed::index_map<local_index_type,
+ global_index_type>;
+
+ std::shared_ptr<const gko::Executor> exec =
+ gko::ReferenceExecutor::create();
+ std::shared_ptr<part_type> part =
+ part_type::build_from_global_size_uniform(exec, 3, 6);
+};
+
+TYPED_TEST_SUITE(IndexMap, gko::test::LocalGlobalIndexTypes,
+ PairTypenameNameGenerator);
+
+
+template <typename T>
+void assert_collection_eq(const gko::segmented_array<T>& a,
+ const gko::segmented_array<T>& b)
+{
+ auto get_flat_array = [](const gko::segmented_array<T>& arr) {
+ return gko::make_const_array_view(arr.get_executor(), arr.get_size(),
+ arr.get_const_flat_data());
+ };
+ GKO_ASSERT_ARRAY_EQ(a.get_offsets(), b.get_offsets());
+ GKO_ASSERT_ARRAY_EQ(get_flat_array(a), get_flat_array(b));
+}
+
+
+TYPED_TEST(IndexMap, CanDefaultConstruct)
+{
+ using map_type = typename TestFixture::map_type;
+
+ auto imap = map_type(this->exec);
+
+ ASSERT_EQ(imap.get_local_size(), 0);
+ ASSERT_EQ(imap.get_non_local_size(), 0);
+ ASSERT_EQ(imap.get_remote_target_ids().get_size(), 0);
+ ASSERT_EQ(imap.get_remote_global_idxs().get_size(), 0);
+ ASSERT_EQ(imap.get_remote_local_idxs().get_size(), 0);
+}
+
+
+TYPED_TEST(IndexMap, CanCopyConstruct)
+{
+ using map_type = typename TestFixture::map_type;
+ using global_index_type = typename TestFixture::global_index_type;
+ gko::array<global_index_type> 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<global_index_type> 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_const_flat_data();
+ auto imap_remote_local_it =
+ imap.get_remote_local_idxs().get_const_flat_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_const_flat_data(),
+ imap_remote_local_it);
+ ASSERT_EQ(move.get_remote_global_idxs().get_const_flat_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<global_index_type> connections(this->exec, {4, 3, 3, 4, 2});
+ auto imap = map_type(this->exec, this->part, 0, connections);
+ auto copy = map_type(this->exec);
+
+ 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<global_index_type> 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_const_flat_data();
+ auto imap_remote_local_it =
+ imap.get_remote_local_idxs().get_const_flat_data();
+ auto imap_remote_target_it = imap.get_remote_target_ids().get_const_data();
+ auto move = map_type(this->exec);
+
+ 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_const_flat_data(),
+ imap_remote_local_it);
+ ASSERT_EQ(move.get_remote_global_idxs().get_const_flat_data(),
+ imap_remote_global_it);
+}

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)

reference/test/distributed/index_map_kernels.cpp

@@ -0,0 +1,200 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/distributed/index_map.hpp>
+
+
+#include <algorithm>
+#include <memory>
+#include <vector>
+
+
+#include <gtest/gtest-typed-test.h>
+#include <gtest/gtest.h>
+
+
+#include <ginkgo/core/base/executor.hpp>
+
+
+#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<local_index_type,
+ global_index_type>;
+ using map_type =
+ gko::experimental::distributed::index_map<local_index_type,
+ global_index_type>;
+
+ IndexMap() : ref(gko::ReferenceExecutor::create()) {}
+
+ std::shared_ptr<const gko::ReferenceExecutor> ref;
+ std::shared_ptr<const part_type> part =
+ part_type::build_from_mapping(ref, {ref, {0, 0, 1, 1, 2, 2}}, 3);
+};
+
+
+TEST_F(IndexMap, CanBuildMapping)
+{
+ gko::array<comm_index_type> target_ids(ref);
+ gko::array<local_index_type> remote_local_idxs(ref);
+ gko::array<global_index_type> remote_global_idxs(ref);
+ gko::array<gko::int64> remote_sizes(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, remote_sizes);
+
+ auto expected_global = gko::array<global_index_type>{ref, {2, 3, 5}};
+ auto expected_local = gko::array<local_index_type>{ref, {0, 1, 1}};
+ auto expected_ids = gko::array<comm_index_type>{ref, {1, 2}};
+ auto expected_sizes = gko::array<gko::int64>{ref, {2, 1}};
+ GKO_ASSERT_ARRAY_EQ(remote_sizes, expected_sizes);
+ GKO_ASSERT_ARRAY_EQ(target_ids, expected_ids);
+ GKO_ASSERT_ARRAY_EQ(remote_global_idxs, expected_global);
+ GKO_ASSERT_ARRAY_EQ(remote_local_idxs, expected_local);
+}
+
+
+TEST_F(IndexMap, CanBuildMappingWithoutRecvConnections)
+{
+ gko::array<comm_index_type> target_ids(ref);
+ gko::array<local_index_type> remote_local_idxs(ref);
+ gko::array<global_index_type> remote_global_idxs(ref);
+ gko::array<gko::int64> remote_sizes(ref);
+
+ gko::kernels::reference::index_map::build_mapping(
+ ref, part.get(), {ref, 0}, target_ids, remote_local_idxs,
+ remote_global_idxs, remote_sizes);
+
+ auto expected_global = gko::array<global_index_type>{ref, 0};
+ auto expected_local = gko::array<local_index_type>{ref, 0};
+ auto expected_ids = gko::array<comm_index_type>{ref, 0};
+ auto expected_sizes = gko::array<gko::int64>{ref};
+ GKO_ASSERT_ARRAY_EQ(remote_sizes, expected_sizes);
+ GKO_ASSERT_ARRAY_EQ(target_ids, expected_ids);
+ GKO_ASSERT_ARRAY_EQ(remote_global_idxs, expected_global);
+ GKO_ASSERT_ARRAY_EQ(remote_local_idxs, expected_local);
+}
+
+
+TEST_F(IndexMap, CanGetLocalWithNonLocalIS)
+{
+ gko::array<global_index_type> global_ids(ref, {1, 1, 4, 0, 4});
+ gko::array<local_index_type> local_ids(ref);
+ auto remote_global_idxs =
+ gko::segmented_array<global_index_type>::create_from_sizes(
+ {ref, {0, 1, 4}}, {ref, {2, 1}});
+ gko::array<comm_index_type> 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<local_index_type> expected(ref, {1, 1, 2, 0, 2});
+ GKO_ASSERT_ARRAY_EQ(local_ids, expected);
+}
+
+
+TEST_F(IndexMap, CanGetLocalWithNonLocalISWithInvalid)
+{
+ gko::array<global_index_type> global_ids(ref, {1, 1, 4, 3, 0, 4});
+ gko::array<local_index_type> local_ids(ref);
+ auto remote_global_idxs =
+ gko::segmented_array<global_index_type>::create_from_sizes(
+ {ref, {0, 1, 4}}, {ref, {2, 1}});
+ gko::array<comm_index_type> 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<local_index_type> expected(ref, {1, 1, 2, -1, 0, 2});
+ GKO_ASSERT_ARRAY_EQ(local_ids, expected);
+}
+
+
+TEST_F(IndexMap, CanGetLocalWithLocalIS)
+{
+ gko::array<global_index_type> global_ids(ref, {2, 3, 3, 2});
+ gko::array<local_index_type> local_ids(ref);
+ auto remote_global_idxs =
+ gko::segmented_array<global_index_type>::create_from_sizes(
+ {ref, {0, 1, 4}}, {ref, {2, 1}});
+ gko::array<comm_index_type> 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<local_index_type> expected(ref, {0, 1, 1, 0});
+ GKO_ASSERT_ARRAY_EQ(local_ids, expected);
+}
+
+
+TEST_F(IndexMap, CanGetLocalWithLocalISWithInvalid)
+{
+ gko::array<global_index_type> global_ids(ref, {2, 4, 5, 3, 3, 2});
+ gko::array<local_index_type> local_ids(ref);
+ auto remote_global_idxs =
+ gko::segmented_array<global_index_type>::create_from_sizes(
+ {ref, {0, 1, 4}}, {ref, {2, 1}});
+ gko::array<comm_index_type> 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<local_index_type> expected(ref, {0, -1, -1, 1, 1, 0});
+ GKO_ASSERT_ARRAY_EQ(local_ids, expected);
+}
+
+
+TEST_F(IndexMap, CanGetLocalWithCombinedIS)
+{
+ gko::array<global_index_type> global_ids(ref, {0, 1, 2, 3, 0, 4, 3});
+ gko::array<local_index_type> local_ids(ref);
+ auto remote_global_idxs =
+ gko::segmented_array<global_index_type>::create_from_sizes(
+ {ref, {0, 1, 4}}, {ref, {2, 1}});
+ gko::array<comm_index_type> 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<local_index_type> expected(ref, {2, 3, 0, 1, 2, 4, 1});
+ GKO_ASSERT_ARRAY_EQ(local_ids, expected);
+}
+
+
+TEST_F(IndexMap, CanGetLocalWithCombinedISWithInvalid)
+{
+ gko::array<global_index_type> global_ids(ref, {0, 1, 2, 3, 0, 4, 5, 3});
+ gko::array<local_index_type> local_ids(ref);
+ auto remote_global_idxs =
+ gko::segmented_array<global_index_type>::create_from_sizes(
+ {ref, {0, 1, 4}}, {ref, {2, 1}});
+ gko::array<comm_index_type> 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<local_index_type> expected(ref, {2, 3, 0, 1, 2, 4, -1, 1});
+ GKO_ASSERT_ARRAY_EQ(local_ids, expected);
+}
+
+
+} // namespace