adds distributed index map core + reference tests

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()

core/test/distributed/CMakeLists.txt

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

core/test/distributed/index_map.cpp

@@ -0,0 +1,172 @@
+// 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_default_constructed_array(const gko::array<T>& a)
+{
+    ASSERT_EQ(a.get_executor(), nullptr);
+    ASSERT_EQ(a.get_size(), 0);
+}
+
+template <typename T>
+void assert_collection_eq(const gko::collection::array<T>& a,
+                          const gko::collection::array<T>& 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<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_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<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();
+
+    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_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);
+}

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,169 @@
+// 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::collection::array<local_index_type> remote_local_idxs(ref);
+    gko::collection::array<global_index_type> 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<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}};
+    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_index_type> global_ids(ref, {1, 1, 4, 0, 4});
+    gko::array<local_index_type> local_ids(ref);
+    gko::collection::array<global_index_type> remote_global_idxs(
+        gko::array<global_index_type>(ref, {0, 1, 4}), std::vector<int>{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);
+    gko::collection::array<global_index_type> remote_global_idxs(
+        gko::array<global_index_type>(ref, {0, 1, 4}), std::vector<int>{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);
+    gko::collection::array<global_index_type> remote_global_idxs(
+        gko::array<global_index_type>(ref, {0, 1, 4}), std::vector<int>{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);
+    gko::collection::array<global_index_type> remote_global_idxs(
+        gko::array<global_index_type>(ref, {0, 1, 4}), std::vector<int>{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);
+    gko::collection::array<global_index_type> remote_global_idxs(
+        gko::array<global_index_type>(ref, {0, 1, 4}), std::vector<int>{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);
+    gko::collection::array<global_index_type> remote_global_idxs(
+        gko::array<global_index_type>(ref, {0, 1, 4}), std::vector<int>{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