|
Uni20
Asynchronous tensor and runtime library.
|
Status: active design note for future block-sparse lowering and placement.
This is a draft design note. It records the intended design for coalescing block-sparse sub-blocks into larger kernel operations, to reduce kernel/launch count while retaining sector identity and avoiding structural zeros. It is design direction, not current implemented behavior.
Related notes:
docs/symmetry/block_sparse_tensor.md — the tensor and layout this operates over.docs/architecture/ordering_and_backend_lowering.md — stream/event lowering; per-launch cost.docs/architecture/execution.md — coalescing as one use of the cost-model knob.docs/architecture/backend_dispatch.md — batched/coalesced kernels as backend capabilities.docs/tensor_network/rabc_contraction_scheduling.md — the R/A/B/C apply this most affects.Block-sparse contractions decompose into many small GEMMs, one per sector combination. A large fraction of these are tiny — in the calibrated R/A/B/C data, roughly half the kernel launches carry near-zero flops (the small-block, high-connectivity "tail"), and that fraction does not shrink as bond dimension m grows. Tiny GEMMs are launch-bound, so the most fundamental remedy is to issue fewer, larger kernels. Coalescing along a single axis does this with no structural zeros: it is a kernel-execution lowering expressed as wider mdspan views over the same buffer storage, with each individual block retaining its sector identity so results map back to the correct sectors. The decision whether to coalesce a group is the same flop/launch crossover the R/A/B/C cost model quantifies, applied to the data structure instead of to placement.
For a block GEMM C[m,n] = Σ_k A[m,k]·B[k,n] over block-multi-indices, a coalesced dense operand ‘A’over a chosen(m-set, k-set)rectangle is dense only if every (m,k)pair in the rectangle is a populated sector. Coalesce along **one** axis (fixk, take a set ofm; or fixm, take a set ofk) and you can choose exactly the populated blocks → no holes. Coalesce along **two** axes and the 2-D setm-set × k-set` generally has missing combinations → structural zeros. So "single axis, no zeros" is the general statement, not a heuristic:
The two single-axis directions are both available and are distinct planning choices: free-axis grouping (stack the valid output rows for a shared contracted index → one taller GEMM) and contraction-axis grouping (concatenate along k → one wider-K GEMM). You may take one or the other, not both at once.
A logical block is an mdspan<T, extents, layout_stride> with its (offset, strides) into a shared buffer; a coalesced operand is a wider view over the same buffer with the group's leading dimension. The same bytes are simultaneously individual block views carrying sector identity and one wider matrix view handed to GEMM. After the GEMM, the per-block views associate each result region with the correct sector.
The difference between the two memory arrangements is whether a gather is needed:
Which to use is adjudicated by the cost model: pack cost vs launch savings. The consequence for the tensor is that the layout's memory arrangement should be chosen so the expected coalescing groups are stride-addressable — interleave the axis you will coalesce. Layout is therefore a contraction-pattern-aware decision, not arbitrary.
Coalescing trades flops for launches. Below a block size it wins (launch-bound); above it, coalescing structural zeros is pure waste (flop-bound). The crossover is roughly where launch_us × N ≈ wasted_flops / gflops, i.e. the same calibrated constants used for placement (see ../architecture/execution.md for the numbers). For pure batched GEMM there is no flop waste, so the crossover is simply where the launch savings exceed any pack cost.
A single-axis coalescing group is simultaneously:
Isend, directly reducing the per-message / finite-tag overhead from ../architecture/ordering_and_backend_lowering.md).The planner chooses groups and emits a single DAG node per coalesced group (plus its scatter map), rather than per-block nodes. This is why coalescing mostly avoids the aliasing-hazard problem: the coalesced op is the node, so individual blocks are not separate concurrent ops on the shared region. It also means coalescing is part of placement planning, using the same cost model — not a runtime decision inside the kernel.
mdspan strided views over shared buffer storage, with per-block views retained so symmetry/sector identity is preserved; it is a lowering, never a representation change.B·C stage, contraction-axis for the A·Y accumulate in R/A/B/C)?