Merge: Add a class for partitions of intervals

This PR adds a class to handle a partition of an interval [0,N) into a number of disjoint sub ranges. Each range is associated with a not necessarily unique part ID. 

Related PR: #909

cmake/create_test.cmake

@@ -163,6 +163,7 @@ ginkgo_build_test_name(${test_name} test_target_name)
 endfunction(ginkgo_create_hip_test)
 
 function(ginkgo_create_common_test test_name)
+    cmake_parse_arguments(PARSE_ARGV 1 common_test "" "" "DISABLE_EXECUTORS;ADDITIONAL_LIBRARIES")
     set(executors)
     if(GINKGO_BUILD_OMP)
         list(APPEND executors omp)
@@ -176,6 +177,9 @@ function(ginkgo_create_common_test test_name)
     if(GINKGO_BUILD_DPCPP)
         list(APPEND executors dpcpp)
     endif()
+    foreach(disabled_exec ${common_test_DISABLE_EXECUTORS})
+        list(REMOVE_ITEM executors ${disabled_exec})
+    endforeach()
     foreach(exec ${executors})
         ginkgo_build_test_name(${test_name} test_target_name)
         # build executor typename out of shorthand
@@ -189,7 +193,7 @@ function(ginkgo_create_common_test test_name)
         target_compile_features(${test_target_name} PUBLIC cxx_std_14)
         target_compile_options(${test_target_name} PRIVATE ${GINKGO_COMPILER_FLAGS})
         target_compile_definitions(${test_target_name} PRIVATE EXEC_TYPE=${exec_type} EXEC_NAMESPACE=${exec})
-        target_link_libraries(${test_target_name} PRIVATE ${ARGN})
+        target_link_libraries(${test_target_name} PRIVATE ${common_test_ADDITIONAL_LIBRARIES})
         # use float for DPC++ if necessary
         if((exec STREQUAL "dpcpp") AND GINKGO_DPCPP_SINGLE_MODE)
             target_compile_definitions(${test_target_name} PRIVATE GINKGO_COMMON_SINGLE_MODE=1)
@@ -209,4 +213,4 @@ function(ginkgo_create_common_and_reference_test test_name)
     target_compile_definitions(${test_target_name} PRIVATE EXEC_TYPE=ReferenceExecutor EXEC_NAMESPACE=reference)
     target_link_libraries(${test_target_name} PRIVATE ${ARGN})
     ginkgo_set_test_target_properties(${test_name}_reference ${test_target_name})
-endfunction()
+endfunction()

common/CMakeLists.txt

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

common/cuda_hip/distributed/partition_kernels.hpp.inc

