gko::experimental::distributed::
Partition#

Ownership map: stores, for each contiguous range of global indices, which rank owns it. Three factory helpers cover the common cases:

  • build_from_contiguous: explicit per-rank ranges.

  • build_from_global_size_uniform: uniform block split of a given global size.

  • build_from_mapping: generic per-row ownership array (e.g. the output of ParMETIS).

The free functions build_partition_from_local_size and build_partition_from_local_range build a partition collectively from what each rank owns: its number of rows, or its range of global rows.

A Matrix takes its row partition, and optionally a different column partition, in read_distributed.

template<typename LocalIndexType = int32, typename GlobalIndexType = int64>
class Partition #

Inherits from

Represents a partition of a range of indices [0, size) into a disjoint set of parts. The partition is stored as a set of consecutive ranges [begin, end) with an associated part ID and local index (number of indices in this part before begin). Global indices are stored as 64 bit signed integers (int64), part-local indices use LocalIndexType, Part IDs use 32 bit signed integers (int).

For example, consider the interval [0, 13) that is partitioned into the following ranges: 

[0,3), [3, 6), [6, 8), [8, 10), [10, 13).
These ranges are distributed on three part with: 
p_0 = [0, 3) + [6, 8) + [10, 13),
p_1 = [3, 6),
p_2 = [8, 10).
The part ids can be queried from the get_part_ids array, and the ranges are represented as offsets, accessed by get_range_bounds, leading to the offset array: 
r = [0, 3, 6, 8, 10, 13]
so that individual ranges are given by [r[i], r[i + 1]). Since each part may be associated with multiple ranges, it is possible to get the starting index for each range that is local to the owning part, see get_range_starting_indices. These indices can be used to easily iterate over part local data. For example, the above partition has the following starting indices 
starting_index[0] = 0,
starting_index[1] = 0,
starting_index[2] = 3,  // second range of part 0
starting_index[3] = 0,
starting_index[4] = 5,  // third range of part 0
which you can use to iterate only over the the second range of part 0 (the third global range) with 
for(int i = 0; i < r[3] - r[2]; ++i){
  data[starting_index[2] + i] = val;
}

Template Parameters:

  • LocalIndexType – The index type used for part-local indices. To prevent overflows, no single part’s size may exceed this index type’s maximum value.

  • GlobalIndexType – The index type used for the global indices. Needs to be at least as large a type as LocalIndexType.

Public Functions

inline size_type get_size() const#

Returns the total number of elements represented by this partition.

Returns:

number elements.

inline size_type get_num_ranges() const noexcept#

Returns the number of ranges stored by this partition. This size refers to the data returned by get_range_bounds().

Returns:

number of ranges.

inline comm_index_type get_num_parts() const noexcept#

Returns the number of parts represented in this partition.

Returns:

number of parts.

bool equals(const Partition &other) const#

Checks whether this partition describes the same distribution as another one, i.e. both span the same global size and assign the same ranges to the same parts.

Comparing two distinct partitions that live on a device executor requires copying their range data to the host, so this is not a cheap operation. Comparing a partition with itself is free.

Parameters:

other – the partition to compare against.

Returns:

true if both partitions describe the same distribution.

inline comm_index_type get_num_empty_parts() const noexcept#

Returns the number of empty parts within this partition.

Returns:

number of empty parts.

inline const global_index_type *get_range_bounds() const noexcept#

Returns the ranges boundary array stored by this partition. range_bounds[i] is the beginning (inclusive) and range_bounds[i + 1] is the end (exclusive) of the ith range.

Returns:

range boundaries array.

inline const comm_index_type *get_part_ids() const noexcept#

Returns the part IDs of the ranges in this partition. For each range from get_range_bounds(), it stores the part ID in the interval [0, get_num_parts() - 1].

Returns:

part ID array.

inline const local_index_type *get_range_starting_indices(
) const noexcept#

Returns the part-local starting index for each range in this partition.

Consider the partition on [0, 10) with 

p_1 = [0-4) + [7-10),
p_2 = [4-7).
Then range_starting_indices[0] = 0, range_starting_indices[1] = 0, range_starting_indices[2] = 4.

Returns:

part-local starting index array.

inline const local_index_type *get_part_sizes() const noexcept#

Returns the part size array. part_sizes[p] stores the total number of indices in part p.

Returns:

part size array.

local_index_type get_part_size(comm_index_type part) const#

Returns the size of a part given by its part ID.

Warning

Triggers a copy from device to host.

Parameters:

part – the part ID.

Returns:

size of part.

inline const segmented_array<size_type> &get_ranges_by_part(
) const#

Returns the range IDs segmented by their part ID.

Returns:

range IDs segmented by part IDs

bool has_connected_parts() const#

Checks if each part has no more than one contiguous range.

Returns:

true if each part has no more than one contiguous range.

bool has_ordered_parts() const#

Checks if the ranges are ordered by their part index.

Implies that the partition is connected.

Returns:

true if the ranges are ordered by their part index.

Public Static Functions

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,
)#

Builds a partition from a given mapping global_index -> part_id.

Parameters:

  • exec – the Executor on which the partition should be built

  • mapping – the mapping from global indices to part IDs.

  • num_parts – the number of parts used in the mapping.

Returns:

a Partition representing the given mapping as a set of ranges

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 = {},
)#

Builds a partition consisting of contiguous ranges, one for each part.

Parameters:

  • exec – the Executor on which the partition should be built

  • ranges – the boundaries of the ranges representing each part. Part part_id[i] contains the indices [ranges[i], ranges[i + 1]). Has to contain at least one element. The first element has to be 0.

  • part_ids – the part ids of the provided ranges. If empty, then it will assume range i belongs to part i.

Returns:

a Partition representing the given contiguous partitioning.

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,
)#

Builds a partition by evenly distributing the global range.

Parameters:

  • exec – the Executor on which the partition should be built

  • num_parts – the number of parst used in this partition

  • global_size – the global size of this partition

Returns:

a Partition where each range has either floor(global_size/num_parts) or floor(global_size/num_parts) + 1 indices.

Building from per-rank data#

template<typename LocalIndexType, typename GlobalIndexType>
std::unique_ptr<Partition<LocalIndexType, GlobalIndexType>> gko::experimental::distributed::build_partition_from_local_range(
std::shared_ptr<const Executor> exec,
mpi::communicator comm,
span local_range,
)#

Builds a partition from a local range.

Warning

This throws, if the resulting partition would contain gaps. That means that for a partition of size n every local range r_i = [s_i, e_i) either s_i != 0 and another local range r_j = [s_j, e_j = s_i) exists, or e_i != n and another local range r_j = [s_j = e_i, e_j) exists.

Parameters:

  • exec – the Executor on which the partition should be built.

  • comm – the communicator used to determine the global partition.

  • local_range – the start and end indices of the local range.

Returns:

a Partition where each range has the individual local_start and local_ends.

template<typename LocalIndexType, typename GlobalIndexType>
std::unique_ptr<Partition<LocalIndexType, GlobalIndexType>> gko::experimental::distributed::build_partition_from_local_size(
std::shared_ptr<const Executor> exec,
mpi::communicator comm,
size_type local_size,
)#

Builds a partition from a local size.

Parameters:

  • exec – the Executor on which the partition should be built.

  • comm – the communicator used to determine the global partition.

  • local_size – the number of the locally owned indices

Returns:

a Partition where each range has the specified local size. More specifically, if this is called on process i with local_size s_i, then the range i has size s_i, and range r_i = [start, start + s_i), where start = sum_j^(i-1) s_j.