[dist] review updates:

- reduce tests
- update docs
- minor refactoring

Co-authored-by: Yu-Hsiang M. Tsai <[email protected]>
Co-authored-by: Pratik Nayak <[email protected]>

core/distributed/device_partition.hpp

@@ -2,8 +2,7 @@
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
-#ifndef GINKGO_PARTITION_HPP
-#define GINKGO_PARTITION_HPP
+#pragma once
 
 #include <ginkgo/core/distributed/partition.hpp>
 
@@ -34,34 +33,7 @@ struct device_partition {
 
 
 /**
- * Create device_segmented_array from a segmented_array.
- */
-template <typename LocalIndexType, typename GlobalIndexType>
-constexpr device_partition<const LocalIndexType, const GlobalIndexType>
-to_device(
-    const experimental::distributed::Partition<LocalIndexType, GlobalIndexType>*
-        partition)
-{
-    auto num_ranges = partition->get_num_ranges();
-    auto num_parts = partition->get_num_parts();
-    return {num_parts,
-            partition->get_num_empty_parts(),
-            partition->get_size(),
-            partition->get_range_bounds(),
-            partition->get_range_bounds() + num_ranges + 1,
-            partition->get_range_starting_indices(),
-            partition->get_range_starting_indices() + num_ranges,
-            partition->get_part_sizes(),
-            partition->get_part_sizes() + num_parts,
-            partition->get_part_ids(),
-            partition->get_part_ids() + num_parts,
-            to_device(partition->get_ranges_by_part())};
-}
-
-/**
- * Explicitly create a const version of device_segmented_array.
- *
- * This is mostly relevant for tests.
+ * Explicitly create a const version of device_partition.
  */
 template <typename LocalIndexType, typename GlobalIndexType>
 constexpr device_partition<const LocalIndexType, const GlobalIndexType>
@@ -87,6 +59,3 @@ to_device_const(
 
 
 }  // namespace gko
-
-
-#endif  // GINKGO_PARTITION_HPP

core/distributed/index_map.cpp

@@ -98,8 +98,8 @@ index_map<LocalIndexType, GlobalIndexType>::map_to_global(
     array<GlobalIndexType> global_ids(exec_);
 
     exec_->run(index_map_kernels::make_map_to_global(
-        to_device(partition_.get()), to_device(remote_global_idxs_), rank_,
-        local_ids, index_space_v, global_ids));
+        to_device_const(partition_.get()), to_device(remote_global_idxs_),
+        rank_, local_ids, index_space_v, global_ids));
 
     return global_ids;
 }

core/distributed/index_map_kernels.hpp

