[Frontend] Build every tile descriptor from axes

[YWHyuk]

A pure refactor of how the ten MLIR templates describe their tiles. The emitted MLIR is byte-identical everywhere except one entry of one attribute, explained and justified below. No template constructs an MLIRMultiDimTile any more.

The problem

A tile descriptor is stored column-wise: one parallel array per space, plus a scalar index into them and a scalar riding alongside.

X_tile_size   = [TILE_I_H, TILE_I_W, TILE_M, TILE_K]
X_tile_stride = [TILE_I_W*TILE_M*TILE_K, TILE_M*TILE_K, 1, TILE_M]   # products of the extents, typed out
X_tile_desc   = MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)   # 3 = the lane axis
X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)         # the constructor arg was dead
X_tile_desc.set_name("input_buffer")                                 # must match the template text
X_dim = [Symbol("index_i_h"), Symbol("index_i_w"), Symbol("tile_m"), Symbol("tile_k")]
X_idx = [X_dim[0]*(I_W+2*PADDING_W)*BATCH*I_C, X_dim[1]*I_C*BATCH, X_dim[2]*I_C, X_dim[3]]

Keeping those columns aligned across an axis reorder, insert or collapse is the caller's job, and that is where the mistakes are:

• the GEMM and BMM reduction variants repeat the same if nr_rdim branch over four columns -- the extents, the SRAM strides, the lane axis and the index expression -- which can disagree;
• apply_divisor(mode="split") inserts an axis into _tile_size and tile_axis_order, bumps vlane_split_axis, and forgets tile_constraint (dormant: it has no callers);
• decompose_transfer re-indexes three arrays by hand and remaps the lane index separately;
• five templates constructed W_tile_desc with X_tile_size, harmless only because the next line overwrote it.

The change

Store the same table row-wise. One iteration dimension carries what it means in each space it is embedded in.

@dataclass(frozen=True)
class Axis:
    extent: int                    # how many elements of this axis the tile covers
    dram_stride: object = 0        # stride of the *access*, not of the tensor
    loop: Optional[str] = None     # the enclosing loop variable that advances it

What is not a property of a single axis stays on the tile: which axis rides the lanes, and the order the axes sit in SRAM.

X_tile_desc, X_idx = build_tile("input_buffer", kernel.vector_lane,
    axes={"i_h": Axis(TILE_I_H, (I_W+2*PADDING_W)*BATCH*I_C, loop="index_i_h"),
          "i_w": Axis(TILE_I_W, I_C*BATCH,                   loop="index_i_w"),
          "m":   Axis(TILE_M,   I_C,                         loop="tile_m"),
          "k":   Axis(TILE_K,   1,                           loop="tile_k")},
    sram_order=("i_h", "i_w", "k", "m"),   # outermost first; m is contiguous
    lane="k")

build_tile() derives the SRAM strides, the lane axis, the lane stride and the DRAM index expression. The strides are no longer typed out, so they cannot drift from the extents. The lane axis is named rather than positional. vlane_stride is 1 in every template, so it becomes a default nobody writes -- the word disappears from all ten.

Two orders, and they can differ

The axes' declared order is the memref's dimension order, which is what linalg sees. sram_order is the order they sit in SRAM. The GEMM and BMM reduction variants declare their output transposed, (N, M) instead of (M, N), while laying M out contiguously either way -- the bytes do not move, only the declaration. That was four separate branches; it is now one:

def y_axes(stride):
    m = Axis(TILE_M, stride[0], loop="index0")
    n = Axis(TILE_N, stride[1], loop="index1")
    return {"n": n, "m": m} if nr_rdim else {"m": m, "n": n}   # only the declaration flips

build_tile("Y_buffer", vector_lane, y_axes(Y_stride), sram_order=("n", "m"), lane="n")

The file tile_axis.py is 83 lines. Every template shrinks.

The DRAM stride is not the tensor's stride

Conv walks a padded, permuted logical layout, so its per-axis DRAM stride is a sympy expression, not layout.stride[dim]. Axis carries the expression directly.

Facts the refactor surfaced

• conv_mt walks one kernel column at a time, so its k_w axis is degenerate. That was a bare 1 in a tile-size list; it is now an axis that says so.
• sdpa's key tile is laid out transposed because key is the stationary operand of the systolic array. That was a stride list and a lane index sitting apart; it is now one sram_order and one lane, with the reason next to them.
• sort's lane axis can have extent 1: a one-dimensional sort runs in a single lane.

The one attribute that changes

The templates disagreed on the SRAM stride of an extent-1 axis. conv writes the product it would have if it were not degenerate; BMM and sdpa write 0. Deriving it from sram_order gives conv's answer, so four BMM tile_stride entries and four sdpa ones change from 0 to that product.

It never reaches an address. Spike (torchsim_mvin_common.h) bounds that axis' loop by its extent and multiplies its stride by the index:

uint64_t block_dim[4] = {p_dim_size[N], p_dim_size[C], p_dim_size[H], p_dim_size[W]};
block_dim[vlane_split_axis] = vlane_stride;
...
for (n = 0; n < block_dim[N]; n++) for (c = 0; c < block_dim[C]; c++) ...
    s_idx = s_outerloop_idx_stride * outerloop_idx
          + block_stride[N]*n + block_stride[C]*c + block_stride[H]*h + block_stride[W]*w;

For a non-lane axis of extent 1 the loop runs once with the index at 0. Each block_stride[i] is rescaled against the lane axis' one independently, so it does not perturb the others; buffer_size and dma_buffer_stride come from dim_size, not from the strides.

Compiling BMM both ways confirms it: the LLVM instruction body is byte-identical, and the two differ in exactly one slot of the dma_desc global -- the SRAM stride of a dimension whose extent is 1.

@dma_desc_0  slot 11:  0 -> 32768
@dma_desc_1  slot 11:  0 -> 12288
@dma_desc_2  slot 11:  0 -> 24576
body identical: True

Verification

Every kernel was regenerated from scratch (kernel directories deleted first, so the content-hash path could not mask a change).

	result
GEMM: mm, mm+relu, addmm, mm+reduction, prologue-fused mm, N==1 edge	byte-identical MLIR
conv: single-batch, single-batch-strided, multi-tile, batched (9 kernels)	byte-identical MLIR
maxpool, cat, sort	byte-identical MLIR
BMM (4), sdpa	only the tile_stride entry above
functional (Spike) vs CPU: GEMM 6, BMM 4, conv 4	pass, max abs diff 4.2e-05
functional: max_pool2d, cat, sort, topk	exact
functional: tests/ops/attention/test_sdpa.py	192 cases, 0 failures
tests/ops/fusion/ (7 tests)	fuse the same kernels

What this does not do

MLIRMultiDimTile is unchanged; build_tile() is an adapter over it. The parallel arrays, the vmap scalars, the tile_axis_order float ranks, apply_divisor's missing tile_constraint update, decompose_transfer's hand re-indexing and compute_tile_size's eight copies of vmap.vlane_stride all survive.

That is deliberate. The callers had to stop passing derived values before the internals could change, and the byte-identity above is only meaningful because nothing else moved. The internals are the next step.