diff --git a/.github/workflows/pytorchsim_test.yml b/.github/workflows/pytorchsim_test.yml index 7cc76a5b..0c6b8f93 100644 --- a/.github/workflows/pytorchsim_test.yml +++ b/.github/workflows/pytorchsim_test.yml @@ -18,8 +18,10 @@ on: # Runner policy: the CPU-only CI image is small enough to pull on GitHub-hosted # runners, so op and model tests run on ubuntu-latest. The memory/time-intensive -# jobs stay on self-hosted: test_deepseek (largest model), test_diffusion (UNet2D -# simulation OOMs the hosted runner), and test_accuracy (accuracy + speedup). +# jobs stay on self-hosted: test_deepseek (largest model), test_swinv2 (SwinV2 +# shifted-window backbone), test_clip (CLIP vision backbone), test_convnextv2 +# (ConvNeXt V2 backbone), test_diffusion (UNet2D simulation OOMs the hosted +# runner), and test_accuracy (accuracy + speedup). jobs: test_add: name: Run test_add.py @@ -57,6 +59,24 @@ jobs: -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/elementwise/test_transcendental.py + test_pointwise: + name: Run test_pointwise.py + runs-on: ubuntu-latest + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_pointwise.py + run: | + echo "Running test_pointwise.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/elementwise/test_pointwise.py + test_activation: name: Run test_activation.py runs-on: ubuntu-latest @@ -725,6 +745,66 @@ jobs: -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ ${{ inputs.image_name }} python3 PyTorchSim/tests/models/DeepSeek/test_deepseek_v3_base.py + test_swinv2: + name: Run test_swinv2.py + # SwinV2 backbone (shifted-window attention) is heavy; keep it on a + # self-hosted runner like the other model tests. + runs-on: self-hosted + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_swinv2.py + run: | + echo "Running test_swinv2.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/models/test_swinv2.py + + test_clip: + name: Run test_clip.py + # CLIP vision backbone; keep it on a self-hosted runner like the other + # model tests (the 32x32 patch conv expands to many tiles at small arrays). + runs-on: self-hosted + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_clip.py + run: | + echo "Running test_clip.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/models/test_clip.py + + test_convnextv2: + name: Run test_convnextv2.py + # ConvNeXt V2 backbone; keep it on a self-hosted runner like the other model + # tests (the depthwise 7x7 conv lowers to one gather kernel per tap). + runs-on: self-hosted + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Run test_convnextv2.py + run: | + echo "Running test_convnextv2.py" + docker run --rm \ + -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \ + ${{ inputs.image_name }} python3 PyTorchSim/tests/models/test_convnextv2.py + test_eager: name: Run test_eager.py runs-on: ubuntu-latest diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 83503c31..a9cd8ff3 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -213,11 +213,11 @@ def load(cls, source_code, # the shared spad. Matches the GEMM tiling gate (max_spad_size = spad/2). spad_budget = extension_config.CONFIG_SPAD_INFO["spad_size"] // 2 if spad_budget < spad_usage: - logger.debug( - f"Scratchpad size exceeded: required {spad_usage} bytes, but only " - f"{spad_budget} bytes (spad/2, double-buffer budget) available." + raise SpadOverflowError( + f"Scratchpad size exceeded: kernel needs {spad_usage} bytes/lane, but only " + f"{spad_budget} bytes/lane (spad/2, double-buffer budget) are available. " + f"tile_size={kwargs.get('tile_size')}, origins={origins}, path={write_path}" ) - raise SpadOverflowError() # Launch tile graph generator gem5_pad_cmd = shlex.split(gem5_cmds[0]) diff --git a/PyTorchSimFrontend/mlir/axis_split.py b/PyTorchSimFrontend/mlir/axis_split.py index 71ec4809..2cf58c8a 100644 --- a/PyTorchSimFrontend/mlir/axis_split.py +++ b/PyTorchSimFrontend/mlir/axis_split.py @@ -29,6 +29,77 @@ def _as_int(x): return None +def _as_digit(expr): + """If ``expr`` is a single-variable *digit extractor* -- any nesting of FloorDiv / + ModularIndexing whose innermost argument is one symbol ``v`` -- return ``(v, div, mod)`` + meaning ``(v // div) % mod`` (``mod is None`` -> a pure ``v // div``). Otherwise None. + + Every such nesting collapses to a single (div, mod) by composing divisors, from four + algebraic identities (v >= 0, all constants positive integers): + FloorDiv(v//a, b) = v // (a*b) + FloorDiv((v//a)%m, b) = (v // (a*b)) % (m//b) if b | m + ModularIndexing(v//a, b,m2)= (v // (a*b)) % m2 + ModularIndexing((v//a)%m, b,m2)= (v // (a*b)) % m2 if b*m2| m + The divisibility guards make every rewrite provably equality-preserving. A multi- + variable inner argument (e.g. torch.roll's v+shift) is not a digit extractor -> None. + """ + if isinstance(expr, sympy.Symbol): + return (expr, 1, None) + if isinstance(expr, FloorDiv): + inner, b = expr.args + b, e = _as_int(b), _as_digit(expr.args[0]) + if e is None or b is None: + return None + v, a, m = e + if m is None: + return (v, a * b, None) + if m % b == 0: + return (v, a * b, m // b) + return None + if isinstance(expr, ModularIndexing): + inner, b, m2 = expr.args + b, m2, e = _as_int(b), _as_int(m2), _as_digit(expr.args[0]) + if e is None or b is None or m2 is None: + return None + v, a, m = e + if m is None: + return (v, a * b, m2) + if m % (b * m2) == 0: + return (v, a * b, m2) + return None + return None + + +def _rebuild_digit(v, a, m): + """Canonical single-level form of ``(v // a) % m``.""" + x = v if a == 1 else FloorDiv(v, a) + return x if m is None else ModularIndexing(v, a, m) + + +def flatten_nested_floormod(expr): + """Collapse nested single-variable FloorDiv/ModularIndexing to one level. + + A composition of aligned reshapes on one iteration variable leaves a nested index like + ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor + simplify_with_ranges reduces, so collect_boundaries skips its cut points (the inner base + is not a bare var) and the affine-only DMA check later rejects it. Rewriting each nested + digit extractor to its single-level (v // A) % M form (via _as_digit) exposes those cut + points to axis-split. General and pattern-free -- no per-shape special cases. + """ + try: + atoms = expr.atoms(FloorDiv, ModularIndexing) + except AttributeError: + return expr + replace = {} + for atom in atoms: + e = _as_digit(atom) + if e is not None: + canon = _rebuild_digit(*e) + if canon != atom: + replace[atom] = canon + return expr.xreplace(replace) if replace else expr + + def collect_boundaries(exprs, var_to_axis, var_ranges): """{axis_index: set(boundary cut points)} for the given index expressions. @@ -39,6 +110,7 @@ def collect_boundaries(exprs, var_to_axis, var_ranges): import collections bset = collections.defaultdict(set) for expr in exprs: + expr = flatten_nested_floormod(expr) # nested digit extractors -> single level for fd in expr.atoms(FloorDiv): base, div = fd.args k = _as_int(div) @@ -196,6 +268,7 @@ def _fold_with_ranges(expr, var_ranges): Iterated to a fixpoint (folding a mod can expose a foldable floor). """ from torch.utils._sympy.value_ranges import bound_sympy, ValueRanges + expr = flatten_nested_floormod(expr) # collapse any residual single-var nested digit ranges = {} for v, sz in var_ranges.items(): e = _as_int(sz) diff --git a/PyTorchSimFrontend/mlir/mlir_bmm_template.py b/PyTorchSimFrontend/mlir/mlir_bmm_template.py index 5323fd7c..2bc8cff8 100644 --- a/PyTorchSimFrontend/mlir/mlir_bmm_template.py +++ b/PyTorchSimFrontend/mlir/mlir_bmm_template.py @@ -152,11 +152,13 @@ """ class MLIRBMMTemplate(MLIRTemplate): + # 3-D frame (batch x M x N): batch + M row indices + the reduced N index. + REDUCTION_EPILOGUE_ALIASING = {"index0": "index0", "index1": "index2", "index2": "index1"} + def __init__(self, input_nodes, layout, input_reorder=None): super().__init__("kernel", input_nodes, layout, input_reorder) self.support_epilogue_fusion = True self.support_prologue_fusion = True - self.support_reduction_fusion = True def render(self, kernel: MLIRTemplateKernel, @@ -179,7 +181,7 @@ 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"} + epilogue_dim_aliasing = self.REDUCTION_EPILOGUE_ALIASING nr_rdim = 1 elif prologue_nodes: template = BMM_PROLOGUE_TEMPLATE diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index bd04a72d..18ad4e4f 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -44,10 +44,16 @@ def reduction_init(reduction_type, dtype): return float(0) if dtype.is_floating_point else int(0) if reduction_type == "prod": return float(1) if dtype.is_floating_point else int(1) + # Integer reductions cannot use a +/-inf identity (invalid as an int constant and + # overflows torch.tensor(inf, dtype=int)); use the dtype's representable extreme. if reduction_type in {"max", "argmax"}: - return "-inf" + if dtype.is_floating_point: + return "-inf" + return 0 if dtype is torch.bool else torch.iinfo(dtype).min if reduction_type in {"min", "argmin"}: - return "inf" + if dtype.is_floating_point: + return "inf" + return 1 if dtype is torch.bool else torch.iinfo(dtype).max if reduction_type in {"welford_reduce"}: return f"0.0" raise AssertionError(reduction_type) @@ -114,6 +120,7 @@ def write_header(self): inductor_ops = torch.ops.inductor assert_size_stride = torch._C._dynamo.guards.assert_size_stride assert_alignment = torch._C._dynamo.guards.assert_alignment + empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu alloc_from_pool = torch.ops.inductor._alloc_from_pool reinterpret_tensor = torch.ops.inductor._reinterpret_tensor custom_async_compile = CustomAsyncCompile() @@ -926,8 +933,11 @@ def index_expr(self, index, dtype): c_type = "uint64_t" new_name = f"index_expr_{compute_vec_size}" if new_name not in self.global_vars_dict: - self.header.writeline(f"{c_type} {new_name}_spad[{compute_vec_size*self.vector_lane}] __attribute__ ((section(\".spad\")));") - self.gem5_header.writeline(f"{c_type} {new_name}_spad[{compute_vec_size}] __attribute__((aligned(64)));") + # The initializer below stores two elements per iteration, so its last + # iteration writes one slot past compute_vec_size. + numel_per_lane = compute_vec_size + 1 + self.header.writeline(f"{c_type} {new_name}_spad[{numel_per_lane}] __attribute__ ((section(\".spad\")));") + self.gem5_header.writeline(f"{c_type} {new_name}_spad[{numel_per_lane*self.vector_lane}] __attribute__((aligned(64)));") self.global_vars.writeline(f"memref.global @{new_name}_spad : {tile_shape}") self.global_vars_dict[new_name] = dict() sram_var = self.spad_cse.generate(self.spad_buffer, f"memref.get_global @{new_name}_spad : {tile_shape}") @@ -1323,16 +1333,23 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride - # Case 4. Tile is 4-D tile (e.g., Convolution epilogue) - elif len(local_dims) == 4: - is_reduction = self.reduction_depth < 3 and not store_reduction - if is_reduction: - raise NotImplementedError("Currently not implemented... ;)") - local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) - local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis - local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride + # Case 4+. Tile is 4-D or higher (Convolution epilogue, gathered attention bias, + # var_mean over an axis whose batch dims got split into many loop vars). else: - local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) + # A reduction tile carries the reduction axis (loop dims >= reduction_depth). + # It must place that axis-group OUTERMOST in the per-lane layout so the 2-D + # [reduction | batch] multi_reduction reduces the reduction axis (not a batch + # axis). Same reorder the 3-D case does, generalized to any rank: an indirect + # attention-bias gather blocks a head*query dim-merge and leaves head in-lane, + # and a var_mean's token axis can split into several batch loop vars -- both + # leave a batch axis inner-to-reduction under the default row-major order. + is_reduction = any(d >= self.reduction_depth for d in local_dims) and not store_reduction + if is_reduction: + r = self.get_nr_rdim() + axis_order = list(range(r, len(local_dims))) + list(range(r - 1, -1, -1)) + local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims], axis_order) + else: + local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride diff --git a/PyTorchSimFrontend/mlir/mlir_conv_common.py b/PyTorchSimFrontend/mlir/mlir_conv_common.py index d577dbd8..b1c93f3b 100644 --- a/PyTorchSimFrontend/mlir/mlir_conv_common.py +++ b/PyTorchSimFrontend/mlir/mlir_conv_common.py @@ -14,7 +14,6 @@ def __init__(self, input_nodes, layout, input_reorder=None, **kwargs): super().__init__("kernel", input_nodes, layout, input_reorder) self.support_epilogue_fusion = True self.support_prologue_fusion = False - self.support_reduction_fusion = False self.stride = kwargs["stride"] self.padding = kwargs["padding"] self.dilation = kwargs["dilation"] diff --git a/PyTorchSimFrontend/mlir/mlir_decomposition.py b/PyTorchSimFrontend/mlir/mlir_decomposition.py index a0030e3b..ea7a7ebb 100644 --- a/PyTorchSimFrontend/mlir/mlir_decomposition.py +++ b/PyTorchSimFrontend/mlir/mlir_decomposition.py @@ -371,4 +371,56 @@ def decompose_native_multi_head_attention( attn_weights_mean = attn_weights.mean(dim=1) # Average over heads return output, attn_weights_mean else: - return (output, None) \ No newline at end of file + return (output, None) + + +# Lower roll ourselves as narrow + cat, then REALIZE the result. roll has no Inductor lowering; it +# is decomposed (torch's default is index_select(fmod(...)), i.e. a modular gather the affine-only +# DMA cannot express). Even our narrow+cat rewrite lets the slice offset fuse into a downstream +# modular reshape, re-creating a shifted index like (p+60)%8. copy_input forces a hard buffer +# boundary so the consumer reads a plain affine index -- a plain .contiguous()/clone is inlined by +# Inductor and does NOT create the boundary. Remove roll from the decomp table so it reaches this +# lowering instead of being decomposed first. +from torch._inductor.decomposition import decompositions as _inductor_decompositions +from torch._inductor import lowering as _ind_lowering +from torch._inductor import ir as _ind_ir + +_inductor_decompositions.pop(aten.roll.default, None) + + +def _roll_lowering(x, shifts, dims=()): + if not isinstance(shifts, (list, tuple)): + shifts = (shifts,) + if not isinstance(dims, (list, tuple)): + dims = (dims,) + slice_l = _ind_lowering.lowerings[aten.slice.Tensor] + cat_l = _ind_lowering.lowerings[aten.cat.default] + view_l = _ind_lowering.lowerings[aten.view.default] + + # Realize the INPUT: roll's producer is often a reshaped view (e.g. swinv2's window_reverse), + # and our slice's offset would otherwise fuse into that reshape's modular index. Reading a + # materialized buffer makes each slice a plain contiguous read. + x = _ind_ir.ExternKernel.copy_input(x) + + def roll_dim(t, shift, dim): + n = int(t.get_size()[dim]) + s = int(shift) % n + if s == 0: + return t + front = slice_l(t, dim, n - s, n) # narrow(t, dim, n-s, s) + back = slice_l(t, dim, 0, n - s) + return cat_l([front, back], dim) + + if len(dims) == 0: + numel = 1 + for d in x.get_size(): + numel *= int(d) + result = view_l(roll_dim(view_l(x, [numel]), shifts[0], 0), list(x.get_size())) + else: + result = x + for shift, dim in zip(shifts, dims): + result = roll_dim(result, shift, dim) + return _ind_ir.ExternKernel.copy_input(result) + + +_ind_lowering.lowerings[aten.roll.default] = _roll_lowering \ No newline at end of file diff --git a/PyTorchSimFrontend/mlir/mlir_gemm_template.py b/PyTorchSimFrontend/mlir/mlir_gemm_template.py index 871c244e..6fb2b9bc 100644 --- a/PyTorchSimFrontend/mlir/mlir_gemm_template.py +++ b/PyTorchSimFrontend/mlir/mlir_gemm_template.py @@ -103,11 +103,13 @@ """ class MLIRGemmTemplate(MLIRTemplate): + # 2-D frame (M rows x N cols): one row index + the reduced N index. + REDUCTION_EPILOGUE_ALIASING = {"index0": "index1", "index1": "index0"} + def __init__(self, input_nodes, layout, input_reorder=None): super().__init__("kernel", input_nodes, layout, input_reorder) self.support_epilogue_fusion = True self.support_prologue_fusion = True - self.support_reduction_fusion = True def render(self, kernel: MLIRTemplateKernel, @@ -130,7 +132,7 @@ def render(self, 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"} + epilogue_dim_aliasing = self.REDUCTION_EPILOGUE_ALIASING nr_rdim = 1 else: template = GEMM_TEMPLATE diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index d21cc1d0..98e0272b 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -13,10 +13,15 @@ def reduction_combine_vec(reduction_type, vector_value, init_value, axis, shape, return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" if reduction_type == "prod": return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" + # element type of the reduced vector (e.g. "vector<8x2xi64>" -> "i64"); int max/min use + # the signed-integer reduction kinds, not the float ones (maximumf/minimumf reject ints). + _is_int = shape.rsplit("x", 1)[-1].rstrip(">").strip().startswith("i") if reduction_type == "max": - return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" + kind = "maxsi" if _is_int else "maximumf" + return f"vector.multi_reduction <{kind}>, %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" if reduction_type == "min": - return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" + kind = "minsi" if _is_int else "minimumf" + return f"vector.multi_reduction <{kind}>, %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" if reduction_type == "any": return f"vector.multi_reduction , %{vector_value}, %{init_value} [{axis}] : {shape} to {reduced_shape}" raise AssertionError(reduction_type) diff --git a/PyTorchSimFrontend/mlir/mlir_scheduling.py b/PyTorchSimFrontend/mlir/mlir_scheduling.py index 8520596c..3cdac26f 100644 --- a/PyTorchSimFrontend/mlir/mlir_scheduling.py +++ b/PyTorchSimFrontend/mlir/mlir_scheduling.py @@ -1,9 +1,5 @@ import os -import math import sympy -from functools import reduce -import operator -from sympy import symbols, sympify from PyTorchSimFrontend import extension_config from PyTorchSimFrontend import extension_codecache from PyTorchSimFrontend.mlir.mlir_codegen_backend import MLIRKernel @@ -13,8 +9,6 @@ from torch._inductor.scheduler import BaseScheduling, FusedSchedulerNode, SchedulerNode, BaseSchedulerNode from torch._inductor.utils import IndentedBuffer from torch._inductor.virtualized import V -from torch._inductor.ir import LoopBody -from torch._inductor import dependencies from torch._inductor.codegen.common import BackendFeature from . import mlir_common @@ -36,33 +30,12 @@ def __init__(self, scheduler): self.max_fusion_size = 5 def can_fuse_with_exceptions(self, node1: BaseSchedulerNode, node2: BaseSchedulerNode) -> bool: + # Monkey patch: Inductor's own can_fuse never considers prologue fusion into a + # template, so intercept that one shape and ask the template about it. if not extension_config.CONFIG_FUSION_PROLOGUE: return self.scheduler.can_fuse_origin(node1, node2) - - # Extract base template node - base_template_node1 = [node for node in node1.get_nodes() if node.is_template()] - base_template_node2 = [node for node in node2.get_nodes() if node.is_template()] - - # Case 3: Prologue(Pointwise) + Tempalte - if len(base_template_node1) == 0 and len(node1.get_nodes())==1 and len(node2.get_nodes())==1 and not node1.is_reduction() and len(base_template_node2) == 1 and extension_config.CONFIG_FUSION_PROLOGUE: - target_node = base_template_node2[0].node - - # Check if template supports prologue fusion - if not getattr(target_node.template, 'support_prologue_fusion', False): - return False - - if len(node1.read_writes.writes) != 1: - return False - if node1.node not in target_node.inputs or any(["view" in str(ori) for ori in node1.node.origins]): #FIXME - return False - - # We don't fuse this edge case... - if base_template_node2[0].group[1][0][0] == 1: - return False - - if list(node1.read_writes.writes)[0].name in [dep.name for dep in node2.read_writes.reads]: - node1 = self.revert_group(node1) - return True + if self._single_prologue_template_pair(node1, node2) is not None: + return self._prologue_plan_for(node1, node2) is not None return self.scheduler.can_fuse_origin(node1, node2) @@ -82,6 +55,51 @@ def can_fuse_vertical(self, node1, node2): def can_fuse_multi_outputs_template(self, node1, node2): return self.can_fuse_horizontal(node1, node2) + def _single_template_epilogue_pair(self, node1, node2): + """node1 is a lone template node and node2 a lone non-template node -> that template node.""" + templates = [n for n in node1.get_nodes() if n.is_template()] + if len(templates) != 1 or len(node1.get_nodes()) != 1: + return None + if len(node2.get_nodes()) != 1 or any(n.is_template() for n in node2.get_nodes()): + return None + return templates[0] + + def _single_prologue_template_pair(self, node1, node2): + """node1 is a lone non-template node feeding a lone template node2 -> that template node.""" + if len(node1.get_nodes()) != 1 or any(n.is_template() for n in node1.get_nodes()): + return None + if node1.is_reduction(): + return None + templates = [n for n in node2.get_nodes() if n.is_template()] + if len(templates) != 1 or len(node2.get_nodes()) != 1: + return None + if not ({d.name for d in node1.read_writes.writes} & {d.name for d in node2.read_writes.reads}): + return None + return templates[0] + + def _epilogue_plan_for(self, node1, node2): + """The FusionPlan the template approved for this pair, or None. Pure.""" + template_node = self._single_template_epilogue_pair(node1, node2) + if template_node is None: + return None + return template_node.node.template.try_fuse_epilogue(node1, [], node2) + + def _prologue_plan_for(self, node1, node2): + """The FusionPlan the template approved for prologue-fusing node1 into node2. Pure.""" + template_node = self._single_prologue_template_pair(node1, node2) + if template_node is None: + return None + return template_node.node.template.try_fuse_prologue(node2, node1) + + def fuse(self, node1, node2): + # Commit hook: templates defer their node mutations into the plan so that + # can_fuse stays a pure predicate. Recomputing the plan here (it is pure and + # cheap) avoids caching it across the many speculative can_fuse calls. + plan = self._epilogue_plan_for(node1, node2) or self._prologue_plan_for(node1, node2) + if plan is not None and plan.remap is not None: + plan.remap() + return super().fuse(node1, node2) + def can_fuse_horizontal(self, node1, node2): if not extension_config.CONFIG_FUSION: return False @@ -104,10 +122,6 @@ def can_fuse_horizontal(self, node1, node2): if '_unsafe_index' in node1.get_nodes()[0].node.origins or "_unsafe_index" in node2.get_nodes()[0].node.origins: return False - # Extract base template node - base_template_node1 = [node for node in node1.get_nodes() if node.is_template()] - base_template_node2 = [node for node in node2.get_nodes() if node.is_template()] - # Case 0: Reduction fusion if ( node1.is_reduction() @@ -124,122 +138,17 @@ def can_fuse_horizontal(self, node1, node2): ) return same_iter and no_dependency - # Case 1: Template + Pointwise fusion - if len(base_template_node1) == 1 and len(node1.get_nodes())==1 and len(node2.get_nodes())==1 and len(base_template_node2) == 0 and not node2.is_reduction(): - # Don't fuse maxpool template code - from PyTorchSimFrontend.mlir.mlir_maxpool_template import MLIRMaxPoolTemplate - - template_node = base_template_node1[0] - epilogue_node = node2 - - # Check if template supports epilogue fusion - if not getattr(template_node.node.template, 'support_epilogue_fusion', False): - return False - - if isinstance(template_node.node.template, MLIRMaxPoolTemplate): - return False - - # Pointwise check - v1_total = math.prod(vars1) if len(vars1) else 0 - v2_total = math.prod(vars2) if len(vars2) else 0 - if v1_total != v2_total: - return False - - # Pattern check: check data dependency between act_node and template_node - template_sched_nodes = list(template_node.get_nodes()) - # Buffers produced by the template (its outputs) - template_writes = { - dep - for n in template_sched_nodes - for dep in n.read_writes.writes - } - # Buffers still required by the activation node (unmet) or read by it - epilogue_unmet = { dep for dep in epilogue_node.unmet_dependencies } - has_dependency = bool(template_writes) and epilogue_unmet.issubset(template_writes) and not bool(reads1 & writes2) - if not has_dependency: - return False - - # Revert act_node.group : simplify_and_reorder() modified _body, _size, group - if template_node.group != epilogue_node.group: - # We don't fuse this case... - if getattr(template_node.node.template, 'support_prologue_fusion', False) and template_node.group[1][0][0] == 1: - return False - - if list(template_node.group[1][0]) != list(epilogue_node.get_nodes()[0].node.data.get_size()): - return False - self.revert_group(epilogue_node) - return True - - # Case 2: Tempalte + Reduction fusion - if len(base_template_node1) == 1 and len(node1.get_nodes())==1 and len(node2.get_nodes())==1 and len(base_template_node2) == 0 and node2.is_reduction() and extension_config.CONFIG_FUSION_REDUCTION_EPILOGUE: - target_node = base_template_node1[0].node - - # Check if template supports reduction fusion - if not getattr(target_node.template, 'support_reduction_fusion', False): - return False - - size_match = node1.get_nodes()[0].node.get_numel() == reduce(operator.mul, node2.get_nodes()[0].node.get_size(), 1) * reduce(operator.mul, node2.get_nodes()[0].node.get_reduction_size(), 1) - target_symbol = symbols("r0_0") - try: - stride = [i.strip()[:-1].split(",")[-1].strip() for i in str(node2.get_nodes()[0].node).split("\n") if "r0" in i][1] - stride = int(sympify(stride).coeff(target_symbol)) - except: - return False - - # We can't fuse dim=-1 & N == 1 - layout_possible = stride != 1 and (1 not in node1.node.get_size()) - # Directed linked? - dependency_check = writes1 & reads2 - dependency_size = all([i.get_numel() == node1.get_nodes()[0].node.get_numel() for i in node2.read_writes.reads]) - return size_match and layout_possible and dependency_check and dependency_size - - # Case 3: Prologue(Pointwise) + Tempalte - # if len(base_template_node1) == 0 and len(node1.get_nodes())==1 and not node1.is_reduction() and len(base_template_node2) == 1 and extension_config.CONFIG_FUSION_PROLOGUE: - # from PyTorchSimFrontend.mlir.mlir_gemm_template import MLIRGemmTemplate - # from PyTorchSimFrontend.mlir.mlir_bmm_template import MLIRBMMTemplate - - # target_node = base_template_node2[0].node - # # Currently only BMM, MM support prologue fusion - # if not isinstance(target_node.template, (MLIRBMMTemplate, MLIRGemmTemplate)): - # return False - - # if len(node1.read_writes.writes) != 1: - # return False - # if node1.node not in target_node.inputs or any(["view" in str(ori) for ori in node1.node.origins]): #FIXME - # return False - - # # We don't fuse this edge case... - # if base_template_node2[0].group[1][0][0] == 1: - # return False - - # if list(node1.read_writes.writes)[0].name in [dep.name for dep in node2.read_writes.reads]: - # node1 = self.revert_group(node1) - # return True + # Template + epilogue (pointwise or reduction): every condition is the template's. + if self._single_template_epilogue_pair(node1, node2) is not None: + return self._epilogue_plan_for(node1, node2) is not None + return False def revert_group(self, act_nodes, args=None, var_ranges=None): - for act_node in act_nodes.get_nodes(): - if args is None or var_ranges is None: - args, var_ranges = dependencies.index_vars_no_squeeze( - act_node.node.data.get_size(), act_node.node.data.get_reduction_size(), prefix="q" - ) - body = LoopBody( - act_node.node.get_store_function(), - (args if act_node.node.get_reduction_type() else args[:1]), - var_ranges, - args[0], - args[1] - ) - index_size = [] - reduce_size = [] - for v, s in var_ranges.items(): - if v in args[0]: - index_size.append(s) - else: - reduce_size.append(s) - node_device = act_node.get_device() - ranges = (index_size, reduce_size) - act_node._sizes, act_node._body, act_node.group = (ranges), body, (node_device, self.group_fn(ranges)) + # Used by axis-split to re-trace a node over split ranges. Fusion reaches the same + # re-derivation through FusionPlan.remap, so it never runs from a can_fuse predicate. + from PyTorchSimFrontend.mlir.mlir_template import realign_node_group + realign_node_group(act_nodes, args, var_ranges) def group_fn(self, sizes): return tuple(tuple(map(V.graph.sizevars.simplify, s)) for s in sizes) diff --git a/PyTorchSimFrontend/mlir/mlir_template.py b/PyTorchSimFrontend/mlir/mlir_template.py index 277c4189..2979f704 100644 --- a/PyTorchSimFrontend/mlir/mlir_template.py +++ b/PyTorchSimFrontend/mlir/mlir_template.py @@ -10,7 +10,9 @@ import operator from collections import OrderedDict -from typing import List, Optional +import logging +from dataclasses import dataclass +from typing import Callable, List, Optional from unittest.mock import patch from PyTorchSimFrontend import extension_config @@ -92,6 +94,65 @@ def as_local(self): finally: self.restore_buffers() +fusion_log = logging.getLogger(__name__) + + +@dataclass +class FusionPlan: + """What the scheduler must do to commit a fusion that a template approved. + + `MLIRTemplate.try_fuse_*` is a pure predicate (the scheduler calls it speculatively + for many candidate pairs), so any node mutation the fusion needs -- a loop merge, a + group re-derivation -- is deferred here and applied by `MLIRScheduling.fuse()` once + the fusion is actually committed. + """ + remap: Optional[Callable] = None # zero-arg thunk, invoked from fuse() + + +def realign_node_group(node_to_realign, args=None, var_ranges=None): + """Re-derive a node's loop body/sizes/group from its data size. + + `simplify_and_reorder()` may have reordered the node's loops away from the template's + iteration order; rebuilding the body puts them back. Used as a `FusionPlan.remap` (so + it runs from `MLIRScheduling.fuse()`, never from a can_fuse predicate) and by + axis-split, which re-traces a node over split ranges. + """ + from torch._inductor.ir import LoopBody + from torch._inductor import dependencies + from torch._inductor.virtualized import V as _V + + for node in node_to_realign.get_nodes(): + if args is None or var_ranges is None: + args, var_ranges = dependencies.index_vars_no_squeeze( + node.node.data.get_size(), node.node.data.get_reduction_size(), prefix="q") + body = LoopBody( + node.node.get_store_function(), + (args if node.node.get_reduction_type() else args[:1]), + var_ranges, args[0], args[1]) + index_size, reduce_size = [], [] + for v, s in var_ranges.items(): + (index_size if v in args[0] else reduce_size).append(s) + ranges = (index_size, reduce_size) + group = tuple(tuple(map(_V.graph.sizevars.simplify, s)) for s in ranges) + node._sizes, node._body, node.group = ranges, body, (node.get_device(), group) + + +def merge_node_loops(node_to_merge): + """Collapse a node's contiguous loops so its iteration fits the template's frame.""" + node_to_merge.get_nodes()[0].merge_loops() + + +def why_no_fuse(template_node, node_to_fuse, reason): + """Record why a fusion was declined, and return the 'declined' value (None). + + Declining is a normal outcome -- upstream CUTLASS/CPP decline every reduction + epilogue outright -- so every decline carries a reason to keep the decision + debuggable (cf. Inductor's WhyNoFuseNames). + """ + fusion_log.debug("cannot fuse %s into template %s: %s", node_to_fuse, template_node, reason) + return None + + class MLIRTemplateKernel(MLIRKernel, BaseMLIRHardwareInfo): def __init__(self, kernel_name, @@ -347,7 +408,13 @@ def conv_multi_tile_mapping(self, M, N, K, K_H, K_W, O_H, O_W, stride, dilation, max_used_spad_size = used_spad_size max_k_h_w = k_h mapping = (k_h, K_W, o_h, o_w, M, N, K) - if max_used_spad_size == 0: + # Raise only when NO tile fits SPAD. The max_used_spad_size / max_k_h_w + # tracking above only records the largest-k_h tile (max_k_h_w starts at + # K_W, so it never fires unless the full-kernel tile fits) and is not + # returned -- the sorted tile_candidates is. Guarding on it wrongly + # rejected convs whose full-kernel tile overflows SPAD (e.g. a batched + # CLIP patch conv) even though smaller-k_h tiles fit. See issue #252. + if not tile_candidates: raise RuntimeError("Cannot find a valid mapping") tile_candidates = sorted(tile_candidates, key=lambda x: x[0], reverse=True) tile_candidates = [v for _, v in tile_candidates] @@ -383,7 +450,7 @@ def conv_single_batch_mapping(self, M, N, K, K_H, K_W, O_H, O_W, stride, dilatio max_used_spad_size = used_spad_size max_k_h_w = k_h * k_w mapping = (k_h, k_w, o_h, M, M, N, K) - if max_used_spad_size == 0: + if not tile_candidates: raise RuntimeError("Cannot find a valid mapping") tile_candidates = sorted(tile_candidates, key=lambda x: x[0], reverse=True) tile_candidates = [v for _, v in tile_candidates] @@ -1389,6 +1456,12 @@ def call_name(self) -> str: class MLIRTemplate(KernelTemplate): index_counter = itertools.count() + # Maps a fused reduction epilogue's loop indices onto this template's tile axes. + # `render()` hands it to the kernel, and its LENGTH is the number of loop ranges the + # epilogue codegen can address (set_ranges asserts len(dim_aliasing) == len(ranges)). + # None means this template cannot absorb a reduction epilogue. + REDUCTION_EPILOGUE_ALIASING = None + def __init__(self, name, input_nodes, layout, input_reorder = None): """ Baseclass for MLIR Templates, derived from KernelTemplate. Not to be instantiated directly. @@ -1410,7 +1483,147 @@ def __init__(self, name, input_nodes, layout, input_reorder = None): # Fusion support flags (default to False) self.support_epilogue_fusion = False self.support_prologue_fusion = False - self.support_reduction_fusion = False + + def try_fuse_epilogue(self, template_node, existing_epilogues, node_to_fuse): + """Decide whether `node_to_fuse` can be fused as this template's epilogue. + + Only the template knows its output coordinate frame, so every fusion condition + that depends on it lives here -- the scheduler only identifies the template and + the direction, then asks. MUST BE PURE: the scheduler calls this speculatively + for many candidate pairs, so it may not mutate any node. Anything that has to + mutate goes into the returned plan's `remap`, which the scheduler applies from + `fuse()` when the fusion is actually committed. + + `existing_epilogues` are the epilogues already fused onto this template (always + empty for now; the argument is here so chained epilogue fusion can be enabled + later without changing every implementation). + + Returns a FusionPlan to fuse, or None to decline (declining is a normal result: + upstream CUTLASS/CPP decline every reduction epilogue outright).""" + if existing_epilogues: + return why_no_fuse(template_node, node_to_fuse, "chained epilogue fusion is not supported yet") + if node_to_fuse.is_reduction(): + if self.REDUCTION_EPILOGUE_ALIASING is None: + return why_no_fuse(template_node, node_to_fuse, "template does not fuse reduction epilogues") + return self._try_fuse_reduction_epilogue(template_node, node_to_fuse) + if not self.support_epilogue_fusion: + return why_no_fuse(template_node, node_to_fuse, "template does not support epilogue fusion") + return self._try_fuse_pointwise_epilogue(template_node, node_to_fuse) + + def _try_fuse_reduction_epilogue(self, template_node, epi): + """Frame-independent conditions for absorbing a reduction epilogue, plus the one + frame-dependent question: does it fit REDUCTION_EPILOGUE_ALIASING? Pure.""" + def why(reason): + return why_no_fuse(template_node, epi, reason) + + if not extension_config.CONFIG_FUSION_REDUCTION_EPILOGUE: + return why("reduction-epilogue fusion is disabled") + if epi.has_aliasing_or_mutation(): + return why("epilogue has aliasing or mutation") + + tnode = template_node.get_nodes()[0].node + esched = epi.get_nodes()[0] + enode = esched.node + + # The epilogue must iterate exactly the template output (rows x reduced cols). + if tnode.get_numel() != math.prod(enode.get_size()) * math.prod(enode.get_reduction_size()): + return why("epilogue does not iterate the whole template output") + if not all(d.get_numel() == tnode.get_numel() for d in epi.read_writes.reads): + return why("epilogue reads a buffer whose size differs from the template output") + + # It must read the template output at exactly one index expression (the CPP backend's + # template_fusion_with_epilogues_supported applies the same rule). Read the expressions + # off the LoopBody: a MemoryDep's index is normalized to a flat contiguous access and + # no longer carries the per-axis strides. + body = esched._body + template_bufs = {d.name for d in template_node.read_writes.writes} + read_exprs = {e for name in template_bufs for e in body.get_all_read_expr(name)} + if not read_exprs: + return why("epilogue does not read the template output") + if len(read_exprs) != 1: + return why("epilogue reads the template output at more than one index") + index = next(iter(read_exprs)) + + # The reduced axis must not be the contiguous (stride-1) column, and the output must + # not carry a size-1 dim. + reduce_vars = body.vars[1] + if len(reduce_vars) != 1: + return why("more than one reduction axis") + if index.coeff(reduce_vars[0]) == 1: + return why("reduces the contiguous (stride-1) axis") + if 1 in template_node.node.get_size(): + return why("template output has a size-1 dimension") + + # The only frame-dependent question: does the epilogue's iteration fit the coordinate + # map this template's reduction-epilogue codegen addresses? `set_ranges` asserts + # len(dim_aliasing) == len(ranges), so the map's length IS the frame's capacity, and + # checking it here is that assert, raised as a decline instead of a crash. See #255. + addressable = len(self.REDUCTION_EPILOGUE_ALIASING) + needed = len(esched._sizes[0]) + len(reduce_vars) + if needed <= addressable: + return FusionPlan() + + # It does not fit as-is. merge_loops() returns a NEW body, so ask whether merging the + # node's contiguous loops would make it fit -- without mutating anything. + merged = len(body.merge_loops().sizes[0]) + len(reduce_vars) + if merged > addressable: + # e.g. ConvNextV2's channels-first LayerNorm reduces the middle C axis, so batch + # and spatial are not contiguous and no loop merge can flatten them. + return why(f"epilogue needs {needed} loop ranges ({merged} after a loop merge) but " + f"this template's reduction epilogue addresses {addressable}") + return FusionPlan(remap=functools.partial(merge_node_loops, epi)) + + def _try_fuse_pointwise_epilogue(self, template_node, epi): + """Generic (frame-independent) pointwise-epilogue conditions. Pure.""" + def why(reason): + return why_no_fuse(template_node, epi, reason) + + # The epilogue must sweep exactly as many elements as the template output. + tmpl_iter, epi_iter = template_node.group[1][0], epi.group[1][0] + if (math.prod(tmpl_iter) if len(tmpl_iter) else 0) != (math.prod(epi_iter) if len(epi_iter) else 0): + return why("epilogue iterates a different number of elements than the template") + + # Everything the epilogue still needs must come from the template's outputs, and + # it must not overwrite anything the template reads. + template_writes = {dep for n in template_node.get_nodes() for dep in n.read_writes.writes} + if not template_writes: + return why("template writes nothing") + if not {dep for dep in epi.unmet_dependencies}.issubset(template_writes): + return why("epilogue depends on buffers the template does not produce") + if {d.name for d in template_node.read_writes.reads} & {d.name for d in epi.read_writes.writes}: + return why("epilogue overwrites a buffer the template reads") + + if template_node.group == epi.group: + return FusionPlan() + + # simplify_and_reorder() moved the epilogue's loops; they can be put back only if + # its data size still matches the template's iteration. + if self.support_prologue_fusion and template_node.group[1][0][0] == 1: + return why("degenerate first iteration dim on a prologue-fusing template") + if list(template_node.group[1][0]) != list(epi.get_nodes()[0].node.data.get_size()): + return why("epilogue loop order cannot be realigned to the template") + return FusionPlan(remap=functools.partial(realign_node_group, epi)) + + def try_fuse_prologue(self, template_node, node_to_fuse): + """Prologue counterpart of `try_fuse_epilogue`. Same purity contract. + + The scheduler has already established that `node_to_fuse` is a lone non-template + node feeding this lone template node.""" + def why(reason): + return why_no_fuse(template_node, node_to_fuse, reason) + + if not self.support_prologue_fusion: + return why("template does not support prologue fusion") + if len(node_to_fuse.read_writes.writes) != 1: + return why("prologue writes more than one buffer") + target = template_node.get_nodes()[0].node + if node_to_fuse.node not in target.inputs: + return why("prologue output is not a template input") + if any("view" in str(origin) for origin in node_to_fuse.node.origins): + return why("prologue is a view") + if template_node.group[1][0][0] == 1: + return why("template has a degenerate first iteration dim") + return FusionPlan(remap=functools.partial(realign_node_group, node_to_fuse)) def generate(self, **kwargs) -> ChoiceCaller: kernel_name = f"mlir_{self.name}" diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py index d32ada38..ccaca976 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_vcix.py @@ -1,7 +1,7 @@ """Python port of the C++ `-test-pytorchsim-to-vcix` conversion pass (TestPyTorchSimToVCIXConversion.cpp). -Lowers `linalg.matmul` and the transcendental math ops (exp/erf/tanh/sin/cos) to +Lowers `linalg.matmul` and the transcendental math ops (exp/erf/tanh/sin/cos/log/atan) to VCIX dialect ops (RISC-V vector custom instructions). The C++ pass is a dialect-conversion (`applyPartialConversion`); the MLIR Python bindings expose no conversion framework, so each matchAndRewrite is reimplemented as imperative IR @@ -14,7 +14,7 @@ `allow_unregistered_dialects` -- so emitting generic vcix ops here is consistent with the current pipeline. -Covers all 6 C++ patterns: linalg.matmul (gemm + conv2d) and exp/erf/tanh/sin/cos. +Covers all 8 C++ patterns: linalg.matmul (gemm + conv2d) and exp/erf/tanh/sin/cos/log/atan. Wired into extension_codecache (run_to_vcix) after fine-grained, before the standard lowering; mlir-opt then runs only -test-loop-padding. Validated structurally against `mlir-opt -test-pytorchsim-to-vcix` (non-constant ops byte-identical incl. dma_wait tag @@ -31,15 +31,21 @@ from ._mlir_util import walk_ops, i32, i64, attr_bool -MARKERS = ("linalg.matmul", "math.exp", "math.erf", "math.tanh", "math.sin", "math.cos") +MARKERS = ("linalg.matmul", "math.exp", "math.erf", "math.tanh", "math.sin", + "math.cos", "math.log", "math.atan") # math op name -> (opcode, imm) for the vcix.v.iv lowering (mirror Math*ToVCIX). +# The sf.vc.v.iv opcode field is only 2 bits (uimm2, 0-3): erf(0)/tanh(1)/ +# sin,cos,log,atan(2)/exp(3). Opcode 2 is shared and disambiguated by imm +# (sin=0, cos=1, log=2, atan=3), matching the spike rs1 / gem5 VS1 sub-opcodes. _MATH_VIV = { "math.exp": (0b000011, 0), "math.erf": (0b000000, 0), "math.tanh": (0b000001, 0), "math.sin": (0b000010, 0), "math.cos": (0b000010, 1), + "math.log": (0b000010, 2), + "math.atan": (0b000010, 3), } diff --git a/README.md b/README.md index c2298376..8b90b87b 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,9 @@ PyTorchSim **supports**: | Stable-diffusion v1 | 🤗 | ✅ | | | Llama 2/3 | 🤗 | ✅ | `tests/models/Llama/` (blocks & decode-style paths) | | DeepSeek-V3 (base) | 🤗 | ✅ | `tests/models/DeepSeek/` — several ops(e.g., gate ops) are not cycle-modeled | +| SwinV2 | 🤗 | ✅ | `tests/models/test_swinv2.py` (shifted-window attention) | +| CLIP (vision) | 🤗 | ✅ | `tests/models/test_clip.py` | +| ConvNeXt V2 | 🤗 | ✅ | `tests/models/test_convnextv2.py` (channels-first LayerNorm, depthwise conv) | | Llama-4 | 🤗 | ⏳ | In development | | Broader model support | — | ⏳ | In development |