cuda::experimental::places::cute_partition_descriptor#

class cute_partition_descriptor#

Canonical runtime description of a structured tensor partition.

Used at polymorphic and ABI boundaries where one concrete C++ partition type cannot be retained. Normal C++ code should construct a typed cute_partition through make_partition().

Public Functions

inline cute_partition_descriptor(
const ::std::vector<layout_leaf> &place_leaves,
const ::std::vector<int> &place_axes,
const ::std::vector<layout_leaf> &local_leaves,
dim4 padded_dims,
dim4 true_dims,
dim4 grid_dims
)#

Construct a partition from flattened leaves (expert form)

Throws std::invalid_argument unless the two modes together tile the padded space exactly (bijectivity - validated in O(leaves)).

Parameters:

  • place_leaves – Leaves of the place mode, leaf 0 fastest; one leaf per used grid axis (at most max_leaves)

  • place_axes – Grid axis associated with each place leaf

  • local_leaves – Leaves of the local mode, leaf 0 fastest (at most max_leaves)

  • padded_dims – Padded tensor extents the strides refer to

  • true_dims – True tensor extents (the predicate)

  • grid_dims – Extents of the grid of places

inline const dim4 &true_dims() const#

True tensor extents (the predicate for the padded space)

inline const dim4 &padded_dims() const#

Padded tensor extents the leaf strides refer to.

inline const dim4 &grid_dims() const#

Extents of the grid of places.

inline ::cuda::std::span<const layout_leaf> place_leaves() const#

Leaves of the place mode (leaf 0 fastest)

inline ::cuda::std::span<const int> place_axes() const#

Grid axis associated with each place leaf.

inline ::cuda::std::span<const layout_leaf> local_leaves() const#

Leaves of the local mode (leaf 0 fastest)

inline size_t num_places() const#

Number of places the partition distributes over (product of place extents; grid axes not bound to any dimension receive coordinate 0 and do not count)

inline size_t tiles_per_place() const#

Number of padded elements owned by each place (product of local extents)

inline pos4 owner(pos4 data_coords) const#

Grid position owning the element at the given coordinates.

Total on all true coordinates (true extents never exceed the padded ones); grid axes not bound to any dimension get coordinate 0.

template<typename S>
inline auto apply(
const S &s,
pos4 place_position,
dim4 grid_dims
) const#

Sub-shape owned by one place, for parallel_for over a grid.

Follows the partitioner contract of the classic partitioners: the place is given as its linear index in the dispatch loop (pos4(i)), and the returned shape enumerates the coordinates that place owns. Coordinates beyond the true extents (padding phantoms, for uneven covers) are excluded by the sub-shape’s predicate rather than by restructuring the iteration (the CuTe predication idiom).

Parameters:

  • s – Shape of the task (must match the partition’s true extents)

  • place_position – Linear place index in .x (dispatch convention)

  • grid_dims – Extents of the grid (must match the partition’s)

template<size_t dims>
inline auto apply(
const box<dims> &b,
pos4 place_position,
dim4 grid_dims
) const#

Sub-shape owned by one place, restricted to a region of the tensor.

The box is not a shape: it is a region within the coordinate space of the tensor this partition was built for (the partition remains the authority on the extents). Each place enumerates its own coordinates and the sub-shape’s predicate keeps those inside the box - so the iteration chunks stay aligned with data ownership, unlike scale-free partitioners that split the box itself.

Parameters:

  • b – Region to iterate (must be contained in [0, true extents))

  • place_position – Linear place index in .x (dispatch convention)

  • grid_dims – Extents of the grid (must match the partition’s)

inline size_t place_offset(size_t place_index) const#

Linear element offset (in the padded space) of a place’s first element, given the place’s linear index in place-mode order (leaf 0 fastest)

inline int cmp(const cute_partition_descriptor &o) const#

Structural comparison (used for data place ordering)

inline bool operator==(const cute_partition_descriptor &o) const#
inline bool operator!=(const cute_partition_descriptor &o) const#

Public Static Attributes

static constexpr size_t max_leaves = cute_partition_max_leaves#

Maximum number of leaves per mode (see cute_partition_max_leaves)