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.

template<typename F>
inline bool for_each_owner_byte_run(
size_t elemsize,
size_t max_runs,
F &&emit
) const#

Enumerate the certified constant-owner byte runs of the flat allocation, merged across adjacent equal-owner segments, in increasing byte order: calls emit(byte_start, byte_len, owner) per run; emit returns false to abort the walk early (the enumeration then also returns false).

Returns false (without emitting) when the walk would exceed max_runs; see try_block_owners for the budget semantics. The certificate for one owner() evaluation per segment is the divisibility chain of the exact tiling’s sorted leaf strides (owner constant on intervals of the padded space aligned to the smallest place-leaf stride).

inline size_t min_owner_run_bytes(size_t elemsize) const#

Certified owner-run granularity in bytes: the smallest place-leaf stride times the element size.

Ownership is constant on padded-linear intervals of this size (see try_block_owners), so it is the “smallest part” of the placement layout:

  • >= the placement-block size: blocks straddle at most one ownership boundary and the analytic plan is exact or near-exact;

  • < the placement-block size: NO block can be provably pure &#8212; the plan is all-straddles with closed-form misplacement, a signal that identity storage is block-hostile for this spec (prefer a relayout);

  • the analytic walk costs ~total_bytes / this many owner() calls, which is what try_block_owners’ max_runs guards.

Returns 0 when there is no place mode (single-owner layouts).

inline size_t min_owner_run_elems() const#

The same granularity in ELEMENTS of the padded space (the walk’s native unit); 0 means no ownership boundary at all (single owner).

inline ::cuda::std::optional<::std::vector<pos4>> try_block_owners(
size_t block_size_bytes,
size_t elemsize,
size_t *misplaced_bytes,
size_t max_runs = 0
) const#

Analytic per-placement-block owners: divide the ownership layout by the placement-block layout, without sampling.

Owners are derived with one owner() evaluation per provably-constant run. The certificate comes from the leaf algebra: in an exact tiling the sorted leaf strides form a divisibility chain, so every place-mode coordinate is constant on intervals of the padded linear space aligned to the smallest place-leaf stride. Blocks straddling an ownership boundary are assigned by exact byte majority and their error is accumulated in *misplaced_bytes (0 means the partition factors through placement blocks and the plan is exact).

The linearization convention matches the composite allocation path (dimension 0 varies fastest, see dim4::index_to_pos).

Parameters:

max_runs – evaluation budget for the walk; 0 (the default) self-scales to max(16 x nblocks, 1 << 16) &#8212; a small constant of owner() evaluations per placement block produced, so the budget grows with the allocation instead of being an arbitrary constant.

Returns:

per-block owners, or nullopt when the run enumeration would exceed max_runs (dense sub-block interleavings such as element-cyclic): callers fall back to sampled majority.

inline ::cuda::std::optional<::std::vector<block_run>> try_block_runs(
size_t block_size_bytes,
size_t elemsize,
size_t max_runs = 0
) const#

Divide the ownership layout by the placement-block layout, returning the quotient DIRECTLY as maximal same-owner block runs.

This is the exact tier’s natural output: one run per (cuMemCreate, cuMemMap) pair, with no per-block materialization and no downstream merge &#8212; the census/merge round-trip of the one-owner-per-block representation exists only for tiers whose primitive is per-block (sampling, straddle majority). Succeeds iff every internal ownership boundary is aligned to the block size (the strict quotient exists, so the plan is exact and misplacement is zero by construction); returns nullopt otherwise &#8212; callers fall through to try_block_owners (census with closed-form majority) and then to sampling.

NB the number of runs equals the number of physical allocations the caller will create: a block-aligned fine interleaving (e.g. block_cyclic at exactly the block size) is EXACT but yields one run per block &#8212; the plan is honest about that cost rather than hiding it.

SCOPE NOTE: this computes the EXTENSION of the quotient (its runs, enumerated by the certified walk with alignment verified a posteriori), not the quotient as a layout object. A symbolic, leaf-level quotient &#8212; existence decided a priori by congruences on the place-leaf strides/offsets, construction by dividing the strides through, O(leaves) regardless of run count &#8212; is possible future work; what stands in the way is that the allocation is the RESTRICTION of the padded-space layout to the true extents, and that restriction is not a pure layout operation (it is exactly what the row walk performs by hand). Until then the walk’s cost is bounded by max_runs and calibrated against the sampled fallback, so the proxy is never a performance regression.

Parameters:

max_runs – walk budget, as in try_block_owners (0 = auto)

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)