From 47cd80ce35cead2f63dcedf12acadd044a1208f4 Mon Sep 17 00:00:00 2001 From: Marcel Koch Date: Tue, 26 Oct 2021 16:24:10 +0200 Subject: [PATCH] add partition builder from global size --- .../unified/distributed/partition_kernels.cpp | 20 ++++ core/device_hooks/common_kernels.inc.cpp | 1 + core/distributed/partition.cpp | 15 +++ core/distributed/partition_kernels.hpp | 7 ++ include/ginkgo/core/distributed/partition.hpp | 13 +++ reference/distributed/partition_kernels.cpp | 21 ++++ .../test/distributed/partition_kernels.cpp | 104 ++++++++++++++++++ test/distributed/partition_kernels.cpp | 52 +++++++++ 8 files changed, 233 insertions(+) diff --git a/common/unified/distributed/partition_kernels.cpp b/common/unified/distributed/partition_kernels.cpp index 6b2c89b6234..88ea639e7a2 100644 --- 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 exec, GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING); +template +void build_ranges_from_global_size(std::shared_ptr exec, + int num_parts, int64 global_size, + Array& 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(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 index 0ab2cf79776..87c13c01786 100644 --- 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 index 40c573b646d..903543b48b9 100644 --- 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::build_from_contiguous( } +template +std::unique_ptr> +Partition::build_from_global_size( + std::shared_ptr exec, comm_index_type num_parts, + global_index_type global_size) +{ + Array ranges(exec, num_parts + 1); + exec->run(partition::make_build_ranges_from_global_size( + num_parts, global_size, ranges)); + return Partition::build_from_contiguous(exec, ranges); +} + + template void Partition::compute_range_starting_indices() { diff --git a/core/distributed/partition_kernels.hpp b/core/distributed/partition_kernels.hpp index 2d316371ed1..4ce6025d066 100644 --- a/core/distributed/partition_kernels.hpp +++ b/core/distributed/partition_kernels.hpp @@ -60,6 +60,11 @@ namespace kernels { const Array& mapping, \ distributed::Partition* partition) +#define GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE(LocalIndexType) \ + void build_ranges_from_global_size( \ + std::shared_ptr exec, int num_parts, \ + int64 global_size, Array& ranges) + #define GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES(LocalIndexType) \ void build_starting_indices(std::shared_ptr exec, \ const global_index_type* range_offsets, \ @@ -82,6 +87,8 @@ namespace kernels { template \ GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING(LocalIndexType); \ template \ + GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE(LocalIndexType); \ + template \ GKO_DECLARE_PARTITION_BUILD_STARTING_INDICES(LocalIndexType); \ template \ GKO_DECLARE_PARTITION_IS_ORDERED(LocalIndexType) diff --git a/include/ginkgo/core/distributed/partition.hpp b/include/ginkgo/core/distributed/partition.hpp index a315dc461e7..e353fb0db72 100644 --- a/include/ginkgo/core/distributed/partition.hpp +++ b/include/ginkgo/core/distributed/partition.hpp @@ -204,6 +204,19 @@ class Partition : public EnablePolymorphicObject>, std::shared_ptr exec, const Array& 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 build_from_global_size( + std::shared_ptr 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 index cf12dbfe31f..d2fbcf433ee 100644 --- a/reference/distributed/partition_kernels.cpp +++ b/reference/distributed/partition_kernels.cpp @@ -93,6 +93,27 @@ void build_from_mapping(std::shared_ptr exec, GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PARTITION_BUILD_FROM_MAPPING); +template +void build_ranges_from_global_size(std::shared_ptr exec, + int num_parts, int64 global_size, + Array& 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((i - 1) < rest); + } +} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE( + GKO_DECLARE_PARTITION_BUILD_FROM_GLOBAL_SIZE); + + template void build_starting_indices(std::shared_ptr exec, const global_index_type* range_offsets, diff --git a/reference/test/distributed/partition_kernels.cpp b/reference/test/distributed/partition_kernels.cpp index 0e04bc9b88c..ed80e5ff0a4 100644 --- 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::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::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::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 index bf94763be3d..ff9aec6bd98 100644 --- 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::build_from_global_size( + this->ref, num_parts, global_size); + auto dpart = + gko::distributed::Partition::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::build_from_global_size( + this->ref, num_parts, global_size); + auto dpart = + gko::distributed::Partition::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::build_from_global_size( + this->ref, num_parts, global_size); + auto dpart = + gko::distributed::Partition::build_from_global_size( + this->exec, num_parts, global_size); + + this->assert_equal(part, dpart); +} + + } // namespace