Part III — Middle end

Layout & data formats

Lesson 06 decided where every buffer lives and when — it collapsed peak memory by reusing storage across non-overlapping lifetimes. What it never touched is the shape of the bytes inside each buffer: given a 4-D activation, which index runs fastest in memory? That choice is invisible to the math and decisive for the hardware. The wrong physical layout makes every downstream kernel read HBM in a strided, uncoalesced pattern, locks you out of the tensor cores, and — worst of all — forces the compiler to splice in transpose kernels between ops that disagree, each one a full extra HBM round trip. This lesson assigns a physical format to every tensor and then propagates those choices across the graph so adjacent transposes cancel.

1 · Why layout decides speed: coalescing and tensor-core fragments

A layout is the rule that maps a tensor's logical indices to linear memory addresses — concretely, the strides: for a 2-D tensor, row-major means address = i·stride_row + j with stride_row = ncols (the j axis is contiguous); column-major makes i contiguous instead. The numbers stored are identical; only the order on the wire changes. Two hardware facts make that order decisive.

Coalescing. HBM is read in 128-byte cache lines, and a warp's 32 threads issue 32 loads at once. If those addresses fall in one line, the controller serves them in one transaction at peak bandwidth; if they are scattered across 32 lines, you pay up to 32× the transactions for the same data (gpu_kernels · 04 coalesced access, gpu_kernels · 02 memory hierarchy). Whether a kernel's threads stride contiguously is a direct consequence of which axis the layout put innermost. A reduction over channels wants channels contiguous; if the layout buried channels behind H·W, every load jumps a full plane and bandwidth craters.

Tensor-core fragments. An mma.sync instruction multiplies a small fixed-shape fragment — for bf16 the canonical tile is 16×8×16, distributed as 8 elements per thread across the 32-lane warp in a specific register layout the hardware expects (gpu_kernels · 09 tensor cores). The fast path that loads those fragments, ldmatrix, assumes the operand tile sits contiguously in shared memory in a known order. If your global layout doesn't let the kernel stage a contiguous, aligned tile into SRAM, it either falls back to the ~15× slower CUDA-core path or inserts a reformat first. Tensor-core eligibility is a layout property, not just a dtype property.

2 · NCHW vs NHWC: where transposes come from

A 4-D convolution activation has axes N (batch), C (channels), H, W (spatial). Two layouts dominate, named by their axis order from outermost to innermost:

• NCHW (channels-major, "channels-first"): the W axis is contiguous, then H, then C. PyTorch's historical default and what many CPU/cuDNN kernels were first written for.
• NHWC (channels-last): the C axis is contiguous. This is what NVIDIA tensor cores want for convolution — a conv is implicitly a matmul over channels, so making C innermost lets the kernel load a fragment of C values per spatial position straight into an mma tile. cuDNN's fastest conv algorithms and TensorRT default to NHWC; PyTorch exposes it as channels_last and it is the recommended format for mixed-precision conv nets.

Here is the trap. Suppose your graph is conv → batchnorm → conv. BatchNorm is a per-channel scale-and-shift; it is happiest with channels contiguous (NHWC) so each channel's mean/var streams in coalesced. The convs, on tensor cores, also want NHWC. So far so good — one global NHWC and everyone is happy. But now imagine a stray op in the chain that the compiler only has an NCHW kernel for (an older custom op, a pooling variant, a framework default). The compiler cannot run an NCHW kernel on NHWC bytes, so it splices in a transpose: a full kernel that reads the whole tensor in one order and writes it back in another. A transpose of a 256 MB activation is a 256 MB read plus a 256 MB write — ≈160 µs on an H100 at ≈3.3 TB/s — to move zero useful FLOPs. Do this between every pair of disagreeing ops and the transposes can cost more than the convs.

eager / naive layout — each op gets its locally preferred format:

  conv(NHWC) --[T]--> bn(NCHW) --[T]--> conv(NHWC)
              ^^^                ^^^
        transpose 256MB     transpose 256MB     ← 2 extra HBM round trips,
        (read+write)        (read+write)          zero useful FLOPs

propagated layout — pick ONE format the whole region tolerates:

  conv(NHWC) -------> bn(NHWC) -------> conv(NHWC)
        no transposes; convs hit tensor cores; bn coalesces on C

