Merge building partition from local sizes/ranges

This merge adds freestanding functions to create a Partition from local sizes, or local ranges. This has to be a freestanding function, because adding it to Partition would require including MPI, which leads to our device libraries requiring MPI (or at least the headers).

common/cuda_hip/distributed/partition_helpers_kernels.hpp.inc

@@ -0,0 +1,54 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2023, the Ginkgo authors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+******************************<GINKGO LICENSE>*******************************/
+
+template <typename GlobalIndexType>
+void sort_by_range_start(
+    std::shared_ptr<const DefaultExecutor> exec,
+    array<GlobalIndexType>& range_start_ends,
+    array<experimental::distributed::comm_index_type>& part_ids)
+{
+    auto num_ranges = range_start_ends.get_num_elems() / 2;
+    auto strided_indices = thrust::make_transform_iterator(
+        thrust::make_counting_iterator(0),
+        [] __host__ __device__(const int i) { return 2 * i; });
+    auto start_it = thrust::make_permutation_iterator(
+        range_start_ends.get_data(), strided_indices);
+    auto end_it = thrust::make_permutation_iterator(
+        range_start_ends.get_data() + 1, strided_indices);
+    auto zip_it = thrust::make_zip_iterator(
+        thrust::make_tuple(end_it, part_ids.get_data()));
+    thrust::stable_sort_by_key(thrust_policy(exec), start_it,
+                               start_it + num_ranges, zip_it);
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_HELPERS_SORT_BY_RANGE_START);

common/unified/CMakeLists.txt

@@ -6,6 +6,7 @@ set(UNIFIED_SOURCES
     components/format_conversion_kernels.cpp
     components/precision_conversion_kernels.cpp
     components/reduce_array_kernels.cpp
+    distributed/partition_helpers_kernels.cpp
     distributed/partition_kernels.cpp
     matrix/coo_kernels.cpp
     matrix/csr_kernels.cpp

common/unified/distributed/partition_helpers_kernels.cpp

@@ -0,0 +1,102 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2023, the Ginkgo authors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+******************************<GINKGO LICENSE>*******************************/
+
+#include "core/distributed/partition_helpers_kernels.hpp"
+
+
+#include "common/unified/base/kernel_launch.hpp"
+#include "common/unified/base/kernel_launch_reduction.hpp"
+
+
+namespace gko {
+namespace kernels {
+namespace GKO_DEVICE_NAMESPACE {
+namespace partition_helpers {
+
+
+template <typename GlobalIndexType>
+void check_consecutive_ranges(std::shared_ptr<const DefaultExecutor> exec,
+                              const array<GlobalIndexType>& range_start_ends,
+                              bool& result)
+{
+    array<uint32> result_uint32{exec, 1};
+    auto num_ranges = range_start_ends.get_num_elems() / 2;
+    // need additional guard because DPCPP doesn't return the initial value for
+    // empty inputs
+    if (num_ranges > 1) {
+        run_kernel_reduction(
+            exec,
+            [] GKO_KERNEL(const auto i, const auto* ranges) {
+                return ranges[2 * i] == ranges[2 * i + 1];
+            },
+            [] GKO_KERNEL(const auto a, const auto b) {
+                return static_cast<uint32>(a && b);
+            },
+            [] GKO_KERNEL(auto x) { return x; }, static_cast<uint32>(true),
+            result_uint32.get_data(), num_ranges - 1,
+            range_start_ends.get_const_data() + 1);
+        result =
+            static_cast<bool>(exec->copy_val_to_host(result_uint32.get_data()));
+    } else {
+        result = true;
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_HELPERS_CHECK_CONSECUTIVE_RANGES);
+
+
+template <typename GlobalIndexType>
+void compress_ranges(std::shared_ptr<const DefaultExecutor> exec,
+                     const array<GlobalIndexType>& range_start_ends,
+                     array<GlobalIndexType>& range_offsets)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(const auto i, const auto* start_ends, auto* offsets) {
+            if (i == 0) {
+                offsets[0] = start_ends[0];
+            }
+            offsets[i + 1] = start_ends[2 * i + 1];
+        },
+        range_offsets.get_num_elems() - 1, range_start_ends.get_const_data(),
+        range_offsets.get_data());
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_HELPERS_COMPRESS_RANGES);
+
+
+}  // namespace partition_helpers
+}  // namespace GKO_DEVICE_NAMESPACE
+}  // namespace kernels
+}  // namespace gko

