From 4fa02b1adeaad660801fe7a11c80db3403596334 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Sat, 11 Jul 2026 17:04:32 +0900 Subject: [PATCH] [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 c90ce2c7..d5e94941 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 @@ -180,15 +180,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 = {"index0":"index0", "index1":"index2", "index2": "index1"} 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 @@ -196,18 +193,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). @@ -218,8 +215,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: @@ -253,35 +252,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 ff576635..4109b386 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 @@ -176,9 +176,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. @@ -241,7 +241,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 29533f5b..8e8bca35 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 @@ -175,9 +175,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. @@ -240,7 +240,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 3f47e39c..ea01263e 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 @@ -176,9 +176,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. @@ -241,7 +241,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 ac4840d5..a3a27028 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 @@ -178,9 +178,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. @@ -243,7 +243,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 3a079d7b..48b594c2 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 @@ -128,14 +128,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 = {"index0":"index1", "index1":"index0"} 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 @@ -146,16 +143,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). @@ -165,8 +162,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] @@ -207,20 +208,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", @@ -228,14 +222,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 12bf7dc5..9237b3e4 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -134,7 +134,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): @@ -527,7 +527,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]) @@ -567,7 +567,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)) @@ -1044,7 +1044,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) @@ -1097,7 +1097,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) @@ -1137,7 +1137,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, @@ -1215,7 +1215,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.