add partition kernels

common/CMakeLists.txt

@@ -1,6 +1,7 @@
 set(UNIFIED_SOURCES
     components/precision_conversion.cpp
     components/reduce_array.cpp
+    distributed/partition_kernels.cpp
     matrix/coo_kernels.cpp
     matrix/csr_kernels.cpp
     matrix/dense_kernels.cpp

common/unified/distributed/partition_kernels.cpp

@@ -0,0 +1,128 @@
+/*******************************<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.hpp"
+
+
+namespace gko {
+namespace kernels {
+namespace GKO_DEVICE_NAMESPACE {
+namespace partition {
+
+
+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 LocalIndexType>
+void build_from_contiguous(std::shared_ptr<const DefaultExecutor> exec,
+                           const Array<global_index_type>& ranges,
+                           distributed::Partition<LocalIndexType>* partition)
+{
+    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, partition->get_range_bounds(),
+        partition->get_part_ids());
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_BUILD_FROM_CONTIGUOUS);
+
+
+template <typename LocalIndexType>
+void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
+                        const Array<comm_index_type>& mapping,
+                        distributed::Partition<LocalIndexType>* partition)
+{
+    Array<size_type> range_index_ranks{exec, mapping.get_num_elems() + 1};
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto mapping, auto output) {
+            const auto prev_part = i > 0 ? mapping[i - 1] : comm_index_type{-1};
+            const auto cur_part = mapping[i];
+            output[i] = cur_part != prev_part ? 1 : 0;
+        },
+        mapping.get_num_elems(), mapping, range_index_ranks);
+    components::prefix_sum(exec, range_index_ranks.get_data(),
+                           mapping.get_num_elems() + 1);
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto size, auto mapping, auto prefix_sum,
+                      auto ranges, auto range_parts) {
+            const auto prev_part = i > 0 ? mapping[i - 1] : comm_index_type{-1};
+            const auto cur_part = i < size ? mapping[i] : comm_index_type{-1};
+            if (cur_part != prev_part) {
+                auto out_idx = prefix_sum[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_index_ranks, partition->get_range_bounds(),
+        partition->get_part_ids());
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING);
+
+
+}  // namespace partition
+}  // namespace GKO_DEVICE_NAMESPACE
+}  // namespace kernels
+}  // namespace gko

core/distributed/partition.cpp

@@ -103,6 +103,7 @@ template <typename LocalIndexType>
 void Partition<LocalIndexType>::compute_block_gather_permutation(
     const bool recompute)
 {
+    return;
     if (block_gather_permutation_.get_num_elems() == 0 || recompute) {
         block_gather_permutation_.resize_and_reset(this->get_size());
         block_gather_permutation_.fill(-1);

cuda/distributed/partition_kernels.cu

@@ -33,41 +33,84 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "core/distributed/partition_kernels.hpp"
 
 
-namespace gko {
-namespace kernels {
-namespace cuda {
-namespace partition {
+#include <thrust/device_ptr.h>
+#include <thrust/execution_policy.h>
+#include <thrust/iterator/zip_iterator.h>
+#include <thrust/scan.h>
+#include <thrust/sort.h>
 
 
-void count_ranges(std::shared_ptr<const DefaultExecutor> exec,
-                  const Array<comm_index_type>& mapping,
-                  size_type& num_ranges) GKO_NOT_IMPLEMENTED;
+#include "common/unified/base/kernel_launch.hpp"
+#include "core/components/fill_array.hpp"
+#include "core/components/prefix_sum.hpp"
 
 
-template <typename LocalIndexType>
-void build_from_contiguous(std::shared_ptr<const DefaultExecutor> exec,
-                           const Array<global_index_type>& ranges,
-                           distributed::Partition<LocalIndexType>* partition)
-    GKO_NOT_IMPLEMENTED;
-
-GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
-    GKO_DECLARE_PARTITION_BUILD_FROM_CONTIGUOUS);
-
-
-template <typename LocalIndexType>
-void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
-                        const Array<comm_index_type>& mapping,
-                        distributed::Partition<LocalIndexType>* partition)
-    GKO_NOT_IMPLEMENTED;
-
-GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING);
+namespace gko {
+namespace kernels {
+namespace cuda {
+namespace partition {
 
 
 template <typename LocalIndexType>
 void build_ranks(std::shared_ptr<const DefaultExecutor> exec,
                  const global_index_type* range_offsets, const int* range_parts,
                  size_type num_ranges, int num_parts, LocalIndexType* ranks,
-                 LocalIndexType* sizes) GKO_NOT_IMPLEMENTED;
+                 LocalIndexType* sizes)
+{
+    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<size_type> permutation{exec, num_ranges};
+    // set sizes to 0 in case of empty parts
+    components::fill_array(exec, sizes, num_parts, LocalIndexType{});
+
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto num_ranges, auto num_parts,
+                      auto range_offsets, auto range_parts, auto range_sizes,
+                      auto tmp_part_ids, auto permutation) {
+            if (i == 0) {
+                // set sentinel value at the end
+                tmp_part_ids[num_ranges] = num_parts;
+            }
+            range_sizes[i] = range_offsets[i + 1] - range_offsets[i];
+            tmp_part_ids[i] = range_parts[i];
+            permutation[i] = static_cast<int64>(i);
+        },
+        num_ranges, 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 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
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto grouped_range_ranks, auto grouped_part_ids,
+                      auto orig_idxs, auto ranks, auto sizes) {
+            auto prev_part =
+                i > 0 ? grouped_part_ids[i - 1] : comm_index_type{-1};
+            auto cur_part = grouped_part_ids[i];
+            auto next_part = grouped_part_ids[i + 1];  // safe due to sentinel
+            if (cur_part != next_part) {
+                sizes[cur_part] = grouped_range_ranks[i];
+            }
+            // write result shifted by one entry to get exclusive prefix sum
+            ranks[orig_idxs[i]] = prev_part == cur_part
+                                      ? grouped_range_ranks[i - 1]
+                                      : LocalIndexType{};
+        },
+        num_ranges, range_sizes, tmp_part_ids, permutation, ranks, sizes);
+}
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_RANKS);
 

dpcpp/distributed/partition_kernels.dp.cpp

@@ -39,35 +39,62 @@ namespace dpcpp {
 namespace partition {
 
 
-void count_ranges(std::shared_ptr<const DefaultExecutor> exec,
-                  const Array<comm_index_type>& mapping,
-                  size_type& num_ranges) GKO_NOT_IMPLEMENTED;
-
-
-template <typename LocalIndexType>
-void build_from_contiguous(std::shared_ptr<const DefaultExecutor> exec,
-                           const Array<global_index_type>& ranges,
-                           distributed::Partition<LocalIndexType>& partition)
-    GKO_NOT_IMPLEMENTED;
-
-GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
-    GKO_DECLARE_PARTITION_BUILD_FROM_CONTIGUOUS);
-
-
-template <typename LocalIndexType>
-void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
-                        const Array<comm_index_type>& mapping,
-                        distributed::Partition<LocalIndexType>& partition)
-    GKO_NOT_IMPLEMENTED;
-
-GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING);
-
-
 template <typename LocalIndexType>
 void build_ranks(std::shared_ptr<const DefaultExecutor> exec,
                  const global_index_type* range_offsets, const int* range_parts,
                  size_type num_ranges, int num_parts, LocalIndexType* ranks,
-                 LocalIndexType* sizes) GKO_NOT_IMPLEMENTED;
+                 LocalIndexType* sizes)
+{
+    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<size_type> permutation{exec, num_ranges};
+    // set sizes to 0 in case of empty parts
+    components::fill_array(exec, sizes, num_parts, LocalIndexType{});
+
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto num_ranges, auto num_parts,
+                      auto range_offsets, auto range_parts, auto range_sizes,
+                      auto tmp_part_ids, auto permutation) {
+            if (i == 0) {
+                // set sentinel value at the end
+                tmp_part_ids[num_ranges] = num_parts;
+            }
+            range_sizes[i] = range_offsets[i + 1] - range_offsets[i];
+            tmp_part_ids[i] = range_parts[i];
+            permutation[i] = static_cast<int64>(i);
+        },
+        num_ranges, num_ranges, num_parts, range_offsets, range_parts,
+        range_sizes, tmp_part_ids, permutation);
+
+    // group sizes by part ID
+    // TODO oneDPL has stable_sort and views::zip
+    // compute inclusive prefix sum for each part
+    // TODO compute "row_ptrs" for tmp_part_ids
+    // TODO compute prefix_sum over range_sizes
+    // TODO compute adjacent differences, set -part_size at part boundaries
+    // TODO compute prefix_sum again
+    // write back the results
+    // TODO this needs to be adapted to the output of the algorithm above
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto i, auto grouped_range_ranks, auto grouped_part_ids,
+                      auto orig_idxs, auto ranks, auto sizes) {
+            auto prev_part =
+                i > 0 ? grouped_part_ids[i - 1] : comm_index_type{-1};
+            auto cur_part = grouped_part_ids[i];
+            auto next_part = grouped_part_ids[i + 1];  // safe due to sentinel
+            if (cur_part != next_part) {
+                sizes[cur_part] = grouped_range_ranks[i];
+            }
+            // write result shifted by one entry to get exclusive prefix sum
+            ranks[orig_idxs[i]] = prev_part == cur_part
+                                      ? grouped_range_ranks[i - 1]
+                                      : LocalIndexType{};
+        },
+        num_ranges, range_sizes, tmp_part_ids, permutation, ranks, sizes);
+}
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_RANKS);