The two transposes are not in the math — the model never asked for them. They are an artifact of letting each op choose its layout in isolation. That is the disagreement layout assignment exists to resolve.

3 · Tiled, blocked, and packed layouts; padding for alignment

NCHW/NHWC are only the row/column-major story lifted to 4-D. Two further families matter for the metal.

Tiled / blocked layouts. Instead of a flat row-major sheet, store the tensor as a grid of small contiguous tiles — e.g. break the channel axis into blocks of 8 and store NCHWc (oneDNN's notation: an outer channel-block axis, spatial, then an inner block of 8 channels contiguous). Each tile then matches the shape a kernel stages into SRAM or feeds to an mma, so the innermost block is exactly fragment-sized and 128-bit-aligned. Tiling the storage means the kernel's natural working set is already contiguous — no gather to assemble a fragment. cuDNN's NC/32HW32 (int8) and the blocked formats in oneDNN/XLA are exactly this: the physical layout is pre-shaped to the consuming instruction.

Packing. Packing reorders/copies a tensor into the exact byte arrangement a kernel's inner loop wants ahead of time — the classic example is GEMM, which packs panels of A and B into contiguous, aligned strips so the microkernel streams them with no stride arithmetic. The pack is a one-time copy whose cost amortizes over the matmul's O(M·N·K) FLOPs; for a weight it can be done once at compile time and reused every inference.

Padding for alignment. Tensor cores and vector loads want dimensions that are multiples of the fragment/transaction size — e.g. an inner dim that is a multiple of 8 (bf16) or 16, and base addresses on 128-byte boundaries. A channel count of 130 is awful: the last fragment is ragged, the kernel either masks (slow) or falls off the tensor-core path. Padding rounds 130 up to 136 (multiple of 8) or 256 with zeros so every tile is full and aligned. The cost is real and must be priced: padding 130→256 channels nearly doubles the bytes moved and the FLOPs computed on garbage zeros. The compiler pads only when the alignment win beats the wasted-traffic cost — small over-pad to hit a fragment boundary, yes; doubling the tensor, usually no.

4 · Layout propagation: a global, not local, assignment

The naive algorithm is greedy: ask each op for its single favorite layout and grant it. That is locally optimal and globally terrible, because every boundary where two neighbors' favorites disagree becomes a transpose. The real objective is to assign a layout to every tensor (every graph edge) so that the total cost — kernel slowdown from non-ideal layouts plus the transposes inserted at every mismatch — is minimized over the whole graph.

Model the graph as nodes (ops) and edges (tensors). Each op has a cost function over the layout of each of its operands and result: a conv is cheap in NHWC, costly in NCHW; a transpose between adjacent edges costs the bytes of the tensor. Choosing a layout for one edge constrains its neighbors, so you cannot decide edges independently — this is a combinatorial assignment over a graph, equivalent in spirit to a min-cost labeling / graph-coloring problem, and it is NP-hard in general. Compilers do not solve it exactly; they use heuristics:

Transpose elimination is the payoff pass. Algebraically, transpose(transpose(x)) with inverse permutations is the identity and is deleted outright; a transpose that meets a commuting elementwise op can be hoisted past it so it merges with the next transpose, and the pair then cancels. After propagation has chosen consistent layouts for long stretches of the graph, elimination sweeps up the few remaining transposes the conflict-resolution step left behind. The conv→bn→conv chain that had two transposes under greedy assignment ends with zero. This is the same correctness invariant from lessons 03–04: a layout change plus its inverse is a provable no-op on the math, which is exactly why the compiler is allowed to delete it.

5 · How XLA, TVM, and cuDNN actually pick layouts

Every production stack runs some version of this assignment-then-eliminate pipeline:

• XLA runs a layout-assignment pass over HLO: it seeds layouts from layout-constrained ops (custom calls, convolutions, the entry/exit of the computation) and propagates them through the graph to minimize copies, then a later algebraic-simplifier/transpose-folding pass removes the transposes that became redundant. Layout is a first-class part of an HLO tensor's type.
• TVM exposes it as the ConvertLayout pass: you name a desired layout for heavy ops (e.g. NHWC for conv), and the pass walks the graph inserting the minimum layout transforms and folding consecutive ones, so the rest of the network is pulled into a consistent format around the anchors.
• cuDNN / TensorRT sit at the kernel level: cuDNN's fastest conv algorithms require NHWC (or blocked int8 formats like NC/32HW32), so the framework's job is to feed them tensors already in that layout — PyTorch's channels_last memory format exists precisely so the whole conv net stays NHWC end-to-end and no per-op transpose is needed. TensorRT chooses per-layer formats during its build/optimization phase and reconciles them, inserting reformat layers only where unavoidable.

The verdict the field has converged on: pick layouts globally around the tensor-core-bound anchors, make everything else conform, and treat a surviving transpose as a bug to be eliminated, not a cost to be tolerated.

6 · Drive it: the layout chooser

The widget below is the conv → bn → conv chain. Each op has a locally preferred layout (the convs want NHWC for tensor cores; toggle BatchNorm to prefer the other format to manufacture a conflict). Choose a single global layout, or flip on propagate to cancel transposes to let the assignment pass conform the agnostic ops and run transpose elimination. Watch the KPIs: with greedy local choices the chain sprouts transposes (red arrows), bandwidth balloons with the extra round trips, and a wrong global format drops the convs off the tensor-core path. Propagation collapses the transposes to zero and the convs light up green.

The shape of the result is the whole lesson: a layout chosen op-by-op looks fine locally and is riddled with hidden transposes globally; one consistent layout, propagated around the tensor-core anchors, pays the data movement exactly once. Notice the conflict toggle — when BatchNorm wants the opposite format, a pure global pick still costs some transpose, but propagation places it where one conversion serves the whole region instead of one per boundary.

Failure modes & checklist

Checkpoint

Where this points next

You now have a graph that is canonicalized, fused, memory-planned, and laid out: every tensor has a physical format chosen to coalesce, feed the tensor cores, and avoid transposes, with the few unavoidable reformats placed where one conversion serves a whole region. But everything so far is still declarative — it specifies what to compute on which bytes, not how the loops run. A single laid-out matmul still has astronomically many loop nests that compute it: which axis is the outer loop, how big are the tiles, what gets staged into SRAM, where do you vectorize and unroll. Those choices vary performance by 10–100× for identical FLOPs and an identical layout. Picking them is scheduling — the algorithm/schedule split, tiling for the memory hierarchy, the loop-nest knobs. That is lesson 08, Scheduling — the loop nest.

Interview prompts

• What is a tensor's "layout" and why does it matter if the math is identical? (§1 — it's the index→address stride map, e.g. row- vs column-major / NCHW vs NHWC; it decides which warp loads coalesce into one 128-byte transaction and whether a tile can become a tensor-core fragment, so it sets bandwidth and tensor-core eligibility despite being a no-op on the numbers.)
• Why do NVIDIA tensor cores prefer NHWC for convolution? (§2 — a conv is implicitly a matmul over channels; making C the contiguous axis lets the kernel load a fragment of channel values per spatial position straight into an mma tile, so cuDNN's fastest conv algos and channels_last want C innermost.)
• Where do transpose kernels come from, and what do they cost? (§2 — from two neighboring ops disagreeing on layout; the compiler inserts a reformat that reads the whole tensor in one order and writes it in another — a full HBM round trip, ≈160 µs for 256 MB — for zero useful FLOPs.)
• Why is layout assignment a global problem rather than per-op? (§4 — granting each op its local favorite is locally optimal but makes every disagreement a transpose; minimizing total cost = kernel slowdown + transposes over all edges is a combinatorial min-cost graph labeling, NP-hard, so compilers anchor-and-propagate heuristically.)
• What is transpose elimination and why is it legal? (§4 — a pass that deletes a transpose composed with its inverse (identity) and pushes transposes through commuting elementwise ops so adjacent ones cancel; legal because a permutation plus its inverse provably preserves the math.)
• When should the compiler pad a dimension, and when not? (§3 — pad a small ragged inner dim up to the mma/transaction multiple to keep fragments full and aligned (cheap win); refuse pads like 130→256 channels that nearly double bytes and FLOPs — the wasted traffic dwarfs the alignment benefit.)
• How do XLA / TVM / cuDNN assign layouts in practice? (§5 — XLA's layout-assignment pass seeds from layout-constrained ops and propagates to minimize copies, then folds transposes; TVM's ConvertLayout pulls the graph into a named format around anchors; cuDNN/TensorRT pick per-kernel formats and rely on the framework (channels_last) to feed them, inserting reformats only where unavoidable.)