passnet-pattern-fusion

---

name: passnet-pattern-fusion description: > Author PassNet pass files: write patterns that actually MATCH the FX graph, pick fusion regions (which ops to absorb into one Triton kernel), and structure multi-pass submissions with the mandatory shared-dispatch architecture. Use when creating or fixing pattern()/replacement_args()/replacement_func() files, when a pass "failed to match", or when deciding how to fuse multiple kernels/ops.

A pass file teaches the compiler: "this subgraph → my Triton kernel". Getting the PATTERN right is binary (match or 0.1); getting the FUSION REGION right decides the speedup. Kernel implementation details live in references/kernel-templates.md; numeric-fidelity recipes in passnet-triton-opt.

1. Pass file contract

import torch                      # 'import triton' etc. as needed
def pattern(a, b, ...):           # subgraph to match — torch IR, mirrors model.py
    ...
    return out                    # EXACTLY ONE returned value (tuple-of-1 ok)
def replacement_args(a, b, ...):  # pure arg shuffle, FX-traced — no math/.shape here!
    return (a, b, "route_name")   # constants (route strings/scalars) may be appended
def replacement_func():           # returns the SAME module-level function every call
    return dispatch_wrapper

Plus pass_dir/sorted_output_pass_rule_names.json: ["PassA", "PassB"] — file stems, priority order, every pass you want loaded MUST be listed.

2. The 10 hard rules (violating any ⇒ no match / crash / dropped pass)

