From e1be975523513940e01717d339a61e390863bd0e Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 21:27:29 +0900 Subject: [PATCH 1/4] [Frontend] Let templates decide their own fusions, from the IR ConvNextV2 with batch > 1 died in codegen with "Index names length mismatch: 2 != 3". The scheduler had approved fusing a reduction epilogue into a GEMM whose 2-D (M rows x N cols) frame cannot express it, and set_ranges then found more loop ranges than the template's coordinate map has entries. It approved it because the eligibility check read the reduction's stride out of the node's repr: stride = [i.strip()[:-1].split(",")[-1].strip() for i in str(node).split("\n") if "r0" in i][1] The intent (reject a reduction over the contiguous, stride-1 axis, which the GEMM reduction template cannot codegen) was right, but picking the second line that mentions "r0" lands on the wrong index expression once the node has more than two dimensions -- so ConvNextV2's channels-first LayerNorm mean was let through. Reading the stride off the LoopBody's index expression instead makes it decline for the right reason. Note a MemoryDep's index cannot be used here: it is normalized to a flat contiguous access and no longer carries the per-axis strides. Fixing that exposed the deeper problem: only a template knows its output coordinate frame, yet can_fuse_horizontal hardcoded a case per template kind and re-guessed that frame. It also mutated nodes (revert_group) from inside a predicate the scheduler calls speculatively for every candidate pair, carried a dead isinstance(MaxPoolTemplate) special case, and declined without ever saying why. So the decision moves to the templates, following how Inductor's own CUTLASS and CPP backends do it: - MLIRTemplate.try_fuse_epilogue / try_fuse_prologue own every condition that depends on the frame -- including the config gates -- and return a FusionPlan or None. They are pure; a node mutation the fusion needs is deferred into the plan's remap and applied from MLIRScheduling.fuse(), the commit hook Inductor calls once it really fuses. The plan is recomputed there rather than cached, since can_fuse runs many times per pair and node identities change as fusions land. - The scheduler keeps only graph facts: which node carries the template, the direction, and that this is a single-node to single-node fusion. Case 1 (pointwise) and Case 2 (reduction) collapse into one delegation. - Every decline logs a reason, like Inductor's WhyNoFuseNames. Declining is a normal answer: upstream declines reduction epilogues outright. - A template declares REDUCTION_EPILOGUE_ALIASING, the coordinate map its reduction-epilogue codegen addresses. render() consumes it, and its length is exactly what set_ranges asserts against -- so the fusion gate is that assert, raised as a decline instead of a crash. Templates without one cannot absorb a reduction epilogue, which retires the support_reduction_fusion flag. GEMM's map has two entries, BMM's three. When the rows do not fit, LoopBody.merge_loops() (which returns a new body, so this stays pure) says whether merging them would; if it would, the merge becomes the plan's remap. tests/ops/fusion covers pointwise, reduction, prologue, conv and attention fusion and is unchanged, as are the kernels each of them fuses. Adds a matmul whose reduction is over the contiguous axis: it must not fuse, and wrongly fusing it returns wrong values rather than failing to compile. ConvNextV2 batch > 1 no longer hits the assertion; it now stops at an unrelated SPAD overflow, so issue #255's report is not fully closed. --- PyTorchSimFrontend/mlir/mlir_bmm_template.py | 6 +- PyTorchSimFrontend/mlir/mlir_conv_common.py | 1 - PyTorchSimFrontend/mlir/mlir_gemm_template.py | 6 +- PyTorchSimFrontend/mlir/mlir_scheduling.py | 205 +++++------------ PyTorchSimFrontend/mlir/mlir_template.py | 211 +++++++++++++++++- tests/ops/fusion/test_matmul_reduction.py | 16 ++ 6 files changed, 290 insertions(+), 155 deletions(-) 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_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_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_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 12bf7dc5..aec36e85 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, @@ -1378,6 +1439,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. @@ -1399,7 +1466,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/tests/ops/fusion/test_matmul_reduction.py b/tests/ops/fusion/test_matmul_reduction.py index d5dde2af..6be02fb8 100644 --- a/tests/ops/fusion/test_matmul_reduction.py +++ b/tests/ops/fusion/test_matmul_reduction.py @@ -25,6 +25,21 @@ def matmul_fused(a, b): test_result("Matmul Reduction Fusion activation", res[0], y[0]) test_result("Matmul Reduction Fusion reduction", res[1], y[1]) +def test_matmul_reduce_last_dim(device, M=256, N=96, K=96): + """Reducing the matmul output's contiguous (last) axis cannot be expressed in the GEMM's + 2-D reduction-epilogue frame, so the reduction has to run as its own kernel. Fusing it + anyway silently produces wrong values, so this only checks the numbers.""" + def matmul_fused(a, b): + result = torch.matmul(a, b) + return result - result.mean(dim=-1, keepdim=True) + torch.manual_seed(0) + input = torch.randn(M, K) + weight = torch.randn(K, N) + opt_fn = torch.compile(dynamic=False)(matmul_fused) + res = opt_fn(input.to(device=device), weight.to(device=device)) + y = matmul_fused(input.to("cpu"), weight.to("cpu")) + test_result("Matmul reduce over contiguous axis", res, y) + def test_matmul_var_mean(device, size=512): def matmul_fused(a, b, c): result = torch.matmul(a, b.T) @@ -76,5 +91,6 @@ def matmul_fused(a, b, c, d): if __name__ == "__main__": device = torch.device("npu:0") test_matmul_reduce(device, 3072, 512, 768) + test_matmul_reduce_last_dim(device) test_matmul_var_mean(device) test_matmul_add_var_mean(device) From 81d1916fca753d218d04bd96949a647d9f9c576b Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 22:41:41 +0900 Subject: [PATCH 2/4] [Frontend] Report which kernel overflowed the scratchpad SpadOverflowError was raised with no arguments and the only diagnostic was a logger.debug line that is off by default, so a failing compile said nothing beyond "SPAD overflow occurred." and gave no way to tell which kernel, which tiling, or by how much. load() already has the measured usage, the budget, the tile size and the kernel's origins in scope. Put them in the exception message. --- PyTorchSimFrontend/extension_codecache.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index 9618b8e8..8e0e5f78 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -208,11 +208,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]) From 0cd3511a6a39c08389dd675a3268fbc8d1f22ce4 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 22:41:41 +0900 Subject: [PATCH 3/4] [Frontend] Size the index-expression scratch buffer per lane Spad globals follow a fixed convention: the MLIR memref carries the full tile shape (all lanes), the .spad C header that Spike links declares the per-lane slice, and the gem5 header declares the full tile. codegen_spad_buffer() emits exactly that, tile_numel_per_lane and tile_numel_per_lane * vector_lane. index_expr() had the two headers swapped. The iota scratch buffer it allocates for vectorized index arithmetic declared compute_vec_size * vector_lane entries in the .spad header and compute_vec_size entries in the gem5 header. So the buffer Spike sees is vector_lane times too large, and the one gem5 sees is one element too small: the initializer stores two elements per iteration, so its last iteration writes one slot past compute_vec_size. The oversize side is what breaks compiles. Since the spad guard tightened to spad/2 for double buffering, a kernel with compute_vec_size 64 spends 64 KB of its 64 KB budget on an iota that only ever uses 64 entries. ConvNextV2 with batch > 1 fails there: its channels-first LayerNorm var_mean kernel needs 81984 bytes/lane, of which 65536 is the iota and only 16448 is real data. Declare compute_vec_size + 1 entries per lane, and vector_lane times that for gem5. The var_mean kernel now measures 16968 bytes/lane. --- PyTorchSimFrontend/mlir/mlir_codegen_backend.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 7705be4a..52a52d3f 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -932,8 +932,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}") From 9fbdf5d34af3979a4b2836080813157fff6edabc Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Thu, 9 Jul 2026 22:51:34 +0900 Subject: [PATCH 4/4] [Model] Enable ConvNeXt V2 backbone (issue #255) Add tests/models/test_convnextv2.py, a self-hosted CI job for it, and a model coverage row. batch=2 with depths=[1,1,1,1] and image_size=64 now matches CPU end to end, max abs diff 5.0e-06. The test keeps the runner on self-hosted like the other model tests: the depthwise 7x7 conv lowers to one gather kernel per tap, so the graph is wide. --- .github/workflows/pytorchsim_test.yml | 25 +++++++++- README.md | 1 + tests/models/test_convnextv2.py | 67 +++++++++++++++++++++++++++ 3 files changed, 91 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_convnextv2.py diff --git a/.github/workflows/pytorchsim_test.yml b/.github/workflows/pytorchsim_test.yml index 4d9823f8..e97ae33a 100644 --- a/.github/workflows/pytorchsim_test.yml +++ b/.github/workflows/pytorchsim_test.yml @@ -19,8 +19,9 @@ 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_swinv2 (SwinV2 -# shifted-window backbone), test_clip (CLIP vision backbone), test_diffusion -# (UNet2D simulation OOMs the hosted runner), and test_accuracy (accuracy + speedup). +# 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 @@ -803,6 +804,26 @@ jobs: -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/README.md b/README.md index a2e51a01..8b90b87b 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ PyTorchSim **supports**: | 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 |