From 04b0669e0d7cabdc6fe128d4df73bfd9cdeb11e4 Mon Sep 17 00:00:00 2001
From: Wonhyuk Yang <wonhyuk@postech.ac.kr>
Date: Fri, 10 Jul 2026 17:38:42 +0900
Subject: [PATCH 1/4] [Frontend] Introduce Axis: one tile dimension, and what
 it means in each space

A tile descriptor is stored column-wise today: one parallel array per space (the
extents, the DRAM strides, the SRAM strides), a scalar index into them
(vlane_split_axis), and a scalar riding alongside (vlane_stride). Keeping the
columns aligned across an axis reorder, insert or collapse is the caller's job,
and that is where the mistakes are. The reduction GEMM repeats the same
"if nr_rdim" branch over four of them. apply_divisor inserts an axis but forgets
tile_constraint. decompose_transfer re-indexes three arrays by hand and remaps
the lane index separately.

Axis stores the same table row-wise. One iteration dimension carries its extent,
the stride of the DRAM access it walks, and 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.

Two orders matter 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, outermost first. The GEMM reduction variant declares its output (N, M)
while still laying M out contiguously; today that is four separate branches over
the extents, the SRAM strides, the lane axis and the index expression.

The DRAM stride is the stride of the access, not of the tensor. Conv walks a
padded, permuted layout, so it is a sympy expression, not a layout stride.

build_tile() derives from that what the templates compute by hand: the SRAM
strides, the lane axis, the lane stride, and the DRAM index expression.

An extent-1 axis is indexed only at 0, so its SRAM stride never reaches an
address: Spike bounds that axis' loop by its extent and multiplies the stride by
the index (torchsim_mvin_common.h), and each stride is rescaled against the lane
axis' one independently, so it does not perturb the others. Deriving it from the
SRAM order gives the product it would have if it were not degenerate, which is
what most templates already write by hand.

No caller yet.
---
 PyTorchSimFrontend/mlir/tile_axis.py | 83 ++++++++++++++++++++++++++++
 1 file changed, 83 insertions(+)
 create mode 100644 PyTorchSimFrontend/mlir/tile_axis.py

diff --git a/PyTorchSimFrontend/mlir/tile_axis.py b/PyTorchSimFrontend/mlir/tile_axis.py
new file mode 100644
index 00000000..78b1e2a0
--- /dev/null
+++ b/PyTorchSimFrontend/mlir/tile_axis.py
@@ -0,0 +1,83 @@
+"""One axis of a tile, and everything that axis means in the spaces it is embedded in.
+
+Today a tile descriptor is stored column-wise: one parallel array per space (the tile
+extents, the DRAM strides, the SRAM strides), a scalar index into them
+(vlane_split_axis) and a scalar riding alongside (vlane_stride). Keeping the columns
+aligned across an axis reorder, insert or collapse is the caller's job, and that is
+where the mistakes are: the reduction GEMM repeats the same `if nr_rdim` branch over
+four of them, `apply_divisor` inserts an axis but forgets `tile_constraint`,
+`decompose_transfer` re-indexes three arrays by hand and remaps the lane index
+separately.
+
+`Axis` stores the same table row-wise: one iteration dimension, carrying what it means
+in DRAM and in the enclosing loop nest. 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.
+"""
+from dataclasses import dataclass
+from typing import Optional
+
+import sympy
+
+from PyTorchSimFrontend.mlir import mlir_common
+
+
+@dataclass(frozen=True)
+class Axis:
+    """One iteration dimension of a tile.
+
+    extent       how many elements of this axis the tile covers
+    dram_stride  distance in DRAM between two neighbours along this axis. This is the
+                 stride of the *access*, not of the tensor: conv walks a padded logical
+                 layout, so it is an int or a sympy expression, not a layout stride.
+    loop         the enclosing loop variable that advances this axis, one tile at a
+                 time; None when the axis does not move in DRAM
+    """
+    extent: int
+    dram_stride: object = 0
+    loop: Optional[str] = None
+
+
+def sram_strides(axes, sram_order):
+    """SRAM strides, in the axes' declared order.
+
+    `sram_order` lists the axis names outermost first, so the last one is contiguous.
+    An extent-1 axis is indexed only at 0, so its stride never reaches an address --
+    Spike bounds that axis' loop by its extent -- but it still gets the stride it would
+    have if it were not degenerate.
+    """
+    stride, init = {}, 1
+    for name in reversed(sram_order):
+        stride[name] = init
+        init *= axes[name].extent
+    return [stride[name] for name in axes]
+
+
+def dram_index(axes):
+    """The tile's DRAM offset, one term per axis, in the axes' declared order."""
+    return [sympy.Integer(0) if a.loop is None else sympy.Symbol(a.loop) * a.dram_stride
+            for a in axes.values()]
+
+
+def build_tile(buffer, vector_lane, axes, sram_order, lane, lane_chunk=1, offset=0):
+    """Build the tile descriptor and the DRAM index expression for one operand.
+
+    `buffer` is the SRAM buffer's name, as the template text spells it. `axes` is an
+    ordered mapping name -> Axis; its order is the memref's dimension order, which is
+    what linalg sees. `sram_order` is the order those axes sit in SRAM, outermost first.
+    The two differ whenever the tile is declared transposed -- the GEMM reduction variant
+    declares (N, M) but still lays M out contiguously.
+
+    The SRAM strides, the lane axis, the lane stride and the DRAM index expression are
+    all derived from that.
+    """
+    assert set(sram_order) == set(axes), f"{buffer}: sram_order does not cover the axes"
+    assert lane in axes, f"{buffer}: lane axis {lane!r} is not an axis"
+
+    names = list(axes)
+    extents = [axes[n].extent for n in names]
+    desc = mlir_common.MLIRMultiDimTile(extents, vector_lane, names.index(lane), lane_chunk)
+    desc.set_tile_size_stride(extents, sram_strides(axes, sram_order))
+    desc.set_name(buffer)
+    desc.offset = offset
+
+    return desc, dram_index(axes)

From cf088ec955e64db2ad5cb781bb1d0b47ddfe64d1 Mon Sep 17 00:00:00 2001
From: Wonhyuk Yang <wonhyuk@postech.ac.kr>
Date: Fri, 10 Jul 2026 17:38:42 +0900
Subject: [PATCH 2/4] [Frontend] Build the GEMM, BMM and conv tile descriptors
 from axes

Each operand took eight lines: the tile extents, the SRAM strides worked out by
hand as products of those extents, a constructor whose tile_size argument the
next line overwrites, the lane axis as a bare integer, the buffer name as a
string, and an index expression built from a list of loop symbols. Half of that
is derived from the rest.

State the decisions -- which extents, which DRAM stride each axis walks, which
loop variable advances it, what order the axes sit in SRAM, which one rides the
lanes -- and derive the rest. The SRAM 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 GEMM and BMM reduction variants declare their output tile transposed, (N, M)
instead of (M, N), while laying M out contiguously either way. That used to be
four separate "if nr_rdim" branches over the extents, the SRAM strides, the lane
axis and the index expression, which could disagree. Now the declaration order
flips and the SRAM order stays (n, m).

This also kills a copy-paste that survived because the API was redundant: five
templates built W_tile_desc with X_tile_size, harmless only because the next
line overwrote it.

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.

Verified by regenerating every kernel from scratch. The emitted MLIR is
byte-identical for mm, mm+relu, addmm, mm+reduction, prologue-fused mm, the N==1
edge, and all four conv variants (single-batch, single-batch-strided, multi-tile,
batched). BMM's degenerate batch axis picks up the SRAM stride it would have if
it were not degenerate, so its transfers differ in that one entry; compiling both
ways gives a byte-identical instruction body and differs in one slot of the DMA
descriptor global. Functional mode matches CPU for all fourteen (max abs diff
4.2e-05), and tests/ops/fusion still fuses the same kernels.
---
 PyTorchSimFrontend/mlir/mlir_bmm_template.py  | 79 +++++++++----------
 .../mlir/mlir_conv_mt_template.py             | 67 ++++++++--------
 .../mlir/mlir_conv_sb_template.py             | 66 +++++++++-------
 .../mlir/mlir_conv_sbs_template.py            | 66 +++++++++-------
 PyTorchSimFrontend/mlir/mlir_conv_template.py | 68 ++++++++--------
 PyTorchSimFrontend/mlir/mlir_gemm_template.py | 71 ++++++++---------
 6 files changed, 211 insertions(+), 206 deletions(-)

