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_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_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()] 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()] 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)