partition.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_PARTITION_HPP_
#define GKO_PUBLIC_CORE_DISTRIBUTED_PARTITION_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/polymorphic_object.hpp>
#include <ginkgo/core/base/segmented_array.hpp>
#include <ginkgo/core/base/types.hpp>
 
 
namespace gko {
namespace experimental {
namespace distributed {
 
 
template <typename LocalIndexType = int32, typename GlobalIndexType = int64>
class Partition : public EnablePolymorphicObject<
                      Partition<LocalIndexType, GlobalIndexType>>,
                  public EnablePolymorphicAssignment<
                      Partition<LocalIndexType, GlobalIndexType>> {
    friend class EnablePolymorphicObject<Partition>;
    static_assert(sizeof(GlobalIndexType) >= sizeof(LocalIndexType),
                  "GlobalIndexType must be at least as large as "
                  "LocalIndexType");
 
public:
    using EnablePolymorphicAssignment<Partition>::convert_to;
    using EnablePolymorphicAssignment<Partition>::move_to;
 
    using local_index_type = LocalIndexType;
    using global_index_type = GlobalIndexType;
 
    size_type get_size() const { return size_; }
 
    size_type get_num_ranges() const noexcept
    {
        return offsets_.get_size() - 1;
    }
 
    comm_index_type get_num_parts() const noexcept { return num_parts_; }
 
    comm_index_type get_num_empty_parts() const noexcept
    {
        return num_empty_parts_;
    }
 
    const global_index_type* get_range_bounds() const noexcept
    {
        return offsets_.get_const_data();
    }
 
    const comm_index_type* get_part_ids() const noexcept
    {
        return part_ids_.get_const_data();
    }
 
    const local_index_type* get_range_starting_indices() const noexcept
    {
        return starting_indices_.get_const_data();
    }
 
    const local_index_type* get_part_sizes() const noexcept
    {
        return part_sizes_.get_const_data();
    }
 
    local_index_type get_part_size(comm_index_type part) const;
 
    const segmented_array<size_type>& get_ranges_by_part() const
    {
        return ranges_by_part_;
    }
 
    bool has_connected_parts() const;
 
    bool has_ordered_parts() const;
 
    static std::unique_ptr<Partition> build_from_mapping(
        std::shared_ptr<const Executor> exec,
        const array<comm_index_type>& mapping, comm_index_type num_parts);
 
    static std::unique_ptr<Partition> build_from_contiguous(
        std::shared_ptr<const Executor> exec,
        const array<global_index_type>& ranges,
        const array<comm_index_type>& part_ids = {});
 
    static std::unique_ptr<Partition> build_from_global_size_uniform(
        std::shared_ptr<const Executor> exec, comm_index_type num_parts,
        global_index_type global_size);
 
private:
    Partition(std::shared_ptr<const Executor> exec,
              comm_index_type num_parts = 0, size_type num_ranges = 0);
 
    static std::unique_ptr<Partition> create(
        std::shared_ptr<const Executor> exec, comm_index_type num_parts = 0,
        size_type num_ranges = 0);
 
    void finalize_construction();
 
    comm_index_type num_parts_;
    comm_index_type num_empty_parts_;
    global_index_type size_;
    array<global_index_type> offsets_;
    array<local_index_type> starting_indices_;
    array<local_index_type> part_sizes_;
    array<comm_index_type> part_ids_;
    segmented_array<size_type> ranges_by_part_;
};
 
 
}  // namespace distributed
}  // namespace experimental
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_PARTITION_HPP_