From 4f5969066e05eeb0d23919f9c3c26b6d5210b622 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 10 Jul 2026 17:38:42 +0900 Subject: [PATCH 1/3] [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 952b7a03e7312e22978cf50f6955ba8a471ef6d3 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 10 Jul 2026 17:38:42 +0900 Subject: [PATCH 2/3] [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 | 72 +++++++++-------- .../mlir/mlir_conv_sbs_template.py | 72 +++++++++-------- PyTorchSimFrontend/mlir/mlir_conv_template.py | 68 ++++++++-------- PyTorchSimFrontend/mlir/mlir_gemm_template.py | 71 ++++++++--------- 6 files changed, 217 insertions(+), 212 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_bmm_template.py b/PyTorchSimFrontend/mlir/mlir_bmm_template.py index 5323fd7c..c90ce2c7 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 @@ -190,53 +191,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 7964da0f..ff576635 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 @@ -148,40 +149,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 dfca23ec..29533f5b 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 @@ -147,39 +148,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]] - - 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]] - - 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")] - - # 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 + # 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_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") + + # 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])} + + 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 f1a42964..3f47e39c 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 @@ -148,39 +149,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]] - - 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]] - - 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")] - - # 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 + # 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_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") + + # 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])} + + 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 8bb64d48..ac4840d5 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 @@ -151,40 +152,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 871c244e..3a079d7b 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 @@ -140,53 +141,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 90d8b226cc471cc5f25840b0ad20bd32500c8e2c Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Fri, 10 Jul 2026 17:38:42 +0900 Subject: [PATCH 3/3] [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 "_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()]