1. ONE output. A pattern returning ≥2 values crashes the variant (harness asserts). If two intermediates are observable outside the region, write two passes that each end at one of them — never one pattern returning both.
2. Mirror model.py's exact call forms. The target graph keeps the source form; your pattern is traced with normalization (see matchability table below). Same op spelled differently = no match (torch.relu ≠ torch.nn.functional.relu; x.transpose(...) (method) ≠ torch.transpose(x, ...)).
3. Literals must be exactly equal: eps 1e-05 ≠ 1e-04, dims, shape lists, scale constants, device(type='cuda') objects in factory calls. Shape literals differing across variants ⇒ that node can't be in a shared pattern (write per-shape passes or exclude the node).
4. Containment: every node inside the pattern except the returned one must be consumed ONLY inside the pattern. If a producer feeds both your target consumer and any outside or returned value, do not hide that producer inside a larger pattern. End the region at the consumer/tail output or write separate one-output passes. If model.py returns an intermediate or feeds it to a node outside your region, end the region there (it becomes the pattern's single output) or exclude it. When an external producer has multiple observable consumers, prefer a one-output consumer-side floor pass or split regions over pulling that producer into a fused region. This is the single-output contract in practice: a producer whose value is externally visible or shared across a fanout cannot disappear inside your pass unless the one value returned by the pattern preserves every outside observation. Do not "simplify" by recomputing only the consumer you care about while hiding the shared producer; that makes the graph semantically incomplete or non-containment-safe.
5. No tmp = None statements, no prints, no asserts inside pattern. Pure alias lines in model.py (tmp_0 = in_0) create NO graph node either — skip them and use the original name's value flow.
6. replacement_args is FX-traced: only reorder/drop arguments and append literal constants. No arithmetic, no .shape, no conditionals on tensors.
7. replacement_func() must be stable (module-level function, f() is f()), and with 2+ pass files ALL must return the SAME object via the shared-dispatch module (§4).
8. No RNG ops (torch.rand*, dropout training=True) inside patterns; eval-mode dropout (training=False) is identity and safe to absorb.
9. AST restrictions (file-wide): never import torch.nn(/functional), torch.ops, torch.autograd, never alias them; in pattern bodies use full dotted torch.nn.functional.xxx(...). Outside pattern/replacement_args, only the whitelisted torch.empty/zeros/ones/full[_like]/as_tensor calls.
10. Wrapper-runtime restrictions (poison dispatch during warmup): inside the kernel wrapper, only allocator ops + metadata (.shape/.stride()/.numel()/.dim()/.device/ .dtype/.data_ptr()) + .to() + Triton launches. NO .contiguous(), no tensor math.

3. Matchability — what a pattern node can bind to

For a normal callable def pattern(...), the pattern is traced with ForceArgsTracer; the target (dynamo) graph is not. Per node:

node kind in model.py	pattern is normalized?	matchable when...	pattern should write...
method call x.view(...), x.mean(dim=-2, keepdim=True), x.to(dtype=...), x.softmax(dim=-1)	NO	always	mirror EXACTLY (incl. kwargs spelling)
C-bound function: torch.conv2d, torch.matmul, torch.cat, torch.sigmoid, torch.arange, torch.sum, F.linear, F.gelu, operator.+ - * /	NO (no inspect signature)	always	mirror EXACTLY (incl. kwargs spelling)
Python-def F.*: relu, silu, softmax, dropout, layer_norm, batch_norm, embedding, interpolate, adaptive_avg_pool2d, pad, normalize, max_pool2d...	YES → full positional, defaults filled	ONLY if model.py call is full-positional with ALL params spelled out	the same call (any style — it normalizes)

Examples from real graphs:

• torch.nn.functional.batch_norm(x, m, v, w, b, False, 0.1, 1e-05) → matchable ✓
• torch.nn.functional.dropout(x, 0.1, False, False) → matchable ✓ (and it's identity!)
• torch.nn.functional.relu(x, inplace = False) → kwargs on a Python-def op: NOT matchable ✗
• torch.nn.functional.softmax(x, 2, _stacklevel = 5) → ✗ (and _stacklevel precedes dtype in the signature — normalization can't reproduce it)
• x.softmax(dim=-1), x.mean(dim = -2, keepdim = True) → method form: matchable ✓
• torch.cat([a, b], dim = 2) → C-bound with kwargs: matchable (mirror the kwargs) ✓
• torch.nn.functional.gelu(x) / F.gelu(x, approximate='none') → C-bound: matchable ✓

Treat kwargs-form Python-level activations and functional calls as suspect boundaries until check_pattern.py proves they match. For small or low-value regions, if they block matching, start the region after that node or end just before it instead of spending evaluations on near-miss patterns.

Manual FX escape hatch for valuable kwargs regions. The harness accepts pattern as a torch.fx.GraphModule. In that case _replace_pattern uses pattern.graph directly and does NOT run ForceArgsTracer, so exact kwargs-form Python functional nodes can match. Use this only when a high-value single-output region is blocked by callable-pattern normalization and a real pre-flight matcher confirms the manual graph matches.

Build the graph inside a function literally named pattern (that function body is AST-exempt), then replace pattern with the returned GraphModule and set __signature__ to the placeholders in order:

import inspect
import operator
import torch

def pattern():
    graph = torch.fx.Graph()
    x = graph.placeholder("x")
    y = graph.placeholder("y")
    act = graph.call_function(
        torch.nn.functional.relu, args=(x,), kwargs={"inplace": True}
    )
    out = graph.call_function(operator.add, args=(act, y), kwargs={})
    graph.output(out)
    gm = torch.fx.GraphModule({}, graph, "ExactKwargPattern")
    gm.__signature__ = inspect.Signature([
        inspect.Parameter("x", inspect.Parameter.POSITIONAL_OR_KEYWORD),
        inspect.Parameter("y", inspect.Parameter.POSITIONAL_OR_KEYWORD),
    ])
    return gm

pattern = pattern()

The manual graph must still obey every other rule: one output, containment, exact targets and literals, matching replacement_args parameters, and Triton-only runtime computation. Do not use it as a speculative default; it is a repair for proven callable-pattern normalization misses.

Manual FX does not make side effects free. Treat inplace=True functional calls as side-effectful even when the graph matches: only absorb them if the replacement mutates the input exactly as eager would under repeated benchmark calls. If that is awkward or uncertain, leave the in-place node in eager and fuse the downstream side-effect-free region.

Do NOT trust memory on which ops are C-bound — VERIFY with the pre-flight checker (passnet-feedback/scripts/check_pattern.py) before every GPU eval; it traces your pattern, runs the real SubgraphMatcher against the real dynamo graph of every variant, and prints near-miss diffs when a pattern doesn't match.

4. Multi-pass architecture (mandatory for ≥2 pass files)

output_pass_replacement_func_limit: 1 keeps only ONE distinct replacement function — passes returning a different object are SILENTLY dropped (watch Loaded N passes in the log). Structure:

pass_dir/_shared_kernels.py (helper module, NOT listed in the json):

import torch
import triton
import triton.language as tl

@triton.jit
def _fused_a_kernel(...): ...
@triton.jit
def _fused_b_kernel(...): ...

def _run_a(x, w):
    out = torch.empty_like(x)
    ...launch _fused_a_kernel...
    return out

def _run_b(x):
    ...
    return out

@torch.fx.wrap
def dispatch_wrapper(*args):
    route = args[-1]
    if route == "fuse_a":
        return _run_a(args[0], args[1])
    if route == "fuse_b":
        return _run_b(args[0])
    raise ValueError(f"unknown route {route}")

each pass_dir/FuseA.py:

import torch
from pass_dir._shared_kernels import dispatch_wrapper

def pattern(x, w):
    ...
    return out

def replacement_args(x, w):
    return (x, w, "fuse_a")

def replacement_func():
    return dispatch_wrapper

• Works because the sample root is on sys.path while passes load; the shared module is imported once → one function object → nothing is dropped.
• Keep _shared_kernels.py AST-clean too (Triton + whitelisted allocators only).
• A single-pass submission may simply define its wrapper in the pass file — but starting with the shared layout costs nothing and survives growth.
• Order in the json = application order. Earlier passes consume nodes; a later pattern overlapping an earlier match will no longer find its nodes. Order independent patterns freely; order overlapping ones from largest to smallest region.

5. Choosing fusion regions

1. From the analyzer output, mark every node matchable/unmatchable for a normal callable pattern, and separately mark high-value kwargs-form Python functional nodes that may be recoverable with a manual FX GraphModule pattern.
2. Greedily grow regions over callable-matchable nodes along dataflow; STOP at: an unrecovered unmatchable node, RNG, a node whose output is consumed outside the region (unless you end the region exactly there), a conv/big-matmul you've decided to keep in aten. Keeping a heavy conv/matmul in aten and fusing its tail is the DEFAULT and usually right. The clearest case where the conv itself becomes the region ANCHOR instead of a wall: a non-overlapping conv — stride == kernel_size, kernel_size > 1, padding == 0 — which is exactly a dense patch matmul (§ kernel-templates §13). Anchoring on a 1×1 conv (already a cuBLAS-optimal pointwise GEMM), a depthwise/grouped conv, or a general matmul usually loses — default to fusing their tails and leaving the op in aten (let a completed eval overrule this if parameters look off-regime). Either way the op is C-bound and matches when you mirror its exact positional call form; never treat "it's a big cuDNN op" as "unmatchable." If the only blocker is callable-pattern normalization of kwargs-form Python functional calls and the wider region has real upside, try a manual FX pattern and require check_pattern.py proof before evaluation. For fanout, draw the boundary at the first value that has multiple external consumers. If each consumer is independently valuable, write separate one-output passes using the shared-dispatch architecture; if the consumers are only cheap layout/view work, treat them as floor/downside-cap candidates rather than performance targets. When a shared or returned producer feeds a substantial downstream consumer, consider a consumer-side region that leaves that producer outside the pattern and treats it as an input. Do not reject the whole sample for global fanout if a local backward slice from the chosen output still has exactly one output and absorbs enough compute to amortize a launch. Prefer a single-output region that absorbs several elementwise, affine, activation, residual, or normalization consumers fed by the same broadcast/layout producer when containment is legal. Repeated round evidence shows this family is more likely to pay the fixed launch/wrapper cost than isolated reductions or tiny tails. For concat-to-normalization-affine prefixes, containment can still be legal even when a downstream activation is unmatchable or in-place: return the normalized or affine prefix tensor as the one crossing value and leave the boundary consumer in eager. Do not pull an in-place activation into the pattern just to make the region look complete. When repeated duplicate compute exists, anchor the replacement on values with real users. Userless/dead duplicate nodes can structurally match but fail replacement because the matched return has no observable consumer. If the profitable opportunity is dead-work elimination, treat it as a graph-rewrite edge case and prove it with the real matcher and smoke test before spending an evaluation.
3. Each region ⇒ one pass ⇒ ideally ONE Triton kernel launch. Two kernels for one region only when a reduction's result feeds elementwise work over a DIFFERENT axis size (e.g. softmax over rows then matmul) — usually better: split into two passes.
4. Region inputs = tensors crossing into it (pattern parameters); region output = the one tensor crossing out.
5. Sanity: replacing N aten ops saves roughly Σ(their eager µs); your kernel costs ~19 µs launch + memory traffic. If the sum is < ~50 µs the region only makes sense as the floor pass / part of something bigger.
6. Layout ops at region EDGES are usually better left outside (they're ~free in eager). Layout ops INSIDE a region (between two compute ops) are absorbed for free via index arithmetic. Special case worth taking even "alone": view/unsqueeze/expand → binary op chains that WRITE a large output — aten's broadcast-elementwise kernels are often 1.5–2.5× off the bandwidth roofline, so a shape-specialized Triton kernel wins well beyond the erased python overhead (measured 2.3× on a 67M-element broadcast-sub). Don't guess: analyze_graph.py --bench gives each node's real eager µs; estimate your kernel as bytes_moved / 1.4 TB/s and compare. Do not force eager metadata-only layout/view/reshape work to materialize just to make a larger pattern. If the pattern output is a returned or externally observed view that eager can keep cheap, materializing it in a wrapper may erase the benefit of absorbing nearby elementwise work; treat that as a downside cap unless a completed eval proves otherwise.

6. Kernel wrapper conventions (poison-proof)

def _run_fused(x, w, b):
    B, N = x.shape                       # metadata: allowed
    out = torch.empty((B, N), dtype=x.dtype, device=x.device)
    grid = (triton.cdiv(B * N, 1024),)
    _kernel[grid](x, w, b, out, B * N, N,
                  x.stride(0), x.stride(1),   # pass strides — NEVER .contiguous()
                  BLOCK=1024)
    return out

• Output dtype MUST equal what eager produces (usually the input dtype; watch explicit .to(torch.float32) nodes in the region).
• If the pattern output is consumed by .view(...) outside the region, your output's contiguity must allow that view — return a fresh contiguous tensor of the eager output's shape (the normal case).
• Handle every variant's shapes/dtypes with ONE wrapper: derive sizes from .shape, never hardcode, unless the pattern itself is shape-literal-specific.
• Scalars you need (eps, slopes, dims) that appear as pattern literals: hardcode them in the kernel for that pass (they're fixed by the match) or pass via replacement_args constants.

See references/kernel-templates.md for ready-to-adapt Triton templates: elementwise chain (stride-aware), row reduction (softmax/layer_norm/mean/L2-norm), BN-inference affine, bias+activation epilogue, cat+elementwise, embedding gather, matmul+epilogue (tl.dot), 2-pass split reductions.

7. Debug: "Pass X failed to match"

1. Run check_pattern.py — it prints the dynamo node list vs your traced pattern nodes.
2. Compare node-by-node: op kind (call_function vs call_method), target identity, arg/kwarg shapes of the WRITTEN forms, literals.
3. Most frequent causes, in order: kwargs-form Python-def F.* op included in a normal callable pattern (drop the node for a cheap region, or switch to an exact manual FX pattern for a valuable one); literal mismatch across variants; intermediate consumed outside (containment); userless/dead nodes selected as the match output; wrong method-vs-function form; _stacklevel-style hidden kwargs; pattern includes a tmp = None-induced phantom; two returned values. If a producer has outside consumers, exclude that producer or split at the consumer output before retrying.
4. [PassMgrBackend] Diagnostic ... best-attempt lines in eval logs name the first mismatching node pair — read them.

8. Round-planning containment checklist

Before handing a region to kernel tuning, answer these with the real graph in mind:

• Does the proposed pattern return exactly one value?
• Are all nodes hidden inside the pattern consumed only inside the pattern, except that single returned value?
• Does any producer inside the region feed another returned value, an outside consumer, or a later observable fanout branch? If yes, exclude it or split into one-output passes.
• Is the region mainly layout/view/reshape/split fanout work? If yes, mark it as a downside-cap/floor candidate unless it bridges compute or avoids a real materialization.
• Does the proposed region force a returned view or metadata-only layout value to become an allocated output? If yes, prefer a consumer-side elementwise/affine region or a narrow floor, and require completed-evaluation proof before keeping the wider materializing pass.
• If multiple passes are needed, do they all return the same shared dispatch wrapper?

Tentative guidance from repeated round behavior: when a larger-looking region is blocked only by the single-output/external-fanout contract, it is usually better to stop with the best correct floor than to add tiny independent kernels that increase launch overhead.