add partition builder from global size

ginkgo-project · Oct 26, 2021 · 47cd80c · 47cd80c
1 parent e0ce767
commit 47cd80c
Show file tree

Hide file tree

Showing 8 changed files with 233 additions and 0 deletions.
diff --git a/common/unified/distributed/partition_kernels.cpp b/common/unified/distributed/partition_kernels.cpp
@@ -122,6 +122,26 @@ void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING);
 
 
+template <typename LocalIndexType>
+void build_ranges_from_global_size(std::shared_ptr<const DefaultExecutor> exec,
+                                   int num_parts, int64 global_size,
+                                   Array<LocalIndexType>& 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 + static_cast<LocalIndexType>(i < rest);
+        },
+        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_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE);
+
+
 }  // namespace partition
 }  // namespace GKO_DEVICE_NAMESPACE
 }  // namespace kernels

diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp
@@ -183,6 +183,7 @@ GKO_STUB(GKO_PARTITION_COUNT_RANGES);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_CONTIGUOUS);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES);
+GKO_STUB_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_PARTITION_IS_ORDERED);
 
 

diff --git a/core/distributed/partition.cpp b/core/distributed/partition.cpp
@@ -44,6 +44,8 @@ namespace partition {
 GKO_REGISTER_OPERATION(count_ranges, partition::count_ranges);
 GKO_REGISTER_OPERATION(build_from_mapping, partition::build_from_mapping);
 GKO_REGISTER_OPERATION(build_from_contiguous, partition::build_from_contiguous);
+GKO_REGISTER_OPERATION(build_ranges_from_global_size,
+                       partition::build_ranges_from_global_size);
 GKO_REGISTER_OPERATION(build_starting_indices,
                        partition::build_starting_indices);
 GKO_REGISTER_OPERATION(is_ordered, partition::is_ordered);
@@ -86,6 +88,19 @@ Partition<LocalIndexType>::build_from_contiguous(
 }
 
 
+template <typename LocalIndexType>
+std::unique_ptr<Partition<LocalIndexType>>
+Partition<LocalIndexType>::build_from_global_size(
+    std::shared_ptr<const Executor> exec, comm_index_type num_parts,
+    global_index_type global_size)
+{
+    Array<global_index_type> ranges(exec, num_parts + 1);
+    exec->run(partition::make_build_ranges_from_global_size(
+        num_parts, global_size, ranges));
+    return Partition<LocalIndexType>::build_from_contiguous(exec, ranges);
+}
+
+
 template <typename LocalIndexType>
 void Partition<LocalIndexType>::compute_range_starting_indices()
 {

diff --git a/core/distributed/partition_kernels.hpp b/core/distributed/partition_kernels.hpp
@@ -60,6 +60,11 @@ namespace kernels {
                             const Array<comm_index_type>& mapping,       \
                             distributed::Partition<LocalIndexType>* partition)
 
+#define GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE(LocalIndexType) \
+    void build_ranges_from_global_size(                              \
+        std::shared_ptr<const DefaultExecutor> exec, int num_parts,  \
+        int64 global_size, Array<LocalIndexType>& ranges)
+
 #define GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES(LocalIndexType)          \
     void build_starting_indices(std::shared_ptr<const DefaultExecutor> exec,  \
                                 const global_index_type* range_offsets,       \
@@ -82,6 +87,8 @@ namespace kernels {
     template <typename LocalIndexType>                            \
     GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING(LocalIndexType);     \
     template <typename LocalIndexType>                            \
+    GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE(LocalIndexType); \
+    template <typename LocalIndexType>                            \
     GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES(LocalIndexType); \
     template <typename LocalIndexType>                            \
     GKO_DECLARE_PARTITION_IS_ORDERED(LocalIndexType)

diff --git a/include/ginkgo/core/distributed/partition.hpp b/include/ginkgo/core/distributed/partition.hpp
@@ -204,6 +204,19 @@ class Partition : public EnablePolymorphicObject<Partition<LocalIndexType>>,
         std::shared_ptr<const Executor> exec,
         const Array<global_index_type>& ranges);
 
+    /**
+     * Builds a partition by evenly distributing the global range.
+     * @param exec  the Executor on which the partition should be built
+     * @param num_parts  the number of parst used in this partition
+     * @param global_size  the global size of this partition
+     * @return  a Partition where each range has either
+     * `floor(global_size/num_parts)` or `floor(global_size/num_parts) + 1`
+     * indices.
+     */
+    static std::unique_ptr<Partition> build_from_global_size(
+        std::shared_ptr<const Executor> exec, comm_index_type num_parts,
+        global_index_type global_size);
+
     /**
      * Creates a partition stored on the given executor with the given number of
      * consecutive ranges and parts.

diff --git a/reference/distributed/partition_kernels.cpp b/reference/distributed/partition_kernels.cpp
@@ -93,6 +93,27 @@ void build_from_mapping(std::shared_ptr<const DefaultExecutor> exec,
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING);
 
 
+template <typename LocalIndexType>
+void build_ranges_from_global_size(std::shared_ptr<const DefaultExecutor> exec,
+                                   int num_parts, int64 global_size,
+                                   Array<LocalIndexType>& ranges)
+{
+    const auto size_per_part = global_size / num_parts;
+    const auto rest = global_size - (num_parts * size_per_part);
+
+    auto* ranges_ptr = ranges.get_data();
+
+    ranges_ptr[0] = 0;
+    for (int i = 1; i < num_parts + 1; ++i) {
+        ranges_ptr[i] = ranges_ptr[i - 1] + size_per_part +
+                        static_cast<LocalIndexType>((i - 1) < rest);
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
+    GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE);
+
+
 template <typename LocalIndexType>
 void build_starting_indices(std::shared_ptr<const DefaultExecutor> exec,
                             const global_index_type* range_offsets,

diff --git a/reference/test/distributed/partition_kernels.cpp b/reference/test/distributed/partition_kernels.cpp
@@ -162,6 +162,110 @@ TYPED_TEST(Partition, BuildsFromRanges)
     EXPECT_EQ(partition->get_part_sizes()[4], 1);
 }
 
+TYPED_TEST(Partition, BuildsFromGlobalSize)
+{
+    using local_index_type = typename TestFixture::local_index_type;
+
+    auto partition =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->ref, 5, 13);
+
+    EXPECT_EQ(partition->get_size(), 13);
+    EXPECT_EQ(partition->get_num_ranges(), 5);
+    EXPECT_EQ(partition->get_num_parts(), 5);
+    EXPECT_EQ(partition->get_const_range_bounds()[0], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[1], 3);
+    EXPECT_EQ(partition->get_const_range_bounds()[2], 6);
+    EXPECT_EQ(partition->get_const_range_bounds()[3], 9);
+    EXPECT_EQ(partition->get_const_range_bounds()[4], 11);
+    EXPECT_EQ(partition->get_const_range_bounds()[5], 13);
+    EXPECT_EQ(partition->get_part_ids()[0], 0);
+    EXPECT_EQ(partition->get_part_ids()[1], 1);
+    EXPECT_EQ(partition->get_part_ids()[2], 2);
+    EXPECT_EQ(partition->get_part_ids()[3], 3);
+    EXPECT_EQ(partition->get_part_ids()[4], 4);
+    EXPECT_EQ(partition->get_range_starting_indices()[0], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[1], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[2], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[3], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[4], 0);
+    EXPECT_EQ(partition->get_part_sizes()[0], 3);
+    EXPECT_EQ(partition->get_part_sizes()[1], 3);
+    EXPECT_EQ(partition->get_part_sizes()[2], 3);
+    EXPECT_EQ(partition->get_part_sizes()[3], 2);
+    EXPECT_EQ(partition->get_part_sizes()[4], 2);
+}
+
+
+TYPED_TEST(Partition, BuildsFromGlobalSizeEmptySize)
+{
+    using local_index_type = typename TestFixture::local_index_type;
+
+    auto partition =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->ref, 5, 0);
+
+    EXPECT_EQ(partition->get_size(), 0);
+    EXPECT_EQ(partition->get_num_ranges(), 5);
+    EXPECT_EQ(partition->get_num_parts(), 5);
+    EXPECT_EQ(partition->get_const_range_bounds()[0], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[1], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[2], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[3], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[4], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[5], 0);
+    EXPECT_EQ(partition->get_part_ids()[0], 0);
+    EXPECT_EQ(partition->get_part_ids()[1], 1);
+    EXPECT_EQ(partition->get_part_ids()[2], 2);
+    EXPECT_EQ(partition->get_part_ids()[3], 3);
+    EXPECT_EQ(partition->get_part_ids()[4], 4);
+    EXPECT_EQ(partition->get_range_starting_indices()[0], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[1], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[2], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[3], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[4], 0);
+    EXPECT_EQ(partition->get_part_sizes()[0], 0);
+    EXPECT_EQ(partition->get_part_sizes()[1], 0);
+    EXPECT_EQ(partition->get_part_sizes()[2], 0);
+    EXPECT_EQ(partition->get_part_sizes()[3], 0);
+    EXPECT_EQ(partition->get_part_sizes()[4], 0);
+}
+
+
+TYPED_TEST(Partition, BuildsFromGlobalSizeWithEmptyParts)
+{
+    using local_index_type = typename TestFixture::local_index_type;
+
+    auto partition =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->ref, 5, 3);
+
+    EXPECT_EQ(partition->get_size(), 3);
+    EXPECT_EQ(partition->get_num_ranges(), 5);
+    EXPECT_EQ(partition->get_num_parts(), 5);
+    EXPECT_EQ(partition->get_const_range_bounds()[0], 0);
+    EXPECT_EQ(partition->get_const_range_bounds()[1], 1);
+    EXPECT_EQ(partition->get_const_range_bounds()[2], 2);
+    EXPECT_EQ(partition->get_const_range_bounds()[3], 3);
+    EXPECT_EQ(partition->get_const_range_bounds()[4], 3);
+    EXPECT_EQ(partition->get_const_range_bounds()[5], 3);
+    EXPECT_EQ(partition->get_part_ids()[0], 0);
+    EXPECT_EQ(partition->get_part_ids()[1], 1);
+    EXPECT_EQ(partition->get_part_ids()[2], 2);
+    EXPECT_EQ(partition->get_part_ids()[3], 3);
+    EXPECT_EQ(partition->get_part_ids()[4], 4);
+    EXPECT_EQ(partition->get_range_starting_indices()[0], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[1], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[2], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[3], 0);
+    EXPECT_EQ(partition->get_range_starting_indices()[4], 0);
+    EXPECT_EQ(partition->get_part_sizes()[0], 1);
+    EXPECT_EQ(partition->get_part_sizes()[1], 1);
+    EXPECT_EQ(partition->get_part_sizes()[2], 1);
+    EXPECT_EQ(partition->get_part_sizes()[3], 0);
+    EXPECT_EQ(partition->get_part_sizes()[4], 0);
+}
+
 
 TYPED_TEST(Partition, IsConnected)
 {

diff --git a/test/distributed/partition_kernels.cpp b/test/distributed/partition_kernels.cpp
@@ -311,4 +311,56 @@ TYPED_TEST(Partition, BuildsFromContiguousWithOnlyOneEmptyPart)
 }
 
 
+TYPED_TEST(Partition, BuildsFromGlobalSize)
+{
+    using local_index_type = typename TestFixture::local_index_type;
+    const int num_parts = 7;
+    const global_index_type global_size = 708;
+
+    auto part =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->ref, num_parts, global_size);
+    auto dpart =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->exec, num_parts, global_size);
+
+    this->assert_equal(part, dpart);
+}
+
+
+TYPED_TEST(Partition, BuildsFromGlobalSizeEmpty)
+{
+    using local_index_type = typename TestFixture::local_index_type;
+    const int num_parts = 7;
+    const global_index_type global_size = 0;
+
+    auto part =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->ref, num_parts, global_size);
+    auto dpart =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->exec, num_parts, global_size);
+
+    this->assert_equal(part, dpart);
+
+}
+
+
+TYPED_TEST(Partition, BuildsFromGlobalSizeMorePartsThanSize)
+{
+    using local_index_type = typename TestFixture::local_index_type;
+    const int num_parts = 77;
+    const global_index_type global_size = 13;
+
+    auto part =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->ref, num_parts, global_size);
+    auto dpart =
+        gko::distributed::Partition<local_index_type>::build_from_global_size(
+            this->exec, num_parts, global_size);
+
+    this->assert_equal(part, dpart);
+}
+
+
 }  // namespace