@@ -0,0 +1,144 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2021, 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>*******************************/
+
+
+namespace kernel {
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void setup_sizes_ids_permutation(
+    std::shared_ptr<const DefaultExecutor> exec, size_type num_ranges,
+    comm_index_type num_parts, const GlobalIndexType* range_offsets,
+    const comm_index_type* range_parts, Array<LocalIndexType>& range_sizes,
+    Array<comm_index_type>& part_ids, Array<GlobalIndexType>& permutation)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto num_ranges, auto num_parts,
+                      auto range_offsets, auto range_parts, auto range_sizes,
+                      auto part_ids, auto permutation) {
+            if (i == 0) {
+                // set sentinel value at the end
+                part_ids[num_ranges] = num_parts;
+            }
+            range_sizes[i] = range_offsets[i + 1] - range_offsets[i];
+            part_ids[i] = range_parts[i];
+            permutation[i] = static_cast<GlobalIndexType>(i);
+        },
+        num_ranges, num_ranges, num_parts, range_offsets, range_parts,
+        range_sizes.get_data(), part_ids.get_data(), permutation.get_data());
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void compute_part_sizes_and_starting_indices(
+    std::shared_ptr<const DefaultExecutor> exec, size_type num_ranges,
+    const Array<LocalIndexType>& range_sizes,
+    const Array<comm_index_type>& part_ids,
+    const Array<GlobalIndexType>& permutation, LocalIndexType* starting_indices,
+    LocalIndexType* part_sizes)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto grouped_starting_indices,
+                      auto grouped_part_ids, auto orig_idxs,
+                      auto starting_indices, auto part_sizes) {
+            auto prev_part = i > 0 ? grouped_part_ids[i - 1]
+                                   : invalid_index<comm_index_type>();
+            auto cur_part = grouped_part_ids[i];
+            auto next_part =
+                grouped_part_ids[i + 1];  // last element has to be num_parts
+            if (cur_part != next_part) {
+                part_sizes[cur_part] = grouped_starting_indices[i];
+            }
+            // write result shifted by one entry to get exclusive prefix sum
+            starting_indices[orig_idxs[i]] =
+                prev_part == cur_part ? grouped_starting_indices[i - 1]
+                                      : LocalIndexType{};
+        },
+        num_ranges, range_sizes.get_const_data(), part_ids.get_const_data(),
+        permutation.get_const_data(), starting_indices, part_sizes);
+}
+
+
+}  // namespace kernel
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void build_starting_indices(std::shared_ptr<const DefaultExecutor> exec,
+                            const GlobalIndexType* range_offsets,
+                            const comm_index_type* range_parts,
+                            size_type num_ranges, comm_index_type num_parts,
+                            comm_index_type& num_empty_parts,
+                            LocalIndexType* starting_indices,
+                            LocalIndexType* part_sizes)
+{
+    if (num_ranges > 0) {
+        Array<LocalIndexType> range_sizes{exec, num_ranges};
+        // num_parts sentinel at the end
+        Array<comm_index_type> tmp_part_ids{exec, num_ranges + 1};
+        Array<GlobalIndexType> permutation{exec, num_ranges};
+        // set part_sizes to 0 in case of empty parts
+        components::fill_array(exec, part_sizes, num_parts, LocalIndexType{});
+
+        kernel::setup_sizes_ids_permutation(
+            exec, num_ranges, num_parts, range_offsets, range_parts,
+            range_sizes, tmp_part_ids, permutation);
+
+        auto tmp_part_id_ptr =
+            thrust::device_pointer_cast(tmp_part_ids.get_data());
+        auto range_sizes_ptr =
+            thrust::device_pointer_cast(range_sizes.get_data());
+        auto permutation_ptr =
+            thrust::device_pointer_cast(permutation.get_data());
+        auto value_it = thrust::make_zip_iterator(
+            thrust::make_tuple(range_sizes_ptr, permutation_ptr));
+        // group range_sizes by part ID
+        thrust::stable_sort_by_key(thrust::device, tmp_part_id_ptr,
+                                   tmp_part_id_ptr + num_ranges, value_it);
+        // compute inclusive prefix sum for each part
+        thrust::inclusive_scan_by_key(thrust::device, tmp_part_id_ptr,
+                                      tmp_part_id_ptr + num_ranges,
+                                      range_sizes_ptr, range_sizes_ptr);
+        // write back the results
+        kernel::compute_part_sizes_and_starting_indices(
+            exec, num_ranges, range_sizes, tmp_part_ids, permutation,
+            starting_indices, part_sizes);
+        num_empty_parts = thrust::count(thrust::device, part_sizes,
+                                        part_sizes + num_parts, 0);
+    } else {
+        num_empty_parts = num_parts;
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES);

common/unified/distributed/partition_kernels.cpp

@@ -0,0 +1,182 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2021, 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_kernels.hpp"
+
+
+#include "common/unified/base/kernel_launch.hpp"
+#include "common/unified/base/kernel_launch_reduction.hpp"
+#include "core/components/prefix_sum_kernels.hpp"
+
+
+namespace gko {
+namespace kernels {
+namespace GKO_DEVICE_NAMESPACE {
+namespace partition {
+
+
+using distributed::comm_index_type;
+
+void count_ranges(std::shared_ptr<const DefaultExecutor> exec,
+                  const Array<comm_index_type>& mapping, size_type& num_ranges)
+{
+    Array<size_type> result{exec, 1};
+    run_kernel_reduction(
+        exec,
+        [] GKO_KERNEL(auto i, auto mapping) {
+            auto cur_part = mapping[i];
+            auto prev_part = i == 0 ? comm_index_type{-1} : mapping[i - 1];
+            return cur_part != prev_part ? 1 : 0;
+        },
+        [] GKO_KERNEL(auto a, auto b) { return a + b; },
+        [] GKO_KERNEL(auto a) { return a; }, size_type{}, result.get_data(),
+        mapping.get_num_elems(), mapping);
+    num_ranges = exec->copy_val_to_host(result.get_const_data());
+}
+
+
+template <typename GlobalIndexType>
+void build_from_contiguous(std::shared_ptr<const DefaultExecutor> exec,
+                           const Array<GlobalIndexType>& ranges,
+                           GlobalIndexType* range_bounds,
+                           comm_index_type* part_ids)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto ranges, auto bounds, auto ids) {
+            if (i == 0) {
+                bounds[0] = 0;
+            }
+            bounds[i + 1] = ranges[i + 1];
+            ids[i] = i;
+        },
+        ranges.get_num_elems() - 1, ranges, range_bounds, part_ids);
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_PARTITION_BUILD_FROM_CONTIGUOUS);
+
+
+template <typename GlobalIndexType>
+void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
+                        const Array<comm_index_type>& mapping,
+                        GlobalIndexType* range_bounds,
+                        comm_index_type* part_ids)
+{
+    Array<size_type> range_starting_index{exec, mapping.get_num_elems() + 1};
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto mapping, auto range_starting_index) {
+            const auto prev_part =
+                i > 0 ? mapping[i - 1] : invalid_index<comm_index_type>();
+            const auto cur_part = mapping[i];
+            range_starting_index[i] = cur_part != prev_part ? 1 : 0;
+        },
+        mapping.get_num_elems(), mapping, range_starting_index);
+    components::prefix_sum(exec, range_starting_index.get_data(),
+                           mapping.get_num_elems() + 1);
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto size, auto mapping,
+                      auto range_starting_index, auto ranges,
+                      auto range_parts) {
+            const auto prev_part =
+                i > 0 ? mapping[i - 1] : invalid_index<comm_index_type>();
+            const auto cur_part =
+                i < size ? mapping[i] : invalid_index<comm_index_type>();
+            if (cur_part != prev_part) {
+                auto out_idx = range_starting_index[i];
+                ranges[out_idx] = i;
+                if (i < size) {
+                    range_parts[out_idx] = cur_part;
+                }
+            }
+        },
+        mapping.get_num_elems() + 1, mapping.get_num_elems(), mapping,
+        range_starting_index, range_bounds, part_ids);
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_PARTITION_BUILD_FROM_MAPPING);
+
+
+template <typename GlobalIndexType>
+void build_ranges_from_global_size(std::shared_ptr<const DefaultExecutor> exec,
+                                   comm_index_type num_parts,
+                                   GlobalIndexType global_size,
+                                   Array<GlobalIndexType>& ranges)
+{
+    const auto size_per_part = global_size / num_parts;
+    const auto rest = global_size - (num_parts * size_per_part);
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto size_per_part, auto rest, auto ranges) {
+            ranges[i] = size_per_part + (i < rest ? 1 : 0);
+        },
+        ranges.get_num_elems() - 1, size_per_part, rest, ranges.get_data());
+    components::prefix_sum(exec, ranges.get_data(), ranges.get_num_elems());
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_PARTITION_BUILD_FROM_GLOBAL_SIZE);
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void has_ordered_parts(
+    std::shared_ptr<const DefaultExecutor> exec,
+    const distributed::Partition<LocalIndexType, GlobalIndexType>* partition,
+    bool* result)
+{
+    const auto part_ids = partition->get_part_ids();
+    const auto num_ranges = partition->get_num_ranges();
+    // it is necessary to use uint32 as a temporary result, since
+    // bool can't be used with suffles
+    Array<uint32> result_uint32{exec, 1};
+    run_kernel_reduction(
+        exec,
+        [] GKO_KERNEL(auto i, const auto part_ids) {
+            return static_cast<uint32>(part_ids[i] < part_ids[i + 1]);
+        },
+        [] GKO_KERNEL(const auto a, const auto b) {
+            return static_cast<uint32>(a && b);
+        },
+        [] GKO_KERNEL(const auto a) { return a; }, uint32(1),
+        result_uint32.get_data(), num_ranges - 1, part_ids);
+    *result = static_cast<bool>(
+        exec->copy_val_to_host(result_uint32.get_const_data()));
+}
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_IS_ORDERED);
+
+
+}  // namespace partition
+}  // namespace GKO_DEVICE_NAMESPACE
+}  // namespace kernels
+}  // namespace gko

core/CMakeLists.txt

@@ -12,6 +12,7 @@ target_sources(ginkgo
     base/mtx_io.cpp
     base/perturbation.cpp
     base/version.cpp
+    distributed/partition.cpp
     factorization/ic.cpp
     factorization/ilu.cpp
     factorization/par_ic.cpp