[dist] review updates:

- reduce tests
- update docs
- minor refactoring
- fix binary search usage in cuda/hip
- refactor
- add tests for segmented array assertion

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

common/cuda_hip/distributed/index_map_kernels.cpp

@@ -302,20 +302,20 @@ void map_to_global(
     device_partition<const LocalIndexType, const GlobalIndexType> partition,
     device_segmented_array<const GlobalIndexType> remote_global_idxs,
     experimental::distributed::comm_index_type rank,
-    const array<LocalIndexType>& local_ids,
+    const array<LocalIndexType>& local_idxs,
     experimental::distributed::index_space is,
-    array<GlobalIndexType>& global_ids)
+    array<GlobalIndexType>& global_idxs)
 {
     auto range_bounds = partition.offsets_begin;
     auto starting_indices = partition.starting_indices_begin;
     const auto& ranges_by_part = partition.ranges_by_part;
-    auto local_ids_it = local_ids.get_const_data();
-    auto input_size = local_ids.get_size();
+    auto local_idxs_it = local_idxs.get_const_data();
+    auto input_size = local_idxs.get_size();
 
     auto policy = thrust_policy(exec);
 
-    global_ids.resize_and_reset(local_ids.get_size());
-    auto global_ids_it = global_ids.get_data();
+    global_idxs.resize_and_reset(local_idxs.get_size());
+    auto global_idxs_it = global_idxs.get_data();
 
     auto map_local = [rank, ranges_by_part, range_bounds, starting_indices,
                       partition] __device__(auto lid) {
@@ -330,11 +330,16 @@ void map_to_global(
         auto local_ranges_size =
             static_cast<int64>(local_ranges.end - local_ranges.begin);
 
-        auto it = binary_search(int64(0), local_ranges_size, [=](const auto i) {
-            return starting_indices[local_ranges.begin[i]] >= lid;
-        });
+        // the binary search finds the first local range, such that the starting
+        // index is larger than lid, thus lid is contained in the local range
+        // before that one
         auto local_range_id =
-            it != local_ranges_size ? it : max(int64(0), it - 1);
+            binary_search(int64(0), local_ranges_size,
+                          [=](const auto i) {
+                              return starting_indices[local_ranges.begin[i]] >
+                                     lid;
+                          }) -
+            1;
         auto range_id = local_ranges.begin[local_range_id];
 
         return static_cast<GlobalIndexType>(lid - starting_indices[range_id]) +
@@ -363,16 +368,16 @@ void map_to_global(
     };
 
     if (is == experimental::distributed::index_space::local) {
-        thrust::transform(policy, local_ids_it, local_ids_it + input_size,
-                          global_ids_it, map_local);
+        thrust::transform(policy, local_idxs_it, local_idxs_it + input_size,
+                          global_idxs_it, map_local);
     }
     if (is == experimental::distributed::index_space::non_local) {
-        thrust::transform(policy, local_ids_it, local_ids_it + input_size,
-                          global_ids_it, map_non_local);
+        thrust::transform(policy, local_idxs_it, local_idxs_it + input_size,
+                          global_idxs_it, map_non_local);
     }
     if (is == experimental::distributed::index_space::combined) {
-        thrust::transform(policy, local_ids_it, local_ids_it + input_size,
-                          global_ids_it, map_combined);
+        thrust::transform(policy, local_idxs_it, local_idxs_it + input_size,
+                          global_idxs_it, map_combined);
     }
 }
 

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

@@ -93,15 +93,15 @@ array<LocalIndexType> index_map<LocalIndexType, GlobalIndexType>::map_to_local(
 template <typename LocalIndexType, typename GlobalIndexType>
 array<GlobalIndexType>
 index_map<LocalIndexType, GlobalIndexType>::map_to_global(
-    const array<LocalIndexType>& local_ids, index_space index_space_v) const
+    const array<LocalIndexType>& local_idxs, index_space index_space_v) const
 {
-    array<GlobalIndexType> global_ids(exec_);
+    array<GlobalIndexType> global_idxs(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_idxs, index_space_v, global_idxs));
 
-    return global_ids;
+    return global_idxs;
 }
 
 

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
@@ -95,8 +95,8 @@ namespace kernels {
         device_partition<const _ltype, const _gtype> partition,  \
         device_segmented_array<const _gtype> remote_global_idxs, \
         experimental::distributed::comm_index_type rank,         \
-        const array<_ltype>& local_ids,                          \
-        experimental::distributed::index_space is, array<_gtype>& global_ids)
+        const array<_ltype>& local_idxs,                         \
+        experimental::distributed::index_space is, array<_gtype>& global_idxs)
 
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                      \

core/test/utils/assertions.hpp

@@ -24,6 +24,7 @@
 #include <ginkgo/core/base/math.hpp>
 #include <ginkgo/core/base/mtx_io.hpp>
 #include <ginkgo/core/base/name_demangling.hpp>
+#include <ginkgo/core/base/segmented_array.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
 
 #include "core/base/batch_utilities.hpp"
@@ -1014,19 +1015,19 @@ ::testing::AssertionResult segmented_array_equal(
                                    second.get_const_flat_data())
             .copy_to_array();
 
-    auto buffer_result = array_equal(first_expression, second_expression,
-                                     view_first, view_second);
-    if (buffer_result == ::testing::AssertionFailure()) {
-        return buffer_result << "Buffers of the segmented arrays mismatch";
-    }
-
     auto offsets_result =
         array_equal(first_expression, second_expression, first.get_offsets(),
                     second.get_offsets());
     if (offsets_result == ::testing::AssertionFailure()) {
         return offsets_result << "Offsets of the segmented arrays mismatch";
     }
 
