review updates

- renaming
- unifying kernels for computing local starting indices
- documentation

Co-authored-by: Aditya Kashi <[email protected]>
Co-authored-by: Yu-Hsiang Tsai <[email protected]>

common/cuda_hip/distributed/partition_kernels.hpp.inc

@@ -0,0 +1,79 @@
+/*******************************<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 "common/unified/distributed/partition_kernels.hpp.inc"
+
+
+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)
+{
+ 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);
+}
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
+ GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES);

common/unified/distributed/partition_kernels.cpp

@@ -44,6 +44,8 @@ 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)
 {
@@ -89,33 +91,37 @@ void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
  GlobalIndexType* range_bounds,
  comm_index_type* part_ids)
 {
- Array<size_type> range_index_ranks{exec, mapping.get_num_elems() + 1};
+ Array<size_type> range_starting_index{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};
+ [] 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];
- output[i] = cur_part != prev_part ? 1 : 0;
+ range_starting_index[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(), 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 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};
+ [] 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 = prefix_sum[i];
+ 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_index_ranks, range_bounds, part_ids);
+ range_starting_index, range_bounds, part_ids);
 }
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_PARTITION_BUILD_FROM_MAPPING);
@@ -157,11 +163,13 @@ void has_ordered_parts(
  [] 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 a && b; },
+ [] 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->template copy_val_to_host(result_uint32.get_const_data()));
+ exec->copy_val_to_host(result_uint32.get_const_data()));
 }
 
 GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(

common/unified/distributed/partition_kernels.hpp.inc

@@ -0,0 +1,92 @@
+/*******************************<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

core/distributed/partition.cpp

@@ -67,7 +67,7 @@ Partition<LocalIndexType, GlobalIndexType>::build_from_mapping(
  exec->run(partition::make_build_from_mapping(*local_mapping.get(),
  result->offsets_.get_data(),
  result->part_ids_.get_data()));
- result->compute_range_starting_indices();
+ result->finalize_construction();
  return result;
 }
 
@@ -84,7 +84,7 @@ Partition<LocalIndexType, GlobalIndexType>::build_from_contiguous(
  exec->run(partition::make_build_from_contiguous(
  *local_ranges.get(), result->offsets_.get_data(),
  result->part_ids_.get_data()));
- result->compute_range_starting_indices();
+ result->finalize_construction();
  return result;
 }
 
@@ -103,28 +103,30 @@ Partition<LocalIndexType, GlobalIndexType>::build_from_global_size_uniform(
 
 
 template <typename LocalIndexType, typename GlobalIndexType>
-void Partition<LocalIndexType,
- GlobalIndexType>::compute_range_starting_indices()
+void Partition<LocalIndexType, GlobalIndexType>::finalize_construction()
 {
  auto exec = offsets_.get_executor();
  exec->run(partition::make_build_starting_indices(
  offsets_.get_const_data(), part_ids_.get_const_data(), get_num_ranges(),
  get_num_parts(), num_empty_parts_, starting_indices_.get_data(),
  part_sizes_.get_data()));
+ size_ = offsets_.get_executor()->copy_val_to_host(
+ offsets_.get_const_data() + get_num_ranges());
 }
 
 
 template <typename LocalIndexType, typename GlobalIndexType>
 bool Partition<LocalIndexType, GlobalIndexType>::has_connected_parts()
 {
- return get_num_parts() - get_num_empty_parts() == get_num_ranges();
+ return this->get_num_parts() - this->get_num_empty_parts() ==
+ this->get_num_ranges();
 }
 
 
 template <typename LocalIndexType, typename GlobalIndexType>
 bool Partition<LocalIndexType, GlobalIndexType>::has_ordered_parts()
 {
- if (has_connected_parts()) {
+ if (this->has_connected_parts()) {
  auto exec = this->get_executor();
  bool has_ordered_parts;
  exec->run(partition::make_has_ordered_parts(this, &has_ordered_parts));

cuda/distributed/partition_kernels.cu

@@ -41,8 +41,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
 #include "common/unified/base/kernel_launch.hpp"
-#include "core/components/fill_array.hpp"
-#include "core/components/prefix_sum.hpp"
+#include "core/components/fill_array_kernels.hpp"
 
 
 namespace gko {
@@ -51,72 +50,7 @@ namespace cuda {
 namespace partition {
 
 
-template <typename LocalIndexType, typename GlobalIndexType>
-void build_starting_indices(std::shared_ptr<const DefaultExecutor> exec,
- const GlobalIndexType* range_offsets,
- const int* range_parts, size_type num_ranges,
- int num_parts, int& num_empty_parts,
- LocalIndexType* ranks, 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<GlobalIndexType> 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<GlobalIndexType>(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);
- num_empty_parts =
- thrust::count(thrust::device, sizes, sizes + num_parts, 0);
-}
-
-GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
- GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES);
+#include "common/cuda_hip/distributed/partition_kernels.hpp.inc"
 
 
 } // namespace partition