common/unified/distributed/partition_kernels.cpp

@@ -66,19 +66,22 @@ void count_ranges(std::shared_ptr<const DefaultExecutor> exec,
 template <typename GlobalIndexType>
 void build_from_contiguous(std::shared_ptr<const DefaultExecutor> exec,
                            const array<GlobalIndexType>& ranges,
+                           const array<comm_index_type>& part_id_mapping,
                            GlobalIndexType* range_bounds,
                            comm_index_type* part_ids)
 {
     run_kernel(
         exec,
-        [] GKO_KERNEL(auto i, auto ranges, auto bounds, auto ids) {
+        [] GKO_KERNEL(auto i, auto ranges, auto mapping, auto bounds, auto ids,
+                      bool uses_mapping) {
             if (i == 0) {
                 bounds[0] = 0;
             }
             bounds[i + 1] = ranges[i + 1];
-            ids[i] = i;
+            ids[i] = uses_mapping ? mapping[i] : i;
         },
-        ranges.get_num_elems() - 1, ranges, range_bounds, part_ids);
+        ranges.get_num_elems() - 1, ranges, part_id_mapping, range_bounds,
+        part_ids, part_id_mapping.get_num_elems() > 0);
 }
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_PARTITION_BUILD_FROM_CONTIGUOUS);

core/CMakeLists.txt

@@ -95,6 +95,7 @@ if(GINKGO_BUILD_MPI)
         PRIVATE
         mpi/exception.cpp
         distributed/matrix.cpp
+        distributed/partition_helpers.cpp
         distributed/vector.cpp
         distributed/preconditioner/schwarz.cpp)
 endif()

core/base/copy_assignable.hpp

@@ -0,0 +1,130 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2023, the Ginkgo authors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+******************************<GINKGO LICENSE>*******************************/
+
+#ifndef GKO_CORE_BASE_COPY_ASSIGNABLE_HPP_
+#define GKO_CORE_BASE_COPY_ASSIGNABLE_HPP_
+
+
+#include <vector>
+
+
+namespace gko {
+namespace detail {
+
+
+template <typename T, typename = void>
+class copy_assignable;
+
+
+/**
+ * Helper class to make a type copy assignable.
+ *
+ * This class wraps an object of a type that has a copy constructor, but not
+ * a copy assignment. This is most often the case for lambdas. The wrapped
+ * object can then be copy assigned, by relying on the copy constructor.
+ *
+ * @tparam T  type with a copy constructor
+ */
+template <typename T>
+class copy_assignable<
+    T, typename std::enable_if<std::is_copy_constructible<T>::value>::type> {
+public:
+    copy_assignable() = default;
+
+    copy_assignable(const copy_assignable& other)
+    {
+        if (this != &other) {
+            *this = other;
+        }
+    }
+
+    copy_assignable(copy_assignable&& other) noexcept
+    {
+        if (this != &other) {
+            *this = std::move(other);
+        }
+    }
+
+    copy_assignable(const T& obj) : obj_{new (buf)(T)(obj)} {}
+
+    copy_assignable(T&& obj) : obj_{new (buf)(T)(std::move(obj))} {}
+
+    copy_assignable& operator=(const copy_assignable& other)
+    {
+        if (this != &other) {
+            if (obj_) {
+                obj_->~T();
+            }
+            obj_ = new (buf)(T)(*other.obj_);
+        }
+        return *this;
+    }
+
+    copy_assignable& operator=(copy_assignable&& other) noexcept
+    {
+        if (this != &other) {
+            if (obj_) {
+                obj_->~T();
+            }
+            obj_ = new (buf)(T)(std::move(*other.obj_));
+        }
+        return *this;
+    }
+
+    ~copy_assignable()
+    {
+        if (obj_) {
+            obj_->~T();
+        }
+    }
+
+    template <typename... Args>
+    decltype(auto) operator()(Args&&... args) const
+    {
+        return (*obj_)(std::forward<Args>(args)...);
+    }
+
+    T const& get() const { return *obj_; }
+
+    T& get() { return *obj_; }
+
+private:
+    //!< Store wrapped object on the stack, should use std::optional in c++17
+    T* obj_{};
+    alignas(T) unsigned char buf[sizeof(T)];
+};
+
+
+}  // namespace detail
+}  // namespace gko
+
+#endif  // GKO_CORE_BASE_COPY_ASSIGNABLE_HPP_