@@ -54,11 +54,11 @@ namespace kernels {
  * 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
+ * - 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
+ *                       ids, and each segment sorted
  * - rank: the part id of this process
  *
  * Any global index that is not in the specified local index space is mapped
@@ -81,7 +81,7 @@ namespace kernels {
  *
  * The relevant input parameter from the index map are:
  *
- * - partition:  the global partition
+ * - partition: the global partition
  * - remote_global_idxs: the remote global indices, segmented by the owning part
  *                       ids
  * - rank: the part id of this process

reference/distributed/partition_helpers.hpp

@@ -87,6 +87,8 @@ GlobalIndexType map_to_global(
     device_partition<const LocalIndexType, const GlobalIndexType> partition,
     size_type range_id)
 {
+    assert(range_id <
+           std::distance(partition.offsets_begin, partition.offsets_end) - 1);
     auto range_bounds = partition.offsets_begin;
     auto starting_indices = partition.starting_indices_begin;
     return static_cast<GlobalIndexType>(idx - starting_indices[range_id]) +

reference/distributed/partition_kernels.cpp

@@ -5,6 +5,7 @@
 #include "core/distributed/partition_kernels.hpp"
 
 #include "core/base/segmented_array.hpp"
+#include "ginkgo/core/base/math.hpp"
 
 
 namespace gko {
@@ -128,7 +129,7 @@ void build_ranges_by_part(std::shared_ptr<const DefaultExecutor> exec,
               });
 
     sizes.resize_and_reset(num_parts);
-    std::fill_n(sizes.get_data(), num_parts, int64(0));
+    std::fill_n(sizes.get_data(), num_parts, zero<int64>());
     for (size_type i = 0; i < num_ranges; ++i) {
         sizes.get_data()[range_parts[i]]++;
     }

reference/test/distributed/index_map_kernels.cpp

@@ -198,22 +198,6 @@ TEST_F(IndexMap, CanGetLocalWithCombinedISWithInvalid)
 }
 
 
-TEST_F(IndexMap, CanGetGlobalWithLocalIS)
-{
-    gko::array<global_index_type> global_ids(ref);
-    gko::array<local_index_type> local_ids(ref, {5, 4, 3, 2, 1, 0, 4});
-    auto remote_global_idxs = gko::segmented_array<global_index_type>{ref};
-
-    gko::kernels::reference::index_map::map_to_global(
-        ref, to_device_const(part_large.get()),
-        to_device_const(remote_global_idxs), 1, local_ids,
-        gko::experimental::distributed::index_space::local, global_ids);
-
-    gko::array<global_index_type> expected(ref, {14, 13, 12, 5, 4, 3, 13});
-    GKO_ASSERT_ARRAY_EQ(global_ids, expected);
-}
-
-
 TEST_F(IndexMap, CanGetGlobalWithLocalISWithInvalid)
 {
     gko::array<global_index_type> global_ids(ref);
@@ -222,87 +206,98 @@ TEST_F(IndexMap, CanGetGlobalWithLocalISWithInvalid)
 
     gko::kernels::reference::index_map::map_to_global(
         ref, to_device_const(part_large.get()),
-        to_device_const(remote_global_idxs), 1, local_ids,
+        to_device_const(remote_global_idxs), 0, local_ids,
         gko::experimental::distributed::index_space::local, global_ids);
 
     auto invalid = gko::invalid_index<global_index_type>();
     gko::array<global_index_type> expected(
-        ref, I<global_index_type>{14, 13, invalid, 12, 5, 4, 3, invalid, 13});
+        ref, I<global_index_type>{17, 16, invalid, 15, 2, 1, 0, invalid, 16});
     GKO_ASSERT_ARRAY_EQ(global_ids, expected);
 }
 
 
-TEST_F(IndexMap, CanGetGlobalWithNonLocalIS)
+TEST_F(IndexMap, CanGetGlobalWithNonLocalISWithInvalid)
 {
     gko::array<global_index_type> global_ids(ref);
-    gko::array<local_index_type> local_ids(ref, {5, 4, 3, 2, 1, 0, 4});
+    gko::array<local_index_type> local_ids(ref, {5, 4, 10, 3, 2, 1, 0, 100, 4});
     auto remote_global_idxs =
         gko::segmented_array<global_index_type>::create_from_sizes(
-            {ref, {0, 1, 2, 17, 16, 15}}, {ref, {2, 4}});
+            {ref, {0, 1, 9, 8, 7, 6}}, {ref, {2, 4}});
 
     gko::kernels::reference::index_map::map_to_global(
         ref, to_device_const(part_large.get()),
         to_device_const(remote_global_idxs), 1, local_ids,
         gko::experimental::distributed::index_space::non_local, global_ids);
 
-    gko::array<global_index_type> expected(ref, {15, 16, 17, 2, 1, 0, 16});
+    auto invalid = gko::invalid_index<global_index_type>();
+    gko::array<global_index_type> expected(
+        ref, I<global_index_type>{6, 7, invalid, 8, 9, 1, 0, invalid, 7});
     GKO_ASSERT_ARRAY_EQ(global_ids, expected);
 }
 
 
-TEST_F(IndexMap, CanGetGlobalWithNonLocalISWithInvalid)
+TEST_F(IndexMap, CanGetGlobalWithCombinedISWithInvalid)
 {
     gko::array<global_index_type> global_ids(ref);
-    gko::array<local_index_type> local_ids(ref, {5, 4, 10, 3, 2, 1, 0, 100, 4});
+    gko::array<local_index_type> local_ids(ref, {2, 5, 133, 6, 10});
     auto remote_global_idxs =
         gko::segmented_array<global_index_type>::create_from_sizes(
-            {ref, {0, 1, 2, 17, 16, 15}}, {ref, {2, 4}});
+            {ref, {0, 1, 9, 8, 7, 6}}, {ref, {2, 4}});
 
     gko::kernels::reference::index_map::map_to_global(
         ref, to_device_const(part_large.get()),
         to_device_const(remote_global_idxs), 1, local_ids,
-        gko::experimental::distributed::index_space::non_local, global_ids);
+        gko::experimental::distributed::index_space::combined, global_ids);
 
     auto invalid = gko::invalid_index<global_index_type>();
     gko::array<global_index_type> expected(
-        ref, I<global_index_type>{15, 16, invalid, 17, 2, 1, 0, invalid, 16});
+        ref, I<global_index_type>{5, 14, invalid, 0, 7});
     GKO_ASSERT_ARRAY_EQ(global_ids, expected);
 }
 
 