diff --git a/PyTorchSimFrontend/mlir/mlir_bmm_template.py b/PyTorchSimFrontend/mlir/mlir_bmm_template.py
index 2bc8cff8..de97b1e6 100644
--- a/PyTorchSimFrontend/mlir/mlir_bmm_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_bmm_template.py
@@ -7,6 +7,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 
 BMM_TEMPLATE = r"""
 // BMM kernel
@@ -192,53 +193,45 @@ def render(self,
             epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
             nr_rdim = 0
 
-        # Prepare tile descriptors
-        vlane_stride = 1
-        vlane_split_axis = 2
-        loop_dim = [sympy.Symbol("index0"), sympy.Symbol("index1"), sympy.Symbol("index2"), sympy.Symbol("index3")]
-        X_tile_size = [1, TILE_M, TILE_K]
-        X_tile_stride = [0, 1, TILE_M]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("X_buffer")
-        X_tile_desc.offset = X.get_layout().offset
-        X_stride = X_tensor.stride()
-        X_idx = [loop_dim[0]*X_stride[0], loop_dim[1]*X_stride[1], loop_dim[3]*X_stride[2]] # To keep index arguemnt order, we used index_list
-
-        W_tile_size = [1, TILE_K, TILE_N]
-        W_tile_stride = [0, 1, TILE_K]
-        W_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        W_tile_desc.set_tile_size_stride(W_tile_size, W_tile_stride)
-        W_tile_desc.set_name("W_buffer")
-        W_tile_desc.offset = W.get_layout().offset
-        W_stride = W_tensor.stride()
-        W_idx = [loop_dim[0]*W_stride[0], loop_dim[3]*W_stride[1], loop_dim[2]*W_stride[2]]
-
-        vlane_split_axis = vlane_split_axis if nr_rdim==0 else 1
-        Y_tile_size = [1, TILE_M, TILE_N] if nr_rdim == 0 else [1, TILE_N, TILE_M]
-        Y_tile_stride=[0, 1, TILE_M] if nr_rdim == 0 else [0, TILE_M, 1]
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("Y_buffer")
+        # Prepare tile descriptors. Batch is the outermost SRAM axis and degenerate (one
+        # slice per tile), N rides the lanes and M is contiguous inside one lane.
+        X_stride, W_stride = X_tensor.stride(), W_tensor.stride()
         Y_stride = Y.get_layout().stride
-        if nr_rdim == 0:
-          Y_idx = [loop_dim[0]*Y_stride[0], loop_dim[1]*Y_stride[1], loop_dim[2]*Y_stride[2]]
-        else:
-          Y_idx = [loop_dim[0]*Y_stride[0], loop_dim[2]*Y_stride[2], loop_dim[1]*Y_stride[1]]
 
-        # Extract Bias info
-        Bias_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Bias_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Bias_tile_desc.set_name("Y_buffer")
+        X_tile_desc, X_idx = build_tile(
+            "X_buffer", kernel.vector_lane,
+            axes={"b": Axis(1,      X_stride[0], loop="index0"),
+                  "m": Axis(TILE_M, X_stride[1], loop="index1"),
+                  "k": Axis(TILE_K, X_stride[2], loop="index3")},
+            sram_order=("b", "k", "m"), lane="k", offset=X.get_layout().offset)
+
+        W_tile_desc, W_idx = build_tile(
+            "W_buffer", kernel.vector_lane,
+            axes={"b": Axis(1,      W_stride[0], loop="index0"),
+                  "k": Axis(TILE_K, W_stride[1], loop="index3"),
+                  "n": Axis(TILE_N, W_stride[2], loop="index2")},
+            sram_order=("b", "n", "k"), lane="n", offset=W.get_layout().offset)
+
+        # The reduction template sweeps N outside M, so its tile is declared (B, N, M).
+        # Only the declaration flips; the SRAM order is (b, n, m) either way.
+        def y_axes(stride):
+            b = Axis(1,      stride[0], loop="index0")
+            m = Axis(TILE_M, stride[1], loop="index1")
+            n = Axis(TILE_N, stride[2], loop="index2")
+            return {"b": b, "n": n, "m": m} if nr_rdim else {"b": b, "m": m, "n": n}
+
+        Y_tile_desc, Y_idx = build_tile(
+            "Y_buffer", kernel.vector_lane, y_axes(Y_stride), sram_order=("b", "n", "m"), lane="n")
+
+        # Extract Bias info. It accumulates into the Y buffer, so it shares Y's axes.
         if Bias is not None:
-          Bias_stride = Bias.get_layout().stride
-          Bias_tile_desc.offset = Bias.get_layout().offset
-          if nr_rdim == 0:
-            Bias_idx = [loop_dim[0]*Bias_stride[0], loop_dim[1]*Bias_stride[1], loop_dim[2]*Bias_stride[2]]
-          else:
-            Bias_idx = [loop_dim[0]*Bias_stride[0], loop_dim[2]*Bias_stride[2], loop_dim[1]*Bias_stride[1]]
+            Bias_tile_desc, Bias_idx = build_tile(
+                "Y_buffer", kernel.vector_lane, y_axes(Bias.get_layout().stride),
+                sram_order=("b", "n", "m"), lane="n", offset=Bias.get_layout().offset)
         else:
-          Bias_idx = None
+            Bias_tile_desc, _ = build_tile(
+                "Y_buffer", kernel.vector_lane, y_axes(Y_stride), sram_order=("b", "n", "m"), lane="n")
+            Bias_idx = None
 
         data_stype = mlir_common.DTYPE_TO_MLIR[X.get_dtype()]
         kernel.render_options = dict(
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
index d5920fd0..9ac8a518 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
@@ -5,6 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 
 CONV_TEMPLATE = r"""
 // Multi Channel Tile Conv2D kernel
@@ -151,40 +152,44 @@ def render(self,
         kernel.loop_extents = {"tile_m": BATCH, "tile_n": O_C, "o_h": O_H, "o_w": O_W,
                                "k_h": K_H, "tile_k": I_C * K_W}
 
-        # Prepare tile descriptors
-        vlane_stride = 1
-        vlane_split_axis = 1
-        X_tile_size = [TILE_I_H, TILE_O_W, TILE_M, TILE_K]
-        X_tile_stride = [TILE_O_W*TILE_M*TILE_K, TILE_M*TILE_K, 1, TILE_M]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("input_buffer")
-        X_dim = [Symbol("index_i_h"), Symbol("o_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*STRIDE_W, X_dim[2]*I_C*(I_W+2*PADDING_W), X_dim[3]]
+        # Prepare tile descriptors. The channel axis rides the lanes; the DRAM strides walk
+        # the padded, permuted layout, so they are expressions rather than tensor strides.
+        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"),
+                  "o_w": Axis(TILE_O_W, I_C*STRIDE_W,                loop="o_w"),
+                  "m":   Axis(TILE_M,   I_C*(I_W+2*PADDING_W),       loop="tile_m"),
+                  "k":   Axis(TILE_K,   1,                           loop="tile_k")},
+            sram_order=("i_h", "o_w", "k", "m"), lane="k")
 
-        W_tile_size = [TILE_K_H, 1, TILE_K, TILE_N]
-        W_tile_stride = [TILE_K * TILE_N, TILE_K * TILE_N, 1, TILE_K]
-        W_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        W_tile_desc.set_tile_size_stride(W_tile_size, W_tile_stride)
-        W_tile_desc.set_name("weight_buffer")
-        W_dim = [Symbol("k_h"), Symbol("k_w"), Symbol("tile_k"), Symbol("tile_n")]
-        W_idx = [W_dim[0]*K_W*I_C*O_C , Symbol("c0"), W_dim[2]*O_C, W_dim[3]]
+        # This kernel walks one kernel column at a time, so k_w is degenerate here.
+        W_tile_desc, W_idx = build_tile(
+            "weight_buffer", kernel.vector_lane,
+            axes={"k_h": Axis(TILE_K_H, K_W*I_C*O_C, loop="k_h"),
+                  "k_w": Axis(1,        1,           loop="c0"),
+                  "k":   Axis(TILE_K,   O_C,         loop="tile_k"),
+                  "n":   Axis(TILE_N,   1,           loop="tile_n")},
+            sram_order=("k_h", "k_w", "n", "k"), lane="n")
 
-        Y_tile_size = [TILE_M, TILE_N, TILE_O_H, TILE_O_W]
-        Y_tile_stride = [1, TILE_M, TILE_O_W * TILE_M * TILE_N, TILE_M * TILE_N] # N, C, H, W
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("output_buffer")
-        Y_dim = [Symbol("tile_m"), Symbol("tile_n"), Symbol("o_h"), Symbol("o_w")]
-        Y_idx = [Y_dim[0]*O_C*O_H*O_W, Y_dim[1]*O_H*O_W, Y_dim[2]*O_W, Y_dim[3]]
+        # N, C, H, W
+        def y_axes(m_stride, n_stride, h_stride, w_stride, loops):
+            return {"m":   Axis(TILE_M,   m_stride, loop=loops[0]),
+                    "n":   Axis(TILE_N,   n_stride, loop=loops[1]),
+                    "o_h": Axis(TILE_O_H, h_stride, loop=loops[2]),
+                    "o_w": Axis(TILE_O_W, w_stride, loop=loops[3])}
 
-        # Extract Bias info
-        Bias_idx = [Number(0), Symbol("tile_n"), Number(0), Number(0)]
-        Bias_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Bias_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Bias_tile_desc.set_name("output_buffer")
-        if Bias is not None:
-          Bias_tile_desc.offset = Bias.get_layout().offset
+        Y_SRAM_ORDER = ("o_h", "o_w", "n", "m")
+        Y_tile_desc, Y_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(O_C*O_H*O_W, O_H*O_W, O_W, 1, ["tile_m", "tile_n", "o_h", "o_w"]),
+            sram_order=Y_SRAM_ORDER, lane="n")
+
+        # Extract Bias info. It accumulates into the output buffer, and only walks channels.
+        Bias_tile_desc, Bias_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(0, 1, 0, 0, [None, "tile_n", None, None]),
+            sram_order=Y_SRAM_ORDER, lane="n",
+            offset=Bias.get_layout().offset if Bias is not None else 0)
 
         data_stype = mlir_common.DTYPE_TO_MLIR[X.get_dtype()]
 
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
index c5c4843a..cf61f3fa 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
@@ -5,6 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 
 CONV_TEMPLATE = r"""
 // Single Batch Conv2D kernel
@@ -150,39 +151,44 @@ def render(self,
         # Real extent of each structural loop iv, for the masked-DMA clamp (def_dma_op).
         kernel.loop_extents = {"tile_n": O_C, "o_h": O_H, "tile_m": O_W,
                                "k_h": K_H, "k_w": K_W, "tile_k": I_C}
-        # Prepare tile descriptors
-        vlane_stride = 1
-        vlane_split_axis = 1
-        X_tile_size = [1, TILE_I_H, TILE_I_W, TILE_K]
-        X_tile_stride = [TILE_I_H * TILE_I_W * TILE_K , TILE_I_W * TILE_K, 1, TILE_I_W]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("input_buffer")
-        X_dim = [Symbol("c0"), Symbol("index_i_h"), Symbol("index_i_w"), Symbol("tile_k")]
-        X_idx = [X_dim[0]*((I_W+2*PADDING_W)*(I_H+2*PADDING_H)*I_C), X_dim[1]*((I_W+2*PADDING_W)*I_C), X_dim[2]*I_C, X_dim[3]]
+        # Prepare tile descriptors. This kernel handles one image, so the batch axis is
+        # degenerate. The channel axis rides the lanes; the DRAM strides walk the padded,
+        # permuted layout, so they are expressions rather than tensor strides.
+        X_tile_desc, X_idx = build_tile(
+            "input_buffer", kernel.vector_lane,
+            axes={"b":   Axis(1,         (I_W+2*PADDING_W)*(I_H+2*PADDING_H)*I_C, loop="c0"),
+                  "i_h": Axis(TILE_I_H,  (I_W+2*PADDING_W)*I_C, loop="index_i_h"),
+                  "i_w": Axis(TILE_I_W,  I_C,                   loop="index_i_w"),
+                  "k":   Axis(TILE_K,    1,                     loop="tile_k")},
+            sram_order=("b", "i_h", "k", "i_w"), lane="k")
 
-        W_tile_size = [TILE_K_H, TILE_K_W, TILE_K, TILE_N]
-        W_tile_stride = [TILE_K_W * TILE_K * TILE_N, TILE_K * TILE_N, 1, TILE_K]
-        W_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        W_tile_desc.set_tile_size_stride(W_tile_size, W_tile_stride)
-        W_tile_desc.set_name("weight_buffer")
-        W_dim = [Symbol("k_h"), Symbol("k_w"), Symbol("tile_k"), Symbol("tile_n")]
-        W_idx = [W_dim[0]*K_W*I_C*O_C , W_dim[1]*I_C*O_C, W_dim[2]*O_C, W_dim[3]]
+        W_tile_desc, W_idx = build_tile(
+            "weight_buffer", kernel.vector_lane,
+            axes={"k_h": Axis(TILE_K_H, K_W*I_C*O_C, loop="k_h"),
+                  "k_w": Axis(TILE_K_W, I_C*O_C,     loop="k_w"),
+                  "k":   Axis(TILE_K,   O_C,         loop="tile_k"),
+                  "n":   Axis(TILE_N,   1,           loop="tile_n")},
+            sram_order=("k_h", "k_w", "n", "k"), lane="n")
 
-        Y_tile_size = [1, TILE_N, TILE_O_H, TILE_M]
-        Y_tile_stride = [TILE_O_H * TILE_M * TILE_N, TILE_M, TILE_M * TILE_N, 1] # N, C, H, W
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("output_buffer")
-        Y_idx = [Number(0), Symbol("tile_n")*O_H*O_W, Symbol("o_h")*O_W, Symbol("tile_m")]
+        # N, C, H, W
+        def y_axes(b_stride, n_stride, h_stride, m_stride, loops):
+            return {"b":   Axis(1,         b_stride, loop=loops[0]),
+                    "n":   Axis(TILE_N,    n_stride, loop=loops[1]),
+                    "o_h": Axis(TILE_O_H,  h_stride, loop=loops[2]),
+                    "m":   Axis(TILE_M,    m_stride, loop=loops[3])}
 
-        # Extract Bias info
-        Bias_idx = [Number(0), Symbol("tile_n"), Number(0), Number(0)]
-        Bias_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Bias_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Bias_tile_desc.set_name("output_buffer")
-        if Bias is not None:
-          Bias_tile_desc.offset = Bias.get_layout().offset
+        Y_SRAM_ORDER = ("b", "o_h", "n", "m")
+        Y_tile_desc, Y_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(0, O_H*O_W, O_W, 1, [None, "tile_n", "o_h", "tile_m"]),
+            sram_order=Y_SRAM_ORDER, lane="n")
+
+        # Extract Bias info. It accumulates into the output buffer, and only walks channels.
+        Bias_tile_desc, Bias_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(0, 1, 0, 0, [None, "tile_n", None, None]),
+            sram_order=Y_SRAM_ORDER, lane="n",
+            offset=Bias.get_layout().offset if Bias is not None else 0)
 
         data_stype = mlir_common.DTYPE_TO_MLIR[X.get_dtype()]
 
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
index e87eb0fe..4c471006 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
@@ -5,6 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 
 CONV_TEMPLATE = r"""
 // Single Batch Conv2D (Stride != 1) kernel