+    auto buffer_result = array_equal(first_expression, second_expression,
+                                     view_first, view_second);
+    if (buffer_result == ::testing::AssertionFailure()) {
+        return buffer_result << "Buffers of the segmented arrays mismatch";
+    }
+
     return ::testing::AssertionSuccess();
 }
 
@@ -1414,6 +1415,18 @@ T* plain_ptr(T* ptr)
     }
 
 
+/**
+ * Checks if two `gko::segmented_array`s are equal.
+ *
+ * Both the flat array buffer and the offsets of both arrays are tested
+ * for equality.
+ *
+ * Has to be called from within a google test unit test.
+ * Internally calls gko::test::assertions::segmented_array_equal().
+ *
+ * @param _array1 first segmented array
+ * @param _array2 second segmented array
+ */
 #define GKO_ASSERT_SEGMENTED_ARRAY_EQ(_array1, _array2)                      \
     {                                                                        \
         ASSERT_PRED_FORMAT2(::gko::test::assertions::segmented_array_equal,  \

core/test/utils/assertions_test.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -8,6 +8,7 @@
 
 #include <gtest/gtest.h>
 
+#include <ginkgo/core/base/segmented_array.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
 
@@ -218,4 +219,52 @@ TEST_F(ArraysNear, CanUseShortNotation)
 }
 
 
+class SegmentedArraysEqual : public ::testing::Test {
+protected:
+    using array = gko::array<double>;
+    using iarray = gko::array<gko::int64>;
+    using segmented_array = gko::segmented_array<double>;
+
+    std::shared_ptr<gko::Executor> exec = gko::ReferenceExecutor::create();
+
+    segmented_array arr1 = segmented_array::create_from_sizes(
+        array{exec, {1, 2, 3, 4, 5}}, iarray{exec, {2, 1, 2}});
+    segmented_array arr2 = segmented_array::create_from_sizes(
+        array{exec, {1, 2, 3, 4, 5}}, iarray{exec, {2, 1, 2}});
+    segmented_array arr3 = segmented_array::create_from_sizes(
+        array{exec, {1, 2, 3, 5, 6}}, iarray{exec, {2, 1, 2}});
+    segmented_array arr4 = segmented_array::create_from_sizes(
+        array{exec, {1, 2, 3, 4, 5}}, iarray{exec, {3, 2}});
+    segmented_array arr5 = segmented_array::create_from_sizes(
+        array{exec, {1, 2, 3, 4, 5}}, iarray{exec, {1, 2, 2}});
+};
+
+
+TEST_F(SegmentedArraysEqual, SucceedsIfEqual)
+{
+    GKO_ASSERT_SEGMENTED_ARRAY_EQ(arr1, arr2);
+}
+
+
+TEST_F(SegmentedArraysEqual, FailsIfValuesDifferent)
+{
+    ASSERT_PRED_FORMAT2(!::gko::test::assertions::segmented_array_equal, arr1,
+                        arr3);
+}
+
+
+TEST_F(SegmentedArraysEqual, FailsIfOffsetsDifferent1)
+{
+    ASSERT_PRED_FORMAT2(!::gko::test::assertions::segmented_array_equal, arr1,
+                        arr4);
+}
+
+
+TEST_F(SegmentedArraysEqual, FailsIfOffsetsDifferent2)
+{
+    ASSERT_PRED_FORMAT2(!::gko::test::assertions::segmented_array_equal, arr1,
+                        arr5);
+}
+
+
 }  // namespace

dpcpp/distributed/index_map_kernels.dp.cpp

@@ -50,9 +50,9 @@ void map_to_global(
     device_partition<const LocalIndexType, const GlobalIndexType> partition,
     device_segmented_array<const GlobalIndexType> remote_global_idxs,
     experimental::distributed::comm_index_type rank,
-    const array<LocalIndexType>& local_ids,
+    const array<LocalIndexType>& local_idxs,
     experimental::distributed::index_space is,
-    array<GlobalIndexType>& global_ids) GKO_NOT_IMPLEMENTED;
+    array<GlobalIndexType>& global_idxs) GKO_NOT_IMPLEMENTED;
 
 GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
     GKO_DECLARE_INDEX_MAP_MAP_TO_GLOBAL);

dpcpp/distributed/partition_kernels.dp.cpp

@@ -140,10 +140,7 @@ void build_ranges_by_part(std::shared_ptr<const DefaultExecutor> exec,
 
     range_ids.resize_and_reset(num_ranges);
     auto range_ids_ptr = range_ids.get_data();
-    // fill range_ids with 0,...,num_ranges - 1
-    run_kernel(
-        exec, [] GKO_KERNEL(auto i, auto rid) { rid[i] = i; }, num_ranges,
-        range_ids_ptr);
+    components::fill_seq_array(exec, range_ids_ptr, num_ranges);
 
     oneapi::dpl::stable_sort(policy, range_ids_ptr, range_ids_ptr + num_ranges,
                              [range_parts](const auto rid_a, const auto rid_b) {

include/ginkgo/core/distributed/index_map.hpp

@@ -85,15 +85,16 @@ struct index_map {
     /**
      * Maps local indices to global indices
      *
-     * @param local_ids  the local indices to map
+     * @param local_idxs  the local indices to map
      * @param index_space_v  the index space in which the passed-in local
      *                       indices are defined
      *
      * @return  the mapped global indices. Any local index, that is not in the
      *          specified index space is mapped to invalid_index
      */
-    array<GlobalIndexType> map_to_global(const array<LocalIndexType>& local_ids,
-                                         index_space index_space_v) const;
+    array<GlobalIndexType> map_to_global(
+        const array<LocalIndexType>& local_idxs,
+        index_space index_space_v) const;
 
     /**
      * \brief get size of index_space::local