A Simplified Layout Algebra for GPU Kernels

Abstract

The layout algebra behind NVIDIA's CuTe library describes how tensor tiles map onto threads and memory. A layout is a shape with a stride for each entry, so a coordinate's address is the sum of its indices times their strides. As formalized by Shah and by Cecka, its central operations (e.g., composition, complement, division, product) are guarded by at least six interacting admissibility conditions. We show that this complexity is an artifact of a representational choice, not of the domain. We present a layout algebra in which every layout names the space its inputs and outputs live in: logical, physical, or thread-value indices. A layout into logical or thread-value space must be a bijection from its $N$ coordinates onto $[0,N)$. A layout into physical space must be alias-free: distinct coordinates land on distinct addresses, except under broadcast, where many threads read one value on purpose and the layout must mark where. Nothing maps out of a physical index. Composing $f$ with $g$ requires only that $f$'s range be exactly $g$'s flattened coordinate set. The index arithmetic can always be emitted naively: compute $f$, split the result by division and remainder into $g$'s coordinates, and apply $g$. Often the composite simplifies to a layout; we show that when it does, there is a coarsest such form, whose digits every other one refines, and we emit that. A warp's or thread's slice is a \emph{restriction}. A restricted layout is no longer composable, which forces the user to settle the whole view of a tile before taking slices. We validate the algebra by reproducing within it CuTe's operations, such as composition, complement, and division. We implement it in OCaml and use the implementation to rederive CuTe's atoms and usage patterns. The decision procedure is proved correct in Rocq.