Skip to content

[Frontend] Build every tile descriptor from axes #298

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YWHyuk wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

[Frontend] Build every tile descriptor from axes #298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4f59690
[Frontend] Introduce Axis: one tile dimension, and what it means in e…
YWHyuk Jul 10, 2026
952b7a0
[Frontend] Build the GEMM, BMM and conv tile descriptors from axes
YWHyuk Jul 10, 2026
90d8b22
[Frontend] Build the sdpa, sort, cat and maxpool tile descriptors fro…
YWHyuk Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 36 additions & 43 deletions PyTorchSimFrontend/mlir/mlir_bmm_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 5 additions & 8 deletions PyTorchSimFrontend/mlir/mlir_cat_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down
67 changes: 36 additions & 31 deletions PyTorchSimFrontend/mlir/mlir_conv_mt_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()]

Expand Down
72 changes: 39 additions & 33 deletions PyTorchSimFrontend/mlir/mlir_conv_sb_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()]

Expand Down
Loading