review updates

- merge /common/unified/.../partitions_kernels.hpp.inc into /commun/cuda_hip/.../partitions_kernsls.hpp.inc
- documentation

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

common/cuda_hip/distributed/partition_kernels.hpp.inc

@@ -31,7 +31,66 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************<GINKGO LICENSE>*******************************/
 
 
-#include "common/unified/distributed/partition_kernels.hpp.inc"
+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>

dpcpp/distributed/partition_kernels.dp.cpp

@@ -39,6 +39,8 @@ namespace dpcpp {
 namespace partition {
 
 
+// TODO: wait until https://github.com/oneapi-src/oneDPL/pull/388 is release to
+// implement it similar to cuda/hip
 template <typename LocalIndexType, typename GlobalIndexType>
 void build_starting_indices(std::shared_ptr<const DefaultExecutor> exec,
  const GlobalIndexType* range_offsets,

include/ginkgo/core/distributed/partition.hpp

@@ -251,9 +251,10 @@ class Partition
  *
  * @param exec the Executor on which the partition should be built
  * @param ranges the boundaries of the ranges representing each part.
- Part i contains the indices [ranges[i], ranges[i + 1]).
- Has to contain at least one element.
-
+ * Part i contains the indices [ranges[i], ranges[i + 1]).
+ * Has to contain at least one element.
+ * The first element has to be 0.
+ *
  * @return a Partition representing the given contiguous partitioning.
  */
  static std::unique_ptr<Partition> build_from_contiguous(