Layout Assignment and the Elimination of Transposes

The exported graph records a shape, not a memory order. A value of shape $[N, C, H, W]$ in the file says how many elements each axis has; it does not say whether the element at $(n, c, h, w)$ sits at offset $((nC + c)H + h)W + w$ or at $((nH + h)W + w)C + c$. Both are valid encodings of the same tensor, they differ in which axis varies fastest, and the choice changes the memory access pattern of every kernel that touches the value. On an accelerator whose matrix unit wants a contiguous inner dimension of a particular width, one of the two orders is fast and the other requires a repacking step first.

Layout assignment is the pass that picks an order for every value in the graph. It is a constraint problem over the edges: each operator has a set of layouts it can consume efficiently, a conversion between two layouts costs a copy, and the pass must choose one layout per value minimising total conversion cost. Done well, it removes almost every transpose from a model. Done badly, it leaves a graph full of copies that each read and write a full activation.

Layout as a permutation, not a property of the tensor

Represent a layout as a permutation of the axes, plus padding. The permutation says the storage order; the padding says how many elements are appended to each row so that the inner dimension is a multiple of the vector width the kernel wants. A tensor of shape $[1, 3, 224, 224]$ in float32 with a channels-last permutation [0, 2, 3, 1] occupies the same 602,112 bytes as in channels-first, but the byte at a given offset corresponds to a different logical element.

A transpose node in the graph is a claim that the output’s logical shape is the input’s shape with axes permuted. If the compiler is free to choose the physical order of both values, a transpose can be compiled to nothing: give the output the same physical byte order as the input and record the permutation in the value’s layout descriptor. Every consumer that reads the output then reads it through the new descriptor. This is the single most valuable thing layout assignment does, and it is why graphs with long transpose chains, which are common in attention implementations that split heads, can compile to kernels with no copies at all.

The catch is that the freedom is not universal. An operator’s implementation is written for a specific layout, or for a small set. A convolution kernel that expects the channel axis innermost cannot consume a channels-first tensor without repacking, and a matrix multiply that treats its operands as row-major cannot consume a column-major one without transposing the access pattern, which it may support natively through a flag or may not.

Propagation and the conversion frontier

The pass proceeds in two directions. A forward pass pushes a required layout from an operator’s inputs to its outputs, where the operator is layout-preserving. A backward pass pushes the constraint from consumers to producers, where the consumer is picky and the producer is flexible. Between them, constraints meet and become conflicts.

def resolve(node, ins, outs, rule):
    """ins, outs: list of (value, layout). rule: (accepted, preserves)."""
    accepted, preserves = rule[node.op]
    if preserves:
        # Elementwise: all operands and the result must agree on one layout.
        seen = [l for _, l in ins + outs if l in accepted]
        if not seen:
            return None                      # needs a conversion somewhere
        chosen = seen[0]
    else:
        # A picky node fixes the layout of every input it consumes.
        chosen = accepted[0]
    return {"conversions": [v for v, l in ins if l != chosen],
            "layout": chosen}

A conflict is resolved by inserting a conversion, which is a copy with a permuted read or write pattern. The interesting decision is where to put it. Inserting the conversion at the boundary of the picky operator costs one copy of that operator’s input. Inserting it earlier may allow a chain of elementwise operators to run entirely in the other layout and pay one copy instead of several. The standard heuristic anchors each connected component of flexible operators to the layout of its heaviest neighbour and converts only at the edges of the component.

Conversions are not free and not equal. A conversion that permutes the two outermost axes of a large tensor is a strided copy whose reads are far apart, and its bandwidth is a fraction of a contiguous copy’s. A conversion that only changes the innermost dimension, such as channels-first to channels-last on a $[N, C, H, W]$ tensor, is a transpose of a matrix of shape $[C, NHW]$ and can be done with a shared-memory tile at close to full bandwidth. A compiler that treats all conversions as equal cost will place them in the wrong spot.

Padding and the vector width

The second half of layout assignment is choosing padded extents. A kernel that processes the inner dimension with 128-bit vector loads wants the row length to be a multiple of four float32 elements. A tensor of shape $[1, 3, 224, 224]$ has an inner extent of 224, which is a multiple of eight and therefore of every common float32 vector width up to 256 bits; a tensor of shape $[1, 3, 225, 225]$ has an extent of 225, which is a multiple of none of them, so every kernel that touches it must handle a tail. Padding the inner extent up to the next multiple of the vector width costs fewer than that width’s worth of elements per row and removes the scalar tail loop from every kernel that reads the tensor.

The cost is that padded elements must be written with a defined value, usually zero, and that operators which reduce over the padded axis must not include them. A reduction over a padded axis that sums garbage produces garbage, and a max-reduction over padding filled with zero produces a wrong maximum for an all-negative row. Compilers handle this by keeping the logical extent in the descriptor and masking the tail in the kernel, which is more work than simply zero-filling and is the reason padding is applied only where the inner dimension is the vectorised one.

The alternative to padding is to choose a tile size that divides the extent, which is possible when the extent is known at compile time and awkward when it is not. This is one of the several places where dynamic shapes cost performance: a kernel compiled for a symbolic inner extent cannot assume alignment, so it must handle the tail on every call.

Weight layouts are decided once and stored

Activations have layouts chosen per graph. Weights have layouts chosen per model and baked into the artifact. A convolution weight of logical shape $[K, C, R, S]$ is stored in whatever order the chosen kernel wants, often with the output-channel axis grouped to match the accumulator tiling, and the conversion happens at export time rather than at load time. This is why a compiled model is tied to a particular kernel library’s expectations, and why a model compiled for one backend cannot simply be loaded by another: the weight bytes are in an order the other backend’s kernels do not read.

The consequence for a serving system is that a model artifact is not portable in the way a graph file is. A graph file plus a layout assignment produces an artifact; the artifact’s weight layout is the output of that assignment. Recompiling for a different device re-runs the assignment and produces different bytes from the same source weights.

/* NCHW -> NHWC for one image, tiled to keep writes contiguous. */
void nchw_to_nhwc(const float *src, float *dst, int C, int H, int W) {
    for (int c0 = 0; c0 < C; c0 += 8) {
        int cn = c0 + 8 <= C ? 8 : C - c0;
        for (int hw = 0; hw < H * W; hw++) {
            for (int c = 0; c < cn; c++)          /* 8 consecutive channels */
                dst[hw * C + c0 + c] = src[(c0 + c) * H * W + hw];
        }
    }
}

When the elimination fails

The transposes that survive layout assignment are the ones where both sides are picky and disagree. A convolution that wants channels-last feeding a matrix multiply that wants row-major is the common case, and it needs a real conversion. In attention, the reshape that splits a fused projection into per-head query, key, and value tensors is layout-preserving only if the head axis is adjacent to the feature axis in storage, which is exactly what the reshape’s semantics assume; a compiler that chose a different physical order for the projection output must insert a conversion before the split.

A second survivor is any transpose that crosses a graph break. The compiler cannot see the consumer, so it cannot choose a layout for the output that would suit it, and the safe choice is the layout the boundary expects, which may differ from the one the producer would prefer.

The pass also cannot fix a layout mismatch that is baked into the input. If the runtime is handed a channels-first image because that is what the preprocessing stage produces, the graph’s first conversion is unavoidable, and the only way to remove it is to change the producer of the data rather than the compiler. That is a boundary between the compiler and the surrounding system, and it is the same boundary that appears whenever the compiler stops being able to see the whole computation.