@@ -151,39 +152,44 @@ def render(self,
         kernel.loop_extents = {"tile_n": O_C, "o_h": O_H, "tile_m": O_W,
                                "k_h": K_H, "k_w": K_W, "tile_k": I_C}
 
-        # Prepare tile descriptors
-        vlane_stride = 1
-        vlane_split_axis = 1
-        X_tile_size = [TILE_I_H, TILE_K_W, TILE_M, TILE_K]
-        X_tile_stride = [TILE_K_W*TILE_M*TILE_K, TILE_M*TILE_K, 1, TILE_M]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("input_buffer")
-        X_dim = [Symbol("index_i_h"), Symbol("k_w"), Symbol("tile_m"), Symbol("tile_k")]
-        X_idx = [X_dim[0]*((I_W+2*PADDING_W)*I_C), X_dim[1]*I_C, X_dim[2]*(I_C*STRIDE_W), X_dim[3]]
+        # Prepare tile descriptors. This kernel handles one image, so the output's batch axis
+        # is degenerate. The channel axis rides the lanes; the DRAM strides walk the padded,
+        # permuted layout, so they are expressions rather than tensor strides.
+        X_tile_desc, X_idx = build_tile(
+            "input_buffer", kernel.vector_lane,
+            axes={"i_h": Axis(TILE_I_H, (I_W+2*PADDING_W)*I_C, loop="index_i_h"),
+                  "k_w": Axis(TILE_K_W, I_C,                   loop="k_w"),
+                  "m":   Axis(TILE_M,   I_C*STRIDE_W,          loop="tile_m"),
+                  "k":   Axis(TILE_K,   1,                     loop="tile_k")},
+            sram_order=("i_h", "k_w", "k", "m"), lane="k")
 
-        W_tile_size = [TILE_K_H, TILE_K_W, TILE_K, TILE_N]
-        W_tile_stride = [TILE_K_W * TILE_K * TILE_N, TILE_K * TILE_N, 1, TILE_K]
-        W_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        W_tile_desc.set_tile_size_stride(W_tile_size, W_tile_stride)
-        W_tile_desc.set_name("weight_buffer")
-        W_dim = [Symbol("k_h"), Symbol("k_w"), Symbol("tile_k"), Symbol("tile_n")]
-        W_idx = [W_dim[0]*K_W*I_C*O_C , W_dim[1]*I_C*O_C, W_dim[2]*O_C, W_dim[3]]
+        W_tile_desc, W_idx = build_tile(
+            "weight_buffer", kernel.vector_lane,
+            axes={"k_h": Axis(TILE_K_H, K_W*I_C*O_C, loop="k_h"),
+                  "k_w": Axis(TILE_K_W, I_C*O_C,     loop="k_w"),
+                  "k":   Axis(TILE_K,   O_C,         loop="tile_k"),
+                  "n":   Axis(TILE_N,   1,           loop="tile_n")},
+            sram_order=("k_h", "k_w", "n", "k"), lane="n")
 
-        Y_tile_size = [1, TILE_N, TILE_O_H, TILE_M]
-        Y_tile_stride = [TILE_O_H * TILE_M * TILE_N, TILE_M, TILE_M * TILE_N, 1] # N, C, H, W
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("output_buffer")
-        Y_idx = [Number(0), Symbol("tile_n")*O_H*O_W, Symbol("o_h")*O_W, Symbol("tile_m")]
+        # N, C, H, W
+        def y_axes(b_stride, n_stride, h_stride, m_stride, loops):
+            return {"b":   Axis(1,        b_stride, loop=loops[0]),
+                    "n":   Axis(TILE_N,   n_stride, loop=loops[1]),
+                    "o_h": Axis(TILE_O_H, h_stride, loop=loops[2]),
+                    "m":   Axis(TILE_M,   m_stride, loop=loops[3])}
 
-        # Extract Bias info
-        Bias_idx = [Number(0), Symbol("tile_n"), Number(0), Number(0)]
-        Bias_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Bias_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Bias_tile_desc.set_name("output_buffer")
-        if Bias is not None:
-          Bias_tile_desc.offset = Bias.get_layout().offset
+        Y_SRAM_ORDER = ("b", "o_h", "n", "m")
+        Y_tile_desc, Y_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(0, O_H*O_W, O_W, 1, [None, "tile_n", "o_h", "tile_m"]),
+            sram_order=Y_SRAM_ORDER, lane="n")
+
+        # Extract Bias info. It accumulates into the output buffer, and only walks channels.
+        Bias_tile_desc, Bias_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(0, 1, 0, 0, [None, "tile_n", None, None]),
+            sram_order=Y_SRAM_ORDER, lane="n",
+            offset=Bias.get_layout().offset if Bias is not None else 0)
 
         data_stype = mlir_common.DTYPE_TO_MLIR[X.get_dtype()]
 
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_template.py b/PyTorchSimFrontend/mlir/mlir_conv_template.py
index 7b1b0856..aa564b43 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_template.py
@@ -5,6 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 
 CONV_TEMPLATE = r"""
 // Conv2D kernel
@@ -154,40 +155,43 @@ def render(self,
         kernel.loop_extents = {"tile_m": BATCH, "tile_n": O_C, "o_h": O_H, "o_w": O_W,
                                "k_h": K_H, "k_w": K_W, "tile_k": I_C}
 
-        # Prepare tile descriptors
-        vlane_stride = 1
-        vlane_split_axis = 1
-        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]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("input_buffer")
-        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]]
+        # Prepare tile descriptors. The channel axis rides the lanes; the DRAM strides walk
+        # the padded, permuted layout, so they are expressions rather than tensor strides.
+        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"), lane="k")
 
-        W_tile_size = [TILE_K_H, TILE_K_W, TILE_K, TILE_N]
-        W_tile_stride = [TILE_K_W * TILE_K * TILE_N, TILE_K * TILE_N, 1, TILE_K]
-        W_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, 3, vlane_stride)
-        W_tile_desc.set_tile_size_stride(W_tile_size, W_tile_stride)
-        W_tile_desc.set_name("weight_buffer")
-        W_dim = [Symbol("k_h"), Symbol("k_w"), Symbol("tile_k"), Symbol("tile_n")]
-        W_idx = [W_dim[0]*K_W*I_C*O_C , W_dim[1]*I_C*O_C, W_dim[2]*O_C, W_dim[3]]
+        W_tile_desc, W_idx = build_tile(
+            "weight_buffer", kernel.vector_lane,
+            axes={"k_h": Axis(TILE_K_H, K_W*I_C*O_C, loop="k_h"),
+                  "k_w": Axis(TILE_K_W, I_C*O_C,     loop="k_w"),
+                  "k":   Axis(TILE_K,   O_C,         loop="tile_k"),
+                  "n":   Axis(TILE_N,   1,           loop="tile_n")},
+            sram_order=("k_h", "k_w", "n", "k"), lane="n")
 
-        Y_tile_size = [TILE_M, TILE_N, TILE_O_H, TILE_O_W]
-        Y_tile_stride = [1, TILE_M, TILE_O_W * TILE_M * TILE_N, TILE_M * TILE_N] # N, C, H, W
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("output_buffer")
-        Y_dim = [Symbol("tile_m"), Symbol("tile_n"), Symbol("o_h"), Symbol("o_w")]
-        Y_idx = [Y_dim[0]*O_C*O_H*O_W, Y_dim[1]*O_H*O_W, Y_dim[2]*O_W, Y_dim[3]]
-        
-        # Extract Bias info
-        Bias_idx = [Number(0), Symbol("tile_n"), Number(0), Number(0)]
-        Bias_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Bias_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Bias_tile_desc.set_name("output_buffer")
-        if Bias is not None:
-          Bias_tile_desc.offset = Bias.get_layout().offset
+        # N, C, H, W
+        def y_axes(m_stride, n_stride, h_stride, w_stride, loops):
+            return {"m":   Axis(TILE_M,   m_stride, loop=loops[0]),
+                    "n":   Axis(TILE_N,   n_stride, loop=loops[1]),
+                    "o_h": Axis(TILE_O_H, h_stride, loop=loops[2]),
+                    "o_w": Axis(TILE_O_W, w_stride, loop=loops[3])}
+
+        Y_SRAM_ORDER = ("o_h", "o_w", "n", "m")
+        Y_tile_desc, Y_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(O_C*O_H*O_W, O_H*O_W, O_W, 1, ["tile_m", "tile_n", "o_h", "o_w"]),
+            sram_order=Y_SRAM_ORDER, lane="n")
+
+        # Extract Bias info. It accumulates into the output buffer, and only walks channels.
+        Bias_tile_desc, Bias_idx = build_tile(
+            "output_buffer", kernel.vector_lane,
+            y_axes(0, 1, 0, 0, [None, "tile_n", None, None]),
+            sram_order=Y_SRAM_ORDER, lane="n",
+            offset=Bias.get_layout().offset if Bias is not None else 0)
 
         data_stype = mlir_common.DTYPE_TO_MLIR[X.get_dtype()]
 
diff --git a/PyTorchSimFrontend/mlir/mlir_gemm_template.py b/PyTorchSimFrontend/mlir/mlir_gemm_template.py
index 6fb2b9bc..94a7654e 100644
--- a/PyTorchSimFrontend/mlir/mlir_gemm_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_gemm_template.py
@@ -9,6 +9,7 @@
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend import extension_config
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 
 GEMM_TEMPLATE = r"""
 // GEMM {% if prologue_nodes -%}prologue fused{%- endif %} {% if epilogue_nodes -%}eilogue fused{%- endif %} kernel
