Skip to content

[Frontend] Derive the fusion aliasing from the tile axes #299

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 1 commit into
base: refactor/tile-axis
Choose a base branch
from
Open

[Frontend] Derive the fusion aliasing from the tile axes #299

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
48 changes: 19 additions & 29 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,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
Expand Down Expand Up @@ -180,34 +180,31 @@ 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
# 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

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).
Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 4 additions & 4 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,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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 4 additions & 4 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,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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 4 additions & 4 deletions PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 4 additions & 4 deletions PyTorchSimFrontend/mlir/mlir_conv_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
44 changes: 18 additions & 26 deletions PyTorchSimFrontend/mlir/mlir_gemm_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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).
Expand All @@ -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]
Expand Down Expand Up @@ -207,35 +208,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":"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",
input_sram_var = "X_buffer",
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,
Expand Down
Loading