-TEST_F(IndexMap, CanGetGlobalWithCombinedIS)
+TEST_F(IndexMap, RoundTripGlobalWithCombinedIS)
 {
-    gko::array<global_index_type> global_ids(ref);
-    gko::array<local_index_type> local_ids(ref, {2, 5, 6, 10});
+    gko::array<global_index_type> result(ref);
+    gko::array<global_index_type> global_ids(ref,
+                                             {5, 14, 14, 0, 3, 4, 8, 7, 12});
+    gko::array<local_index_type> local_ids(ref);
     auto remote_global_idxs =
         gko::segmented_array<global_index_type>::create_from_sizes(
-            {ref, {0, 1, 2, 17, 16, 15}}, {ref, {2, 4}});
+            {ref, {0, 1, 6, 7, 8, 9}}, {ref, {2, 4}});
+    gko::array<comm_index_type> remote_target_ids(ref, {0, 2});
 
+    gko::kernels::reference::index_map::map_to_local(
+        ref, part_large.get(), remote_target_ids,
+        to_device_const(remote_global_idxs), 1, global_ids,
+        gko::experimental::distributed::index_space::combined, local_ids);
     gko::kernels::reference::index_map::map_to_global(
         ref, to_device_const(part_large.get()),
         to_device_const(remote_global_idxs), 1, local_ids,
-        gko::experimental::distributed::index_space::combined, global_ids);
+        gko::experimental::distributed::index_space::combined, result);
 
-    gko::array<global_index_type> expected(ref, {5, 14, 0, 16});
-    GKO_ASSERT_ARRAY_EQ(global_ids, expected);
+    GKO_ASSERT_ARRAY_EQ(result, global_ids);
 }
 
 
-TEST_F(IndexMap, CanGetGlobalWithCombinedISWithInvalid)
+TEST_F(IndexMap, RoundTripLocalWithCombinedIS)
 {
+    gko::array<local_index_type> result(ref);
     gko::array<global_index_type> global_ids(ref);
-    gko::array<local_index_type> local_ids(ref, {2, 5, 133, 6, 10});
+    gko::array<local_index_type> local_ids(ref, {2, 5, 4, 6, 3, 3, 10});
     auto remote_global_idxs =
         gko::segmented_array<global_index_type>::create_from_sizes(
-            {ref, {0, 1, 2, 17, 16, 15}}, {ref, {2, 4}});
+            {ref, {0, 1, 6, 7, 8, 9}}, {ref, {2, 4}});
+    gko::array<comm_index_type> remote_target_ids(ref, {0, 2});
 
     gko::kernels::reference::index_map::map_to_global(
         ref, to_device_const(part_large.get()),
         to_device_const(remote_global_idxs), 1, local_ids,
         gko::experimental::distributed::index_space::combined, global_ids);
+    gko::kernels::reference::index_map::map_to_local(
+        ref, part_large.get(), remote_target_ids,
+        to_device_const(remote_global_idxs), 1, global_ids,
+        gko::experimental::distributed::index_space::combined, result);
 
-    auto invalid = gko::invalid_index<global_index_type>();
-    gko::array<global_index_type> expected(
-        ref, I<global_index_type>{5, 14, invalid, 0, 16});
-    GKO_ASSERT_ARRAY_EQ(global_ids, expected);
+    GKO_ASSERT_ARRAY_EQ(result, local_ids);
 }

test/distributed/index_map_kernels.cpp

@@ -120,8 +120,6 @@ class IndexMap : public CommonTestFixture {
         std::random_device rd;
         std::uniform_int_distribution<> engine_dist(
             0, std::numeric_limits<int>::max());
-        auto seed = engine_dist(rd);
-        std::cout << "seed = " << seed << std::endl;
         engine.seed(490729788);
 
         auto connections =