@@ -142,53 +143,43 @@ def render(self,
         TOG_latency = M if SUB_TILE_M > M else SUB_TILE_M
         kernel.loop_size =[TOG_latency, SUB_TILE_N, SUB_TILE_K]
 
-        # Prepare tile descriptors
-        vlane_stride = 1
-        vlane_split_axis = 1
-        X_tile_size = [TILE_M, TILE_K]
-        X_tile_stride = [1, TILE_M]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("X_buffer")
-        X_tile_desc.offset = X.get_layout().offset
+        # Prepare tile descriptors. N rides the lanes, M is contiguous inside one lane.
         X_stride = X.get_layout().stride
-        X_idx = [sympy.Symbol("index0") * X_stride[0], sympy.Symbol("index2") * X_stride[1]] # To keep index arguemnt order, we used index_list
-
-        W_tile_size = [TILE_K, TILE_N]
-        W_tile_stride = [1, TILE_K]
-        W_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        W_tile_desc.set_tile_size_stride(W_tile_size, W_tile_stride)
-        W_tile_desc.set_name("W_buffer")
-        W_tile_desc.offset = W.get_layout().offset
         W_stride = W.get_layout().stride if N>1 else [Y.get_layout().stride[0], 0]
-        W_idx = [sympy.Symbol("index2") * W_stride[0], sympy.Symbol("index1") * W_stride[1]]
-
-        vlane_split_axis = vlane_split_axis if nr_rdim==0 else 0
-        Y_tile_size = [TILE_M, TILE_N] if nr_rdim == 0 else [TILE_N, TILE_M]
-        Y_tile_stride=[1, TILE_M] if nr_rdim == 0 else [TILE_M, 1]
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("Y_buffer")
         Y_stride = Y.get_layout().stride if N>1 else [Y.get_layout().stride[0], 0]
-        if nr_rdim == 0:
-            Y_idx = [sympy.Symbol("index0") * Y_stride[0], sympy.Symbol("index1") * Y_stride[1]]
-        else:
-            Y_idx = [sympy.Symbol("index1") * Y_stride[1], sympy.Symbol("index0") * Y_stride[0]]
 
-        # Extract Bias info
+        X_tile_desc, X_idx = build_tile(
+            "X_buffer", kernel.vector_lane,
+            axes={"m": Axis(TILE_M, X_stride[0], loop="index0"),
+                  "k": Axis(TILE_K, X_stride[1], loop="index2")},
+            sram_order=("k", "m"), lane="k", offset=X.get_layout().offset)
+
+        W_tile_desc, W_idx = build_tile(
+            "W_buffer", kernel.vector_lane,
+            axes={"k": Axis(TILE_K, W_stride[0], loop="index2"),
+                  "n": Axis(TILE_N, W_stride[1], loop="index1")},
+            sram_order=("n", "k"), lane="n", offset=W.get_layout().offset)
+
+        # The reduction template sweeps N outside M, so its tile is declared (N, M).
+        # Only the declaration flips; the SRAM order is (n, m) either way.
+        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}
+
+        Y_tile_desc, Y_idx = build_tile(
+            "Y_buffer", kernel.vector_lane, y_axes(Y_stride), sram_order=("n", "m"), lane="n")
+
+        # Extract Bias info. It accumulates into the Y buffer, so it shares Y's axes.
         Bias = None if len(self.input_nodes) == 2 else self.input_nodes[2]
