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/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]) 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 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}") 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/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 |