-        Bias_tile_desc = mlir_common.MLIRMultiDimTile(Y_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Bias_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Bias_tile_desc.set_name("Y_buffer")
         if Bias is not None:
-          Bias_stride = Bias.get_layout().stride
-          Bias_tile_desc.offset = Bias.get_layout().offset
-          if nr_rdim == 0:
-            Bias_idx = [sympy.Symbol("index0") * Bias_stride[0], sympy.Symbol("index1") * Bias_stride[1]]
-          else:
-            Bias_idx = [sympy.Symbol("index1") * Bias_stride[1], sympy.Symbol("index0") * Bias_stride[0]]
+            Bias_tile_desc, Bias_idx = build_tile(
+                "Y_buffer", kernel.vector_lane, y_axes(Bias.get_layout().stride),
+                sram_order=("n", "m"), lane="n", offset=Bias.get_layout().offset)
         else:
-          Bias_idx = None
+            Bias_tile_desc, _ = build_tile(
+                "Y_buffer", kernel.vector_lane, y_axes(Y_stride), sram_order=("n", "m"), lane="n")
+            Bias_idx = None
 
         data_stype = mlir_common.DTYPE_TO_MLIR[X.get_dtype()]
 

From 0a79e01096546d2a7ea09b74bfcbbdc155469589 Mon Sep 17 00:00:00 2001
From: Wonhyuk Yang <wonhyuk@postech.ac.kr>
Date: Fri, 10 Jul 2026 17:38:42 +0900
Subject: [PATCH 3/4] [Frontend] Build the sdpa, sort, cat and maxpool tile
 descriptors from axes

The last four templates. No template constructs an MLIRMultiDimTile any more.

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 name, with the reason next to them.

sort's lane axis can have extent 1, since a one-dimensional sort runs in a single
lane. cat names its tiles after the input rather than "<name>_buffer".

maxpool, cat and sort emit byte-identical MLIR. sdpa's degenerate batch axis
picks up the SRAM stride it would have if it were not degenerate, so four
tile_stride entries change, the same way BMM's do. All four match CPU in
functional mode: max_pool2d, cat, sort and topk are exact, and
tests/ops/attention/test_sdpa.py passes for every head and sequence length it
covers. GEMM, BMM and conv keep the kernels they had, and tests/ops/fusion still
fuses the same ones.
---
 PyTorchSimFrontend/mlir/mlir_cat_template.py  | 13 ++-
 .../mlir/mlir_maxpool_template.py             | 27 +++----
 PyTorchSimFrontend/mlir/mlir_sdpa_template.py | 79 ++++++++-----------
 PyTorchSimFrontend/mlir/mlir_sort_template.py | 31 ++++----
 4 files changed, 64 insertions(+), 86 deletions(-)

diff --git a/PyTorchSimFrontend/mlir/mlir_cat_template.py b/PyTorchSimFrontend/mlir/mlir_cat_template.py
index b922e51b..e0ee448a 100644
--- a/PyTorchSimFrontend/mlir/mlir_cat_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_cat_template.py
@@ -6,6 +6,7 @@
 from torch._inductor.ir import IRNode
 
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplate, MLIRTemplateKernel
 
 
@@ -262,14 +263,10 @@ def _build_tile_descriptors(
             excluded_dims = set()
 
         def make_tile_desc(tile_sz, vector_lane, name, offset):
-            desc = mlir_common.MLIRMultiDimTile(
-                tile_sz, vector_lane,
-                vlane_split_axis=len(tile_sz) - 1,
-                vlane_stride=1
-            )
-            desc.set_tile_size(tile_sz)
-            desc.set_name(name)
-            desc.offset = offset
+            # A plain row-major tile: the innermost axis is contiguous and rides the lanes.
+            axes = {f"d{i}": Axis(sz) for i, sz in enumerate(tile_sz)}
+            desc, _ = build_tile(name, vector_lane, axes, sram_order=tuple(axes),
+                                 lane=f"d{len(tile_sz) - 1}", offset=offset)
             return desc
 
         output_offset = output_node.get_layout().offset
diff --git a/PyTorchSimFrontend/mlir/mlir_maxpool_template.py b/PyTorchSimFrontend/mlir/mlir_maxpool_template.py
index 3658f992..0e2e3749 100644
--- a/PyTorchSimFrontend/mlir/mlir_maxpool_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_maxpool_template.py
@@ -5,6 +5,7 @@
 from torch._inductor.ir import Buffer
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 import sympy
 
 # This template only represents the DMA operations
@@ -55,22 +56,18 @@ def render(self,
         BCH = B * C * H
         kernel.loop_size = None
 
-        # Prepare tile descriptors
-        vlane_stride = 1 # Used dummy value
-        vlane_split_axis = 1
-        X_tile_size = [in_tile, in_tile]
-        X_tile_stride = [1, in_tile]
-        X_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        X_tile_desc.set_tile_size_stride(X_tile_size, X_tile_stride)
-        X_tile_desc.set_name("X_buffer")
-        X_idx = [sympy.Symbol("index0"), sympy.Symbol("index1")*W] # To keep index arguemnt order, we used index_list
+        # Prepare tile descriptors. Rows ride the lanes, columns are contiguous in a lane.
+        X_tile_desc, X_idx = build_tile(
+            "X_buffer", kernel.vector_lane,
+            axes={"col": Axis(in_tile, 1, loop="index0"),
+                  "row": Axis(in_tile, W, loop="index1")},
+            sram_order=("row", "col"), lane="row")
 
-        Y_tile_size = [out_tile, out_tile]
-        Y_tile_stride = [1, out_tile]
-        Y_tile_desc = mlir_common.MLIRMultiDimTile(X_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        Y_tile_desc.set_tile_size_stride(Y_tile_size, Y_tile_stride)
-        Y_tile_desc.set_name("W_buffer")
-        Y_idx = [sympy.Symbol("index0"), sympy.Symbol("index1")*W]
+        Y_tile_desc, Y_idx = build_tile(
+            "W_buffer", kernel.vector_lane,
+            axes={"col": Axis(out_tile, 1, loop="index0"),
+                  "row": Axis(out_tile, W, loop="index1")},
+            sram_order=("row", "col"), lane="row")
 
         kernel.render_options = dict(
             KERNEL_NAME=self.name,
diff --git a/PyTorchSimFrontend/mlir/mlir_sdpa_template.py b/PyTorchSimFrontend/mlir/mlir_sdpa_template.py
index a3ae6192..d4041291 100644
--- a/PyTorchSimFrontend/mlir/mlir_sdpa_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_sdpa_template.py
@@ -12,6 +12,7 @@
 
 from PyTorchSimFrontend import extension_config
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplate
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 
@@ -372,14 +373,11 @@ def render(self,
 
         # Hardware constraint: The tile split axis is restricted.
         # To accommodate this, we compute (key @ query.t) instead of (query @ key.t).
-        # SRAM settings
-        vlane_split_axis = 1
-        q_tile_size = [1, tile_l, tile_e]
-        q_tile_stride = [0, tile_e, 1]
-        q_tile_desc = mlir_common.MLIRMultiDimTile(q_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        q_tile_desc.set_tile_size_stride(q_tile_size, q_tile_stride)
-        q_tile_desc.set_name("q_buffer")
-        q_tile_desc.offset = query.get_layout().offset
+        # SRAM settings. Batch is degenerate; the sequence axis rides the lanes.
+        q_tile_desc, _ = build_tile(
+            "q_buffer", kernel.vector_lane,
+            axes={"b": Axis(1), "l": Axis(tile_l), "e": Axis(tile_e)},
+            sram_order=("b", "l", "e"), lane="l", offset=query.get_layout().offset)
         # DRAM settings
         q_stride = q_tensor.stride()
 
@@ -387,67 +385,54 @@ def render(self,
         # the split axis of the first operand differs from a standard linear algebra matmul.
         # The first operand (key) must be split along the column axis.
         # This logic aligns with the relationship between the dot product's summation direction and the hardware's accumulation direction in the SA.
-        # SRAM settings
-        vlane_split_axis = 2
-        k_tile_size = [1, tile_s, tile_e]
-        k_tile_stride = [0, 1, tile_s]
-        k_tile_desc = mlir_common.MLIRMultiDimTile(k_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        k_tile_desc.set_tile_size_stride(k_tile_size, k_tile_stride)
-        k_tile_desc.set_name("k_buffer")
-        k_tile_desc.offset = key.get_layout().offset
+        # SRAM settings. The embedding axis rides the lanes here, and the sequence axis is
+        # the contiguous one -- key is the stationary operand, so it is laid out transposed.
+        k_tile_desc, _ = build_tile(
+            "k_buffer", kernel.vector_lane,
+            axes={"b": Axis(1), "s": Axis(tile_s), "e": Axis(tile_e)},
+            sram_order=("b", "e", "s"), lane="e", offset=key.get_layout().offset)
         # DRAM settings
         k_stride = k_tensor.stride()
 
         # Since we compute mul = key @ query.t, we perform out.t = (value.t @ Softmax(mul).t).t,
         # which simplifies to (value.t @ Softmax(mul))
         # SRAM settings
-        vlane_split_axis = 1
-        v_tile_size = [1, tile_s, tile_e]
-        v_tile_stride = [0, tile_e, 1]
-        v_tile_desc = mlir_common.MLIRMultiDimTile(v_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        v_tile_desc.set_tile_size_stride(v_tile_size, v_tile_stride)
-        v_tile_desc.set_name("v_buffer")
-        v_tile_desc.offset = value.get_layout().offset
+        v_tile_desc, _ = build_tile(
+            "v_buffer", kernel.vector_lane,
+            axes={"b": Axis(1), "s": Axis(tile_s), "e": Axis(tile_e)},
+            sram_order=("b", "s", "e"), lane="s", offset=value.get_layout().offset)
         # DRAM settings
         v_stride = v_tensor.stride()
 
         # Output is also stored in transposed format to match the value.t @ Softmax(mul) operation.
         # SRAM settings
-        vlane_split_axis = 1
-        out_tile_size = [1, tile_l, tile_e]
-        out_tile_stride=[0, tile_e, 1]
-        out_tile_desc = mlir_common.MLIRMultiDimTile(out_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        out_tile_desc.set_tile_size_stride(out_tile_size, out_tile_stride)
-        out_tile_desc.set_name("out_buffer")
+        out_tile_desc, _ = build_tile(
+            "out_buffer", kernel.vector_lane,
+            axes={"b": Axis(1), "l": Axis(tile_l), "e": Axis(tile_e)},
+            sram_order=("b", "l", "e"), lane="l")
         # DRAM settings
         out_stride = out.get_layout().stride[1:]
 
         # Intermediate buffers
 
         # For mul = key @ query.t
-        vlane_split_axis = 1
-        mul_tile_size = [tile_s, tile_l]
-        mul_tile_stride = [tile_l, 1]
-        mul_tile_desc = mlir_common.MLIRMultiDimTile(mul_tile_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        mul_tile_desc.set_tile_size_stride(mul_tile_size, mul_tile_stride)
-        mul_tile_desc.set_name("mul_buffer")
         #FIXME. What is the offset? -> It doesn't matter at this time.
+        mul_tile_desc, _ = build_tile(
+            "mul_buffer", kernel.vector_lane,
+            axes={"s": Axis(tile_s), "l": Axis(tile_l)},
+            sram_order=("s", "l"), lane="l")
 
         # For storing maximum values per row
-        vlane_split_axis = 0
-        max_size = [tile_l, 2]
-        max_stride = [2, 1]
-        max_desc = mlir_common.MLIRMultiDimTile(max_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        max_desc.set_tile_size_stride(max_size, max_stride)
-        max_desc.set_name("max_buffer")
+        max_desc, _ = build_tile(
+            "max_buffer", kernel.vector_lane,
+            axes={"l": Axis(tile_l), "pair": Axis(2)},
+            sram_order=("l", "pair"), lane="l")
 
         # For storing summation per row
-        vlane_split_axis = 0
-        sum_size = [tile_l, 2]
-        sum_stride = [2, 1]
-        sum_desc = mlir_common.MLIRMultiDimTile(sum_size, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        sum_desc.set_tile_size_stride(sum_size, sum_stride)
-        sum_desc.set_name("sum_buffer")
+        sum_desc, _ = build_tile(
+            "sum_buffer", kernel.vector_lane,
+            axes={"l": Axis(tile_l), "pair": Axis(2)},
+            sram_order=("l", "pair"), lane="l")
 
         # For reduction
         chunk_size = 16
diff --git a/PyTorchSimFrontend/mlir/mlir_sort_template.py b/PyTorchSimFrontend/mlir/mlir_sort_template.py
index 338f9636..d4f65b1a 100644
--- a/PyTorchSimFrontend/mlir/mlir_sort_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_sort_template.py
@@ -7,6 +7,7 @@
 from torch._inductor.codegen import common
 
 from PyTorchSimFrontend.mlir import mlir_common
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplate, MLIRTemplateKernel
 from PyTorchSimFrontend.mlir.mlir_common import LoopLevel
 
@@ -260,22 +261,20 @@ def render(
         # indent for DMA ops = 2 (inside func) + 2 per outer loop
         indent_size = 2 + len(output_dim) * 2 + 4
 
-        vlane_stride = 1
-        vlane_split_axis = 0
-        x_tile_desc = mlir_common.MLIRMultiDimTile(tile_sizes, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        x_tile_desc.set_tile_size_stride(tile_sizes, [sort_size, 1])
-        x_tile_desc.set_name("X_buffer")
-        x_tile_desc.offset = x_layout.offset
-
-        xi_tile_desc = mlir_common.MLIRMultiDimTile(tile_sizes, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        xi_tile_desc.set_tile_size_stride(tile_sizes, [sort_size, 1])
-        xi_tile_desc.set_name("XI_buffer")
-        xi_tile_desc.offset = xi_layout.offset
-
-        yv_tile_desc = mlir_common.MLIRMultiDimTile(tile_sizes, kernel.vector_lane, vlane_split_axis, vlane_stride)
-        yv_tile_desc.set_tile_size_stride(tile_sizes, [sort_size, 1])
-        yv_tile_desc.set_name("YV_buffer")
-        yv_tile_desc.offset = yv_layout.offset
+        # One row per lane; the sorted axis is contiguous inside a lane. Every operand
+        # shares that shape and differs only in its DRAM strides.
+        def sort_axes(dram_stride):
+            return {"tile": Axis(tile_sizes[0], dram_stride[0]),
+                    "sort": Axis(tile_sizes[1], dram_stride[1])}
+
+        def sort_tile(buffer, dram_stride, offset):
+            desc, _ = build_tile(buffer, kernel.vector_lane, sort_axes(dram_stride),
+                                 sram_order=("tile", "sort"), lane="tile", offset=offset)
+            return desc
+
+        x_tile_desc = sort_tile("X_buffer", x_dram_stride, x_layout.offset)
+        xi_tile_desc = sort_tile("XI_buffer", xi_dram_stride, xi_layout.offset)
+        yv_tile_desc = sort_tile("YV_buffer", yv_dram_stride, yv_layout.offset)
 
         data_stype = mlir_common.DTYPE_TO_MLIR[x.get_dtype()]
         idx_stype = mlir_common.DTYPE_TO_MLIR[xi.get_dtype()]

From e8f5bbec299876fcdb578d236215af6299acb3cb Mon Sep 17 00:00:00 2001
From: Wonhyuk Yang <wonhyuk@postech.ac.kr>
Date: Sat, 11 Jul 2026 17:04:32 +0900
Subject: [PATCH 4/4] [Frontend] Derive the fusion aliasing from the tile axes

A fused epilogue or prologue renames its loop variables to the template's, in the
tile's declared order -- that is what set_ranges consumes, and the order the
reduction MVOUT reads the DRAM strides in. Every template spelled that list out by
hand as a `dim_aliasing` dict:

    epilogue_dim_aliasing = {"index0": "index1", "index1": "index0"}   # GEMM reduction
    prologue_dim_aliasing = {"index0": "index2", "index1": "index1"}   # GEMM weight
    dim_aliasing          = {"index0": "c0", "index1": "tile_n", ...}  # conv single-batch

Sixteen of these across gemm, bmm and the four conv templates. Each is exactly the
loop name of every axis in the corresponding tile, in declared order -- the same
tile the template already builds. So they are all `aliasing(tile.axes)`, a list.

The keys of these dicts were never read; every consumer took `.values()`. So
`dim_aliasing` becomes a plain list, and the six consumer sites drop the
`.values()`. `input_dim_aliasing` and `weight_dim_aliasing`, carried in
prologue_info, were never read at all and are deleted.

conv single-batch declared its degenerate batch axis with loop=None while its
aliasing named it "c0"; since the axis has stride 0, `Symbol("c0")*0` and
`Integer(0)` are the same DRAM index term, so naming it "c0" is byte-identical and
lets the aliasing derive.

sdpa keeps its hand-written aliasing: its epilogue frame is the kernel output rank
(4), not its 3-axis output tile, so it is not the tile's loop order. Left as is.

Verified by regenerating every kernel from scratch: byte-identical MLIR for the
gemm, bmm and conv cases, functional mode matches CPU (max abs diff 4.2e-05), and
tests/ops/fusion fuses the same kernels. The equivalence of all sixteen dicts to
the derivation was first confirmed by asserting it across the suite, then the
dicts were removed.
---
 PyTorchSimFrontend/mlir/mlir_bmm_template.py  | 48 ++++++++-----------
 .../mlir/mlir_conv_mt_template.py             |  8 ++--
 .../mlir/mlir_conv_sb_template.py             |  8 ++--
 .../mlir/mlir_conv_sbs_template.py            |  8 ++--
 PyTorchSimFrontend/mlir/mlir_conv_template.py |  8 ++--
 PyTorchSimFrontend/mlir/mlir_gemm_template.py | 44 +++++++----------
 PyTorchSimFrontend/mlir/mlir_template.py      | 14 +++---
 PyTorchSimFrontend/mlir/tile_axis.py          | 10 ++++
 8 files changed, 70 insertions(+), 78 deletions(-)

diff --git a/PyTorchSimFrontend/mlir/mlir_bmm_template.py b/PyTorchSimFrontend/mlir/mlir_bmm_template.py
index de97b1e6..660d0e07 100644
--- a/PyTorchSimFrontend/mlir/mlir_bmm_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_bmm_template.py
@@ -7,7 +7,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
-from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile, aliasing
 
 BMM_TEMPLATE = r"""
 // BMM kernel
@@ -182,15 +182,12 @@ def render(self,
         nr_reduction_nodes = [node for node in epilogue_nodes if node.is_reduction()] if epilogue_nodes is not None else []
         if nr_reduction_nodes:
             template = BMM_REDUCTION_TEMPLATE
-            epilogue_dim_aliasing = self.REDUCTION_EPILOGUE_ALIASING
             nr_rdim = 1
         elif prologue_nodes:
             template = BMM_PROLOGUE_TEMPLATE
-            epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
             nr_rdim = 0
         else:
             template = BMM_TEMPLATE
-            epilogue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2": "index2"}
             nr_rdim = 0
 
         # Prepare tile descriptors. Batch is the outermost SRAM axis and degenerate (one
@@ -198,18 +195,18 @@ def render(self,
         X_stride, W_stride = X_tensor.stride(), W_tensor.stride()
         Y_stride = Y.get_layout().stride
 
-        X_tile_desc, X_idx = build_tile(
-            "X_buffer", kernel.vector_lane,
-            axes={"b": Axis(1,      X_stride[0], loop="index0"),
+        X_axes = {"b": Axis(1,      X_stride[0], loop="index0"),
                   "m": Axis(TILE_M, X_stride[1], loop="index1"),
-                  "k": Axis(TILE_K, X_stride[2], loop="index3")},
+                  "k": Axis(TILE_K, X_stride[2], loop="index3")}
+        X_tile_desc, X_idx = build_tile(
+            "X_buffer", kernel.vector_lane, X_axes,
             sram_order=("b", "k", "m"), lane="k", offset=X.get_layout().offset)
 
-        W_tile_desc, W_idx = build_tile(
-            "W_buffer", kernel.vector_lane,
-            axes={"b": Axis(1,      W_stride[0], loop="index0"),
+        W_axes = {"b": Axis(1,      W_stride[0], loop="index0"),
                   "k": Axis(TILE_K, W_stride[1], loop="index3"),
-                  "n": Axis(TILE_N, W_stride[2], loop="index2")},
+                  "n": Axis(TILE_N, W_stride[2], loop="index2")}
+        W_tile_desc, W_idx = build_tile(
+            "W_buffer", kernel.vector_lane, W_axes,
             sram_order=("b", "n", "k"), lane="n", offset=W.get_layout().offset)
 
         # The reduction template sweeps N outside M, so its tile is declared (B, N, M).
@@ -220,8 +217,10 @@ def y_axes(stride):
             n = Axis(TILE_N, stride[2], loop="index2")
             return {"b": b, "n": n, "m": m} if nr_rdim else {"b": b, "m": m, "n": n}
 
+        Y_axes = y_axes(Y_stride)
         Y_tile_desc, Y_idx = build_tile(
-            "Y_buffer", kernel.vector_lane, y_axes(Y_stride), sram_order=("b", "n", "m"), lane="n")
+            "Y_buffer", kernel.vector_lane, Y_axes, sram_order=("b", "n", "m"), lane="n")
+        epilogue_dim_aliasing = aliasing(Y_axes)
 
         # Extract Bias info. It accumulates into the Y buffer, so it shares Y's axes.
         if Bias is not None:
@@ -255,35 +254,26 @@ def y_axes(stride):
 
         if prologue_nodes:
           prologue_output_name = list(prologue_nodes[0].read_writes.writes)[0].name
-          if prologue_output_name == X.get_name():
-            # Input fusion case
-            prologue_var = "X"
-            prologue_sram_var = "X_buffer"
-            prologue_tile_desc = X_tile_desc
-            prologue_dim_aliasing = {"index0":"index0", "index1":"index1", "index2":"index3"}
-            is_input_fused = True
+          is_input_fused = prologue_output_name == X.get_name()
+          if is_input_fused:
+              prologue_var, prologue_sram_var = "X", "X_buffer"
+              prologue_tile_desc, prologue_dim_aliasing = X_tile_desc, aliasing(X_axes)
           else:
-            # Weight fusion case
-            prologue_var = "W"
-            prologue_sram_var = "W_buffer"
-            prologue_tile_desc = W_tile_desc
-            prologue_dim_aliasing = {"index0":"index0", "index1":"index3", "index2":"index2"}
-            is_input_fused = False
- 
+              prologue_var, prologue_sram_var = "W", "W_buffer"
+              prologue_tile_desc, prologue_dim_aliasing = W_tile_desc, aliasing(W_axes)
+
           kernel.prologue_info = dict (
               input_dram_var = "X",
               input_sram_var = "X_buffer",
               input_tile_desc = X_tile_desc,
               input_idx = X_idx,
               input_subtile_size = [1, TILE_M, TILE_K], # TODO. Curently, Subtiling is not supported for prologue template
-              input_dim_aliasing = {"index0":"index0", "index1":"index1", "index2":"index3"},
 
               weight_dram_var = "W",
               weight_sram_var = "W_buffer",
               weight_tile_desc = W_tile_desc,
               weight_idx = W_idx,
               weight_subtile_size = [1, TILE_K, TILE_N], # TODO. Curently, Subtiling is not supported for prologue template
-              weight_dim_aliasing = {"index0":"index0", "index1":"index3", "index2":"index2"},
 
               # Descriptor for fusion
               dram_var = prologue_var,
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
index 9ac8a518..74d06d5b 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
@@ -5,7 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
-from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile, aliasing
 
 CONV_TEMPLATE = r"""
 // Multi Channel Tile Conv2D kernel
@@ -179,9 +179,9 @@ def y_axes(m_stride, n_stride, h_stride, w_stride, loops):
                     "o_w": Axis(TILE_O_W, w_stride, loop=loops[3])}
 
         Y_SRAM_ORDER = ("o_h", "o_w", "n", "m")
+        Y_axes = y_axes(O_C*O_H*O_W, O_H*O_W, O_W, 1, ["tile_m", "tile_n", "o_h", "o_w"])
         Y_tile_desc, Y_idx = build_tile(
-            "output_buffer", kernel.vector_lane,
-            y_axes(O_C*O_H*O_W, O_H*O_W, O_W, 1, ["tile_m", "tile_n", "o_h", "o_w"]),
+            "output_buffer", kernel.vector_lane, Y_axes,
             sram_order=Y_SRAM_ORDER, lane="n")
 
         # Extract Bias info. It accumulates into the output buffer, and only walks channels.
@@ -244,7 +244,7 @@ def y_axes(m_stride, n_stride, h_stride, w_stride, loops):
             dram_var = "Y",
             dram_idx = Y_idx,
             dram_tile_desc = Y_tile_desc,
-            dim_aliasing = {"index0":"tile_m", "index1":"tile_n", "index2":"o_h", "index3":"o_w"}
+            dim_aliasing = aliasing(Y_axes)
         )
         kernel.exception_nodes["X"] = {"numel" : (I_W+2*PADDING_W)*(I_H+2*PADDING_H)*I_C*BATCH}
         code = self._template_from_string(conv_template).render(**kernel.render_options)
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
index cf61f3fa..d4975e99 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
@@ -5,7 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
-from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile, aliasing
 
 CONV_TEMPLATE = r"""
 // Single Batch Conv2D kernel
@@ -178,9 +178,9 @@ def y_axes(b_stride, n_stride, h_stride, m_stride, loops):
                     "m":   Axis(TILE_M,    m_stride, loop=loops[3])}
 
         Y_SRAM_ORDER = ("b", "o_h", "n", "m")
+        Y_axes = y_axes(0, O_H*O_W, O_W, 1, ["c0", "tile_n", "o_h", "tile_m"])
         Y_tile_desc, Y_idx = build_tile(
-            "output_buffer", kernel.vector_lane,
-            y_axes(0, O_H*O_W, O_W, 1, [None, "tile_n", "o_h", "tile_m"]),
+            "output_buffer", kernel.vector_lane, Y_axes,
             sram_order=Y_SRAM_ORDER, lane="n")
 
         # Extract Bias info. It accumulates into the output buffer, and only walks channels.
@@ -243,7 +243,7 @@ def y_axes(b_stride, n_stride, h_stride, m_stride, loops):
             dram_var = "Y",
             dram_idx = Y_idx,
             dram_tile_desc = Y_tile_desc,
-            dim_aliasing = {"index0":"c0", "index1":"tile_n", "index2":"o_h", "index3":"tile_m"}
+            dim_aliasing = aliasing(Y_axes)
         )
         kernel.exception_nodes["X"] = {"numel" : (I_W+2*PADDING_W)*(I_H+2*PADDING_H)*I_C*BATCH}
         code = self._template_from_string(conv_template).render(**kernel.render_options)
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
index 4c471006..c1e7ffa2 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
@@ -5,7 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
-from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile, aliasing
 
 CONV_TEMPLATE = r"""
 // Single Batch Conv2D (Stride != 1) kernel
@@ -179,9 +179,9 @@ def y_axes(b_stride, n_stride, h_stride, m_stride, loops):
                     "m":   Axis(TILE_M,   m_stride, loop=loops[3])}
 
         Y_SRAM_ORDER = ("b", "o_h", "n", "m")
+        Y_axes = y_axes(0, O_H*O_W, O_W, 1, ["c0", "tile_n", "o_h", "tile_m"])
         Y_tile_desc, Y_idx = build_tile(
-            "output_buffer", kernel.vector_lane,
-            y_axes(0, O_H*O_W, O_W, 1, [None, "tile_n", "o_h", "tile_m"]),
+            "output_buffer", kernel.vector_lane, Y_axes,
             sram_order=Y_SRAM_ORDER, lane="n")
 
         # Extract Bias info. It accumulates into the output buffer, and only walks channels.
@@ -244,7 +244,7 @@ def y_axes(b_stride, n_stride, h_stride, m_stride, loops):
             dram_var = "Y",
             dram_idx = Y_idx,
             dram_tile_desc = Y_tile_desc,
-            dim_aliasing = {"index0":"c0", "index1":"tile_n", "index2":"o_h", "index3":"tile_m"}
+            dim_aliasing = aliasing(Y_axes)
         )
         kernel.exception_nodes["X"] = {"numel" : (I_W+2*PADDING_W)*(I_H+2*PADDING_H)*I_C*BATCH}
         code = self._template_from_string(conv_template).render(**kernel.render_options)
diff --git a/PyTorchSimFrontend/mlir/mlir_conv_template.py b/PyTorchSimFrontend/mlir/mlir_conv_template.py
index aa564b43..56d159bc 100644
--- a/PyTorchSimFrontend/mlir/mlir_conv_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_conv_template.py
@@ -5,7 +5,7 @@
 from PyTorchSimFrontend.mlir.mlir_template import MLIRTemplateKernel
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend.mlir import mlir_common
-from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile, aliasing
 
 CONV_TEMPLATE = r"""
 // Conv2D kernel
@@ -181,9 +181,9 @@ def y_axes(m_stride, n_stride, h_stride, w_stride, loops):
                     "o_w": Axis(TILE_O_W, w_stride, loop=loops[3])}
 
         Y_SRAM_ORDER = ("o_h", "o_w", "n", "m")
+        Y_axes = y_axes(O_C*O_H*O_W, O_H*O_W, O_W, 1, ["tile_m", "tile_n", "o_h", "o_w"])
         Y_tile_desc, Y_idx = build_tile(
-            "output_buffer", kernel.vector_lane,
-            y_axes(O_C*O_H*O_W, O_H*O_W, O_W, 1, ["tile_m", "tile_n", "o_h", "o_w"]),
+            "output_buffer", kernel.vector_lane, Y_axes,
             sram_order=Y_SRAM_ORDER, lane="n")
 
         # Extract Bias info. It accumulates into the output buffer, and only walks channels.
@@ -246,7 +246,7 @@ def y_axes(m_stride, n_stride, h_stride, w_stride, loops):
             dram_var = "Y",
             dram_idx = Y_idx,
             dram_tile_desc = Y_tile_desc,
-            dim_aliasing = {"index0":"tile_m", "index1":"tile_n", "index2":"o_h", "index3":"o_w"}
+            dim_aliasing = aliasing(Y_axes)
         )
         kernel.exception_nodes["X"] = {"numel" : (I_W+2*PADDING_W)*(I_H+2*PADDING_H)*I_C*BATCH}
         code = self._template_from_string(conv_template).render(**kernel.render_options)
diff --git a/PyTorchSimFrontend/mlir/mlir_gemm_template.py b/PyTorchSimFrontend/mlir/mlir_gemm_template.py
index 94a7654e..b4bef044 100644
--- a/PyTorchSimFrontend/mlir/mlir_gemm_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_gemm_template.py
@@ -9,7 +9,7 @@
 from torch._inductor.ir import IRNode
 from PyTorchSimFrontend import extension_config
 from PyTorchSimFrontend.mlir import mlir_common
-from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile
+from PyTorchSimFrontend.mlir.tile_axis import Axis, build_tile, aliasing
 
 GEMM_TEMPLATE = r"""
 // GEMM {% if prologue_nodes -%}prologue fused{%- endif %} {% if epilogue_nodes -%}eilogue fused{%- endif %} kernel
@@ -130,14 +130,11 @@ def render(self,
         if (M == 0) or (N == 0) or (K == 0): # exception for MoE
             template = EMPTY_TEMPLATE
             nr_rdim = 0
-            epilogue_dim_aliasing = {}
         elif n_epilogue_node>=1 and epilogue_nodes[0].is_reduction():
             template = GEMM_REDUCTION_TEMPLATE
-            epilogue_dim_aliasing = self.REDUCTION_EPILOGUE_ALIASING
             nr_rdim = 1
         else:
             template = GEMM_TEMPLATE
-            epilogue_dim_aliasing = {"index0":"index0", "index1":"index1"}
             nr_rdim = 0
 
         TOG_latency = M if SUB_TILE_M > M else SUB_TILE_M
@@ -148,16 +145,16 @@ def render(self,
         W_stride = W.get_layout().stride if N>1 else [Y.get_layout().stride[0], 0]
         Y_stride = Y.get_layout().stride if N>1 else [Y.get_layout().stride[0], 0]
 
+        X_axes = {"m": Axis(TILE_M, X_stride[0], loop="index0"),
+                  "k": Axis(TILE_K, X_stride[1], loop="index2")}
         X_tile_desc, X_idx = build_tile(
-            "X_buffer", kernel.vector_lane,
-            axes={"m": Axis(TILE_M, X_stride[0], loop="index0"),
-                  "k": Axis(TILE_K, X_stride[1], loop="index2")},
+            "X_buffer", kernel.vector_lane, X_axes,
             sram_order=("k", "m"), lane="k", offset=X.get_layout().offset)
 
+        W_axes = {"k": Axis(TILE_K, W_stride[0], loop="index2"),
+                  "n": Axis(TILE_N, W_stride[1], loop="index1")}
         W_tile_desc, W_idx = build_tile(
-            "W_buffer", kernel.vector_lane,
-            axes={"k": Axis(TILE_K, W_stride[0], loop="index2"),
-                  "n": Axis(TILE_N, W_stride[1], loop="index1")},
+            "W_buffer", kernel.vector_lane, W_axes,
             sram_order=("n", "k"), lane="n", offset=W.get_layout().offset)
 
         # The reduction template sweeps N outside M, so its tile is declared (N, M).
@@ -167,8 +164,12 @@ def y_axes(stride):
             n = Axis(TILE_N, stride[1], loop="index1")
             return {"n": n, "m": m} if nr_rdim else {"m": m, "n": n}
 
+        Y_axes = y_axes(Y_stride)
         Y_tile_desc, Y_idx = build_tile(
-            "Y_buffer", kernel.vector_lane, y_axes(Y_stride), sram_order=("n", "m"), lane="n")
+            "Y_buffer", kernel.vector_lane, Y_axes, sram_order=("n", "m"), lane="n")
+        # The epilogue renames its loop vars to the Y tile's, in declared order; the empty
+        # (MoE) kernel has no epilogue and no store, so it carries no aliasing.
+        epilogue_dim_aliasing = [] if template is EMPTY_TEMPLATE else aliasing(Y_axes)
 
         # Extract Bias info. It accumulates into the Y buffer, so it shares Y's axes.
         Bias = None if len(self.input_nodes) == 2 else self.input_nodes[2]
@@ -209,20 +210,13 @@ def y_axes(stride):
         )
         if prologue_nodes:
             prologue_output_name = list(prologue_nodes[0].read_writes.writes)[0].name
-            if prologue_output_name == X.get_name():
-                # Input fusion case
-                prologue_var = "X"
-                prologue_sram_var = "X_buffer"
-                prologue_tile_desc = X_tile_desc
-                prologue_dim_aliasing = {"index0":"index0", "index1":"index2"}
-                is_input_fused = True
+            is_input_fused = prologue_output_name == X.get_name()
+            if is_input_fused:
+                prologue_var, prologue_sram_var = "X", "X_buffer"
+                prologue_tile_desc, prologue_dim_aliasing = X_tile_desc, aliasing(X_axes)
             else:
-                # Weight fusion case
-                prologue_var = "W"
-                prologue_sram_var = "W_buffer"
-                prologue_tile_desc = W_tile_desc
-                prologue_dim_aliasing = {"index0":"index2", "index1":"index1"}
-                is_input_fused = False
+                prologue_var, prologue_sram_var = "W", "W_buffer"
+                prologue_tile_desc, prologue_dim_aliasing = W_tile_desc, aliasing(W_axes)
 
             kernel.prologue_info = dict (
                 input_dram_var = "X",
@@ -230,14 +224,12 @@ def y_axes(stride):
                 input_tile_desc = X_tile_desc,
                 input_idx = X_idx,
                 input_subtile_size = [TILE_M, TILE_K],
-                input_dim_aliasing = {"index0":"index0", "index1":"index2"},
 
                 weight_dram_var = "W",
                 weight_sram_var = "W_buffer",
                 weight_tile_desc = W_tile_desc,
                 weight_idx = W_idx,
                 weight_subtile_size = [TILE_K, TILE_N],
-                weight_dim_aliasing = {"index0":"index2", "index1":"index1"},
 
                 # Descriptor for fusion
                 dram_var = prologue_var,
diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py
index aec36e85..4acd18cf 100644
--- a/PyTorchSimFrontend/mlir/mlir_template.py
+++ b/PyTorchSimFrontend/mlir/mlir_template.py
@@ -195,7 +195,7 @@ def __init__(self,
         self.reduction_epilogue_result = {}
         self.reduction_mean = []
         # Dim info
-        self.dim_aliasing = {}
+        self.dim_aliasing = []   # loop names in tile-declared order (see tile_axis.aliasing)
         self.reason = reason
 
     def reset(self, reason):
@@ -588,7 +588,7 @@ def codegen_template_code(self, render, template_node, prologue_nodes, epilogue_
                     ).group
                     prologue_tile_desc = kernel.set_tile_size(kernel.prologue_info, prologue=True)
                     kernel.kernel_group.set_tile_info(prologue_tile_desc)
-                    vars, reduction_vars = kernel.set_ranges(group, reduction_group, list(self.dim_aliasing.values()))
+                    vars, reduction_vars = kernel.set_ranges(group, reduction_group, self.dim_aliasing)
                     for node in prologue_nodes:
                         # Reuse created spad
                         read_list = sorted([i.name for i in node.read_writes.reads])
@@ -628,7 +628,7 @@ def codegen_template_code(self, render, template_node, prologue_nodes, epilogue_
                     _, (group, reduction_group) = max(
                         epilogue_nodes, key=lambda x: int(x.is_reduction())
                     ).group
-                    vars, reduction_vars = kernel.set_ranges(group, reduction_group, list(self.dim_aliasing.values()))
+                    vars, reduction_vars = kernel.set_ranges(group, reduction_group, self.dim_aliasing)
                     for node in epilogue_nodes:
                         node.codegen((vars, reduction_vars))
 
@@ -1105,7 +1105,7 @@ def load_epilogue(self, name: str, index: sympy.Expr):
         # Want to use tile_desc from epilogue_info
         with self.override_buffer_cse(buffer=self.applys, cse=self.apply_cse):
             index_var = self.parse_indices(index)
-        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.values()]
+        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing]
         vlane_split_axis = self.kernel_group.tile_desc.vmap.vlane_split_axis
         vlane_stride = self.kernel_group.tile_desc.vmap.vlane_stride
         tile_shape = self.kernel_group.tile_desc.get_mlir_shape(mlir_dtype)
@@ -1158,7 +1158,7 @@ def store_epilogue(self, name: str, index: sympy.Expr, value, *args, **kwargs):
 
         with self.override_buffer_cse(buffer=self.applys, cse=self.apply_cse):
             index_var = self.parse_indices(index)
-        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.values()]
+        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing]
         vlane_split_axis = self.kernel_group.tile_desc.vmap.vlane_split_axis
         vlane_stride = self.kernel_group.tile_desc.vmap.vlane_stride
         tile_shape = self.kernel_group.tile_desc.get_mlir_shape(mlir_dtype)
@@ -1198,7 +1198,7 @@ def store_epilogue(self, name: str, index: sympy.Expr, value, *args, **kwargs):
                 pass
         tile_sizes = self.kernel_group.tile_desc.get_tile_size()
         clamp_axes = [(d, iv, 0, iv_extent[iv], int(tile_sizes[d]))
-                      for d, iv in enumerate(self.dim_aliasing.values())
+                      for d, iv in enumerate(self.dim_aliasing)
                       if d < len(tile_sizes) and iv in iv_extent]
         masked_bounds = self._emit_clamp(clamp_axes, self.dma_stores)
         code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var,
@@ -1276,7 +1276,7 @@ def store_reduction_epilogue(self, name, index, value):
 
         with self.override_buffer_cse(buffer=self.reductions_suffix, cse=self.apply_cse):
             index_var = self.parse_indices(index, comments="// Store reduction")
-        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing.values()][:-1] # Assume that there is only one reduction axis
+        dram_stride = [index.coeff(sympy.Symbol(val)) for val in self.dim_aliasing][:-1] # Assume that there is only one reduction axis
         vlane_split_axis = self.kernel_group.tile_desc.vmap.vlane_split_axis
         vlane_stride = self.kernel_group.tile_desc.vmap.vlane_stride
 
diff --git a/PyTorchSimFrontend/mlir/tile_axis.py b/PyTorchSimFrontend/mlir/tile_axis.py
index 78b1e2a0..ad047363 100644
--- a/PyTorchSimFrontend/mlir/tile_axis.py
+++ b/PyTorchSimFrontend/mlir/tile_axis.py
@@ -58,6 +58,16 @@ def dram_index(axes):
             for a in axes.values()]
 
 
+def aliasing(axes):
+    """The loop name of each axis, in the tile's declared order.
+
+    This is exactly what a fused epilogue or prologue needs to rename its loop
+    variables to (set_ranges), and the order the reduction MVOUT reads the DRAM
+    strides in -- so the hand-written `dim_aliasing` dicts are just this list.
+    """
+    return [a.loop for a in axes.values()]
+
+
 def build_tile(buffer, vector_lane, axes, sram_order, lane, lane_chunk=1, offset=0):
     """Build the tile descriptor and the DRAM index expression for one operand.