From 13aaae9655dea353f454669480d5a1db4b4a092f Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 17:33:35 -0700 Subject: [PATCH 1/3] transpiler: op routing, fusion & JAX-capture correctness aten_ops.py: match the longest op prefix first so minimum/maximum/slice_scatter/select_scatter/addmm aren't shadowed by a shorter prefix (origin mis-routed addmm->add, minimum->min, maximum->max, slice_scatter->slice, select_scatter->index). importers.py: fail closed on scaled addmm (beta/alpha != 1) instead of silently dropping the scale. fusion/rms_norm.py: only fuse RMSNorm when the mean reduces the last dim, so a non-last-dim mean-square chain isn't mis-fused into a last-dim kernel with a wrong-sized weight. capture_jax.py: stride-aware slice aliasing so a strided slice (e.g. x[::2]) isn't treated as a no-op alias, plus avoid registering a duplicate gather-squeeze reshape. Signed-off-by: Noah Cylich --- python/cactus/transpile/aten_ops.py | 11 ++++++++- python/cactus/transpile/capture_jax.py | 28 ++++++++++++---------- python/cactus/transpile/fusion/rms_norm.py | 19 +++++++++++++++ python/cactus/transpile/importers.py | 5 ++++ 4 files changed, 49 insertions(+), 14 deletions(-) diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py index 4cb8967f8..90cb587ec 100644 --- a/python/cactus/transpile/aten_ops.py +++ b/python/cactus/transpile/aten_ops.py @@ -123,6 +123,8 @@ def canonical_torch_op(target: Any) -> str: AtenOpRule("aten.mean", "mean"), AtenOpRule("aten.all", "min"), AtenOpRule("aten.var", "variance"), + AtenOpRule("aten.minimum", "minimum"), + AtenOpRule("aten.maximum", "maximum"), AtenOpRule("aten.min", "min"), AtenOpRule("aten.amin", "min"), AtenOpRule("aten.max", "max"), @@ -131,6 +133,8 @@ def canonical_torch_op(target: Any) -> str: AtenOpRule("aten.split_with_sizes", "split_with_sizes"), AtenOpRule("aten.chunk", "chunk"), AtenOpRule("aten.ones", "ones"), + AtenOpRule("aten.slice_scatter", "slice_scatter"), + AtenOpRule("aten.select_scatter", "select_scatter"), AtenOpRule("aten.slice", "slice"), AtenOpRule("aten.select", "index"), AtenOpRule("aten.gather", "gather"), @@ -151,6 +155,11 @@ def canonical_torch_op(target: Any) -> str: ) +_PREFIX_RULES_BY_LENGTH: tuple[AtenOpRule, ...] = tuple( + sorted(_PREFIX_RULES, key=lambda rule: len(rule.torch_prefix), reverse=True) +) + + def canonical_ir_op(target: Any) -> str: torch_op = canonical_torch_op(target) exact = _EXACT_RULES.get(torch_op) @@ -165,7 +174,7 @@ def canonical_ir_op(target: Any) -> str: "aten.lt": "less", "aten.le": "less_equal", }[torch_op] - for rule in _PREFIX_RULES: + for rule in _PREFIX_RULES_BY_LENGTH: if torch_op.startswith(rule.torch_prefix): return rule.ir_op return torch_op diff --git a/python/cactus/transpile/capture_jax.py b/python/cactus/transpile/capture_jax.py index 7033f74ca..bef4b50e3 100644 --- a/python/cactus/transpile/capture_jax.py +++ b/python/cactus/transpile/capture_jax.py @@ -1358,7 +1358,8 @@ def _node_for_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> IRNode | limits = tuple(int(value) for value in params["limit_indices"]) raw_strides = params.get("strides") strides = (1,) * len(starts) if raw_strides is None else tuple(int(value) for value in raw_strides) - changed_axes = [axis for axis, (start, limit) in enumerate(zip(starts, limits, strict=True)) if start != 0 or limit != getattr(eqn.invars[0].aval, "shape", ())[axis]] + input_shape = getattr(eqn.invars[0].aval, "shape", ()) + changed_axes = [axis for axis, (start, limit, stride) in enumerate(zip(starts, limits, strides, strict=True)) if start != 0 or limit != input_shape[axis] or stride != 1] if not changed_axes: _alias_output(ctx, eqn.outvars[0], inputs[0]) return None @@ -1400,18 +1401,19 @@ def _node_for_eqn(graph: IRGraph, ctx: _JaxImportContext, eqn: Any) -> IRNode | if index_shape and index_shape[-1] == 1: squeezed_id = f"{indices_id}__squeezed" squeeze_shape = tuple(int(dim) for dim in index_shape[:-1]) - squeeze_node = IRNode( - node_id, - "reshape", - [indices_id], - [squeezed_id], - attrs={"shape": squeeze_shape}, - meta=meta, - ) - _register_node(graph, squeeze_node, out_avals=(eqn.invars[1].aval,)) - graph.values[squeezed_id].shape = squeeze_shape - if indices_id in ctx.gather_index_aliases: - ctx.gather_index_aliases[squeezed_id] = ctx.gather_index_aliases[indices_id] + if squeezed_id not in graph.values: + squeeze_node = IRNode( + node_id, + "reshape", + [indices_id], + [squeezed_id], + attrs={"shape": squeeze_shape}, + meta=meta, + ) + _register_node(graph, squeeze_node, out_avals=(eqn.invars[1].aval,)) + graph.values[squeezed_id].shape = squeeze_shape + if indices_id in ctx.gather_index_aliases: + ctx.gather_index_aliases[squeezed_id] = ctx.gather_index_aliases[indices_id] indices_id = squeezed_id node_id = ctx.node_id("gather_embedding") indices_id = ctx.gather_index_aliases.get(indices_id, indices_id) diff --git a/python/cactus/transpile/fusion/rms_norm.py b/python/cactus/transpile/fusion/rms_norm.py index 8951e3c42..322833744 100644 --- a/python/cactus/transpile/fusion/rms_norm.py +++ b/python/cactus/transpile/fusion/rms_norm.py @@ -107,8 +107,27 @@ def _extract_rsqrt_chain(graph: IRGraph, value_id: str) -> dict[str, object] | N pow2_node = producer(graph, mean_node.inputs[0]) if pow2_node is None or pow2_node.op != "pow" or float(pow2_node.attrs.get("exponent", 0.0)) != 2.0: return None + if not _reduces_last_dim(graph, mean_node, pow2_node.inputs[0]): + return None return { "pow_input_value_id": pow2_node.inputs[0], "eps": float(add_node.attrs.get("value", 0.0)), "node_ids": (pow_node.id, add_node.id, mean_node.id, pow2_node.id), } + + +def _reduces_last_dim(graph: IRGraph, mean_node: IRNode, input_value_id: str) -> bool: + value = graph.values.get(input_value_id) + if value is None or value.shape is None: + return False + rank = len(value.shape) + if rank == 0: + return False + axis = mean_node.attrs.get("axis") + if isinstance(axis, (list, tuple)): + if len(axis) != 1: + return False + axis = axis[0] + if not isinstance(axis, int): + return False + return axis % rank == rank - 1 diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py index 17ed8d2a7..d606d32f5 100644 --- a/python/cactus/transpile/importers.py +++ b/python/cactus/transpile/importers.py @@ -986,6 +986,11 @@ def import_linear(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[in def import_addmm(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[int, ...] | None, dtype: str | None, torch_op: str) -> None: + kwargs = getattr(node, "kwargs", {}) + beta = _extract_numeric_literal(node.args[3]) if len(node.args) > 3 else _extract_numeric_literal(kwargs.get("beta")) + alpha = _extract_numeric_literal(node.args[4]) if len(node.args) > 4 else _extract_numeric_literal(kwargs.get("alpha")) + if (beta is not None and float(beta) != 1.0) or (alpha is not None and float(alpha) != 1.0): + ctx.fail(f"unsupported scaled addmm (beta/alpha != 1) for {torch_op}: {node.args!r} {kwargs!r}") ir_node = IRNode( id=node_id(node), op="addmm", From 4d4f7451ee363291c5eac34bf88f662623aca1dd Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 21:32:53 -0700 Subject: [PATCH 2/3] test: #706 reachability repros (reverted in next commit; not for merge) Runnable repros that reach each path #706 touches and show the dropped ones are unreachable: aten_ops mis-routing, scaled-addmm fail-closed, rms_norm non-last-dim mis-fusion, capture_jax strided-slice aliasing, plus the unreachable SDPA mask/dropout and no-training batchnorm cases. This commit is reverted immediately so it stays referenceable by SHA without entering the squash/merge. Signed-off-by: Noah Cylich --- .../tests/transpile/tools/repro_706/README.md | 32 +++++ .../transpile/tools/repro_706/origin_vs_pr.sh | 53 +++++++ .../transpile/tools/repro_706/repro_jax.py | 29 ++++ .../tools/repro_706/repro_pytorch.py | 134 ++++++++++++++++++ 4 files changed, 248 insertions(+) create mode 100644 python/tests/transpile/tools/repro_706/README.md create mode 100644 python/tests/transpile/tools/repro_706/origin_vs_pr.sh create mode 100644 python/tests/transpile/tools/repro_706/repro_jax.py create mode 100644 python/tests/transpile/tools/repro_706/repro_pytorch.py diff --git a/python/tests/transpile/tools/repro_706/README.md b/python/tests/transpile/tools/repro_706/README.md new file mode 100644 index 000000000..f9aff9aeb --- /dev/null +++ b/python/tests/transpile/tools/repro_706/README.md @@ -0,0 +1,32 @@ +# PR #706 reachability repros + +Scripts that reach (or show as unreachable) each code path touched by #706, on the +real capture/import pipeline. This commit is intentionally reverted in the next +commit so the repros are referenceable by SHA but are **not** part of the squash/merge. + +## Run + +```bash +# PyTorch paths (routing, batchnorm export form, scaled-addmm, SDPA mask, rms_norm) +python python/tests/transpile/tools/repro_706/repro_pytorch.py + +# JAX path (strided-slice aliasing) +pip install "jax[cpu]" +python python/tests/transpile/tools/repro_706/repro_jax.py + +# origin/main vs this branch, one file swapped at a time +bash python/tests/transpile/tools/repro_706/origin_vs_pr.sh +``` + +## What each shows + +| Path | Kept? | Repro result | +|---|---|---| +| `aten_ops` longest-prefix | keep | origin mis-routes `addmm→add`, `minimum→min`, `maximum→max`, `slice_scatter→slice`, `select_scatter→index`; a `torch.minimum/maximum` model imports to `{minimum,maximum,add}` here | +| `importers` scaled-addmm | keep | `torch.addmm(b,a,c, beta=2.0)` fails closed | +| `fusion/rms_norm` guard | keep | FP16 channel-wise (non-last) RMS fuses into a last-dim kernel with weight `(5,)` for a channel dim of 4 without the guard; guard skips it, last-dim still fuses | +| `capture_jax` strided slice | keep | `x[::2]` → origin aliases full length-6 input (`ops=[]`), branch gives length-3 `slice` | +| `lower.py` | drop | only the non-FP16 branch changes; FP16 engine never executes it | +| `import_semantics` dropout | drop | `dropout_p` always 0 after `model.eval()` | +| `import_semantics` mask | drop | `attrs["mask"]` only set for a non-tensor literal; SDPA `attn_mask` is always Tensor/None (3 forms tested, none set it) | +| `importers` batchnorm no-training | drop | `_native_batch_norm_legit_no_training` only appears after `run_decompositions`, which only feeds CoreML; the IR importer always sees 8-arg `aten.batch_norm.default` | diff --git a/python/tests/transpile/tools/repro_706/origin_vs_pr.sh b/python/tests/transpile/tools/repro_706/origin_vs_pr.sh new file mode 100644 index 000000000..b89685568 --- /dev/null +++ b/python/tests/transpile/tools/repro_706/origin_vs_pr.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Origin-vs-PR comparison for PR #706. Run from the repo root with this branch +# checked out. Swaps a single file to origin/main, runs the probe, then restores. +# The cactus package is editable-installed against this tree, so a file swap + +# fresh interpreter is the way to compare; worktrees would all import this tree. +set -euo pipefail +cd "$(git rev-parse --show-toplevel)" +BR=audit/transpiler-lowering +F=python/cactus/transpile + +restore() { git checkout -q "$BR" -- "$1"; } + +echo "########## aten_ops: ORIGIN routing (expect add/min/max/slice/index) ##########" +git checkout -q origin/main -- "$F/aten_ops.py" +python -c " +from cactus.transpile.normalize import normalize_target +import torch +for op in ['addmm','minimum','maximum','slice_scatter','select_scatter']: + print(' ORIGIN', op, '->', normalize_target(getattr(torch.ops.aten, op).default))" +restore "$F/aten_ops.py" + +echo "########## rms_norm: WITH guard, non-last must NOT fuse / last must fuse ##########" +# this branch ships the guard; show it explicitly +python -c " +import torch, torch.nn as nn +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.optimize_graph import fuse_rms_norm +def rms(x,d,w,e=1e-6): return x*torch.rsqrt(x.pow(2).mean(d,keepdim=True)+e)*w +class Chan(nn.Module): + def __init__(s): super().__init__(); s.w=nn.Parameter(torch.ones(4,1,1,dtype=torch.float16)) + def forward(s,x): return rms(x,1,s.w) +class Last(nn.Module): + def __init__(s): super().__init__(); s.w=nn.Parameter(torch.ones(8,dtype=torch.float16)) + def forward(s,x): return rms(x,-1,s.w) +for name,m,ex in [('non-last',Chan(),torch.randn(2,4,5,5,dtype=torch.float16)),('last',Last(),torch.randn(2,3,8,dtype=torch.float16))]: + g=capture_model(m.half().eval(),(ex,)).ir_graph + ch=fuse_rms_norm(g); rn=sum(1 for n in g.order if g.nodes[n].op=='rms_norm') + print(' WITH-GUARD',name,'fused=',ch,'rms_norm=',rn)" + +echo "########## capture_jax: ORIGIN aliases strided slice (expect ops=[] shape=(6,)) ##########" +git checkout -q origin/main -- "$F/capture_jax.py" +python python/tests/transpile/tools/repro_706/repro_jax.py 2>/dev/null || echo " (needs jax[cpu])" +restore "$F/capture_jax.py" + +echo "########## batch_norm: run_decompositions turns it into the no_training op ##########" +python -c " +import torch, torch.nn as nn +m=nn.BatchNorm2d(4).eval() +ep=torch.export.export(m,(torch.randn(1,4,8,8),)) +print(' default export :', [str(n.target) for n in ep.graph.nodes if 'batch_norm' in str(n.target)]) +ep2=ep.run_decompositions() +print(' after decompose :', [str(n.target) for n in ep2.graph.nodes if 'batch_norm' in str(n.target)]) +print(' (decompose only runs on the NPU->CoreML path, never into import_captured_to_ir)')" diff --git a/python/tests/transpile/tools/repro_706/repro_jax.py b/python/tests/transpile/tools/repro_706/repro_jax.py new file mode 100644 index 000000000..0b4bc56b1 --- /dev/null +++ b/python/tests/transpile/tools/repro_706/repro_jax.py @@ -0,0 +1,29 @@ +"""Reachability repro for PR #706 (JAX capture path). + + pip install "jax[cpu]" + python python/tests/transpile/tools/repro_706/repro_jax.py + +Shows the stride-aware slice-aliasing fix in capture_jax.py: a strided slice +(x[::2]) must not be treated as a no-op alias. +""" +from __future__ import annotations +import numpy as np +import jax, jax.numpy as jnp + +from cactus.transpile.capture_jax import capture_jax_function + + +def describe(tag, fn, args): + ir = capture_jax_function(fn, args) + ops = [ir.nodes[n].op for n in ir.order] + out_id, in_id = ir.outputs[0], ir.inputs[0] + out_shape = getattr(ir.values.get(out_id), "shape", None) + print(f"[{tag}] ops={ops} out_shape={out_shape} output_is_input_alias={out_id == in_id}") + + +# Strided slice over a length-6 vector -> expected length 3. +# origin: stride ignored in changed_axes -> start=0,limit=full -> aliased no-op (WRONG). +# this branch: stride!=1 -> a real slice node, length 3. +x = jnp.arange(6.0) +print("eager jax x[::2] =", np.asarray(jax.jit(lambda v: v[::2])(x))) +describe("strided-slice x[::2]", lambda v: v[::2], (x,)) diff --git a/python/tests/transpile/tools/repro_706/repro_pytorch.py b/python/tests/transpile/tools/repro_706/repro_pytorch.py new file mode 100644 index 000000000..1cd3c89a0 --- /dev/null +++ b/python/tests/transpile/tools/repro_706/repro_pytorch.py @@ -0,0 +1,134 @@ +"""Reachability repros for PR #706 (PyTorch capture/import path). + +Run from the repo root with the branch checked out: + + python python/tests/transpile/tools/repro_706/repro_pytorch.py + +Each section reaches a code path the PR touches (or shows it is unreachable), +on the *current* checkout. Pair with origin_vs_pr.sh to see the origin behaviour. +""" +from __future__ import annotations +import traceback +from collections import Counter +import torch, torch.nn as nn, torch.nn.functional as F + +from cactus.transpile.capture_pytorch import capture_model +from cactus.transpile.optimize_graph import fuse_rms_norm +from cactus.transpile.normalize import normalize_target +from cactus.transpile.import_semantics import apply_import_semantics + + +def opc(g): + return Counter(g.nodes[n].op for n in g.order) + + +def line(t): + print(f"\n========== {t} ==========") + + +# ---- 1. aten_ops longest-prefix routing (KEEP: reached by natural ops) ---- +line("1. aten_ops routing (this branch)") +for op in ["addmm", "minimum", "maximum", "slice_scatter", "select_scatter"]: + t = getattr(torch.ops.aten, op).default + print(f" normalize_target(aten.{op}) -> {normalize_target(t)}") + +class MinMax(nn.Module): + def forward(self, a, b): + return torch.minimum(a, b) + torch.maximum(a, b) +c = capture_model(MinMax(), (torch.randn(3), torch.randn(3))) +print(" MinMax model IR ops:", dict(opc(c.ir_graph))) + + +# ---- 2. BatchNorm: which aten op does the IR import path actually see? ---- +# REMOVED from PR: the no_training (7-arg) branch is unreachable here because +# the capture path never runs run_decompositions (that only happens on the +# NPU->CoreML path). torch.export emits the 8-arg aten.batch_norm.default. +line("2. eval-mode BatchNorm export form") +class ConvBN(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(3, 4, 3, padding=1) + self.bn = nn.BatchNorm2d(4) + def forward(self, x): + return self.bn(self.conv(x)) +ep = torch.export.export(ConvBN().eval(), (torch.randn(1, 3, 8, 8),)) +print(" exported BN target(s):", [str(n.target) for n in ep.graph.nodes if "batch_norm" in str(n.target)]) +print(" (run_decompositions would instead emit _native_batch_norm_legit_no_training -> only feeds CoreML)") + + +# ---- 3. scaled addmm fail-closed (KEEP: reached) ---- +line("3. scaled addmm (beta=2.0)") +class ScaledAddmm(nn.Module): + def forward(self, bias, a, b): + return torch.addmm(bias, a, b, beta=2.0) +try: + capture_model(ScaledAddmm(), (torch.randn(4), torch.randn(4, 5), torch.randn(5, 4))) + print(" imported WITHOUT failing (origin behaviour: silently drops beta)") +except Exception as e: + print(" fail-closed OK:", type(e).__name__, "->", str(e).split(":")[-1].strip()[:80]) + + +# ---- 4. SDPA literal mask -> attrs['mask']? (REMOVED: never produced) ---- +line("4. SDPA mask -> attrs['mask']? (3 natural forms)") +def check(name, mod, args): + g = capture_model(mod, args).ir_graph + apply_import_semantics(g) + for n in g.order: + nd = g.nodes[n] + if nd.op == "attention": + print(f" [{name}] attrs={sorted(nd.attrs)} mask_in_attrs={'mask' in nd.attrs} n_inputs={len(nd.inputs)}") + return + print(f" [{name}] no attention node") + +class MaskTensorArg(nn.Module): + def forward(self, q, k, v, m): + return F.scaled_dot_product_attention(q, k, v, attn_mask=m) +check("tensor-mask-as-input", MaskTensorArg(), (torch.randn(1, 1, 4, 8),) * 3 + (torch.zeros(1, 1, 4, 4),)) + +class MaskInlineConst(nn.Module): + def forward(self, q, k, v): + return F.scaled_dot_product_attention(q, k, v, attn_mask=torch.zeros(4, 4)) +check("inline-const-mask", MaskInlineConst(), (torch.randn(1, 1, 4, 8),) * 3) + +class CausalFlag(nn.Module): + def forward(self, q, k, v): + return F.scaled_dot_product_attention(q, k, v, is_causal=True) +check("is_causal-flag", CausalFlag(), (torch.randn(1, 1, 4, 8),) * 3) + + +# ---- 6. rms_norm guard (KEEP: non-last-dim mis-fuses without it) ---- +# NOTE: fuse_rms_norm skips FP32 inputs (kept unfused on purpose), so the inputs +# must be FP16 to reach the fusion -- as they are in real transpilation. +def rms(x, dim, w, eps=1e-6): + return x * torch.rsqrt(x.pow(2).mean(dim, keepdim=True) + eps) * w + +line("6a. LAST-dim RMSNorm, FP16") +class RMSLast(nn.Module): + def __init__(self): + super().__init__() + self.w = nn.Parameter(torch.ones(8, dtype=torch.float16)) + def forward(self, x): + return rms(x, -1, self.w) +g = capture_model(RMSLast().half().eval(), (torch.randn(2, 3, 8, dtype=torch.float16),)).ir_graph +changed = fuse_rms_norm(g) +print(" fuse_rms_norm changed?", changed, "| rms_norm nodes:", opc(g).get("rms_norm", 0)) + +line("6b. NON-last-dim RMSNorm (channel dim=1), FP16") +class RMSChan(nn.Module): + def __init__(self): + super().__init__() + self.w = nn.Parameter(torch.ones(4, 1, 1, dtype=torch.float16)) + def forward(self, x): + return rms(x, 1, self.w) +g = capture_model(RMSChan().half().eval(), (torch.randn(2, 4, 5, 5, dtype=torch.float16),)).ir_graph +print(" mean axes:", [(n, g.nodes[n].attrs.get("axis")) for n in g.order if g.nodes[n].op == "mean"]) +changed = fuse_rms_norm(g) +for n in g.order: + if g.nodes[n].op == "rms_norm": + wv = g.values.get(g.nodes[n].inputs[1]) + print(" fused rms_norm weight shape:", getattr(wv, "shape", None), "(channel=4, last dim=5)") +print(" fuse_rms_norm changed?", changed, + "| this branch ships the guard so changed=False (correct);", + "checkout origin/main rms_norm.py to see changed=True with weight (5,)") + +print("\n[done]") From 8dbe3bec7f602ffe3ab49029e9b5112a9a9aab7d Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 21:46:47 -0700 Subject: [PATCH 3/3] Revert "test: #706 reachability repros (reverted in next commit; not for merge)" This reverts commit 4d4f7451ee363291c5eac34bf88f662623aca1dd. Signed-off-by: Noah Cylich --- .../tests/transpile/tools/repro_706/README.md | 32 ----- .../transpile/tools/repro_706/origin_vs_pr.sh | 53 ------- .../transpile/tools/repro_706/repro_jax.py | 29 ---- .../tools/repro_706/repro_pytorch.py | 134 ------------------ 4 files changed, 248 deletions(-) delete mode 100644 python/tests/transpile/tools/repro_706/README.md delete mode 100644 python/tests/transpile/tools/repro_706/origin_vs_pr.sh delete mode 100644 python/tests/transpile/tools/repro_706/repro_jax.py delete mode 100644 python/tests/transpile/tools/repro_706/repro_pytorch.py diff --git a/python/tests/transpile/tools/repro_706/README.md b/python/tests/transpile/tools/repro_706/README.md deleted file mode 100644 index f9aff9aeb..000000000 --- a/python/tests/transpile/tools/repro_706/README.md +++ /dev/null @@ -1,32 +0,0 @@ -# PR #706 reachability repros - -Scripts that reach (or show as unreachable) each code path touched by #706, on the -real capture/import pipeline. This commit is intentionally reverted in the next -commit so the repros are referenceable by SHA but are **not** part of the squash/merge. - -## Run - -```bash -# PyTorch paths (routing, batchnorm export form, scaled-addmm, SDPA mask, rms_norm) -python python/tests/transpile/tools/repro_706/repro_pytorch.py - -# JAX path (strided-slice aliasing) -pip install "jax[cpu]" -python python/tests/transpile/tools/repro_706/repro_jax.py - -# origin/main vs this branch, one file swapped at a time -bash python/tests/transpile/tools/repro_706/origin_vs_pr.sh -``` - -## What each shows - -| Path | Kept? | Repro result | -|---|---|---| -| `aten_ops` longest-prefix | keep | origin mis-routes `addmm→add`, `minimum→min`, `maximum→max`, `slice_scatter→slice`, `select_scatter→index`; a `torch.minimum/maximum` model imports to `{minimum,maximum,add}` here | -| `importers` scaled-addmm | keep | `torch.addmm(b,a,c, beta=2.0)` fails closed | -| `fusion/rms_norm` guard | keep | FP16 channel-wise (non-last) RMS fuses into a last-dim kernel with weight `(5,)` for a channel dim of 4 without the guard; guard skips it, last-dim still fuses | -| `capture_jax` strided slice | keep | `x[::2]` → origin aliases full length-6 input (`ops=[]`), branch gives length-3 `slice` | -| `lower.py` | drop | only the non-FP16 branch changes; FP16 engine never executes it | -| `import_semantics` dropout | drop | `dropout_p` always 0 after `model.eval()` | -| `import_semantics` mask | drop | `attrs["mask"]` only set for a non-tensor literal; SDPA `attn_mask` is always Tensor/None (3 forms tested, none set it) | -| `importers` batchnorm no-training | drop | `_native_batch_norm_legit_no_training` only appears after `run_decompositions`, which only feeds CoreML; the IR importer always sees 8-arg `aten.batch_norm.default` | diff --git a/python/tests/transpile/tools/repro_706/origin_vs_pr.sh b/python/tests/transpile/tools/repro_706/origin_vs_pr.sh deleted file mode 100644 index b89685568..000000000 --- a/python/tests/transpile/tools/repro_706/origin_vs_pr.sh +++ /dev/null @@ -1,53 +0,0 @@ -#!/bin/bash -# Origin-vs-PR comparison for PR #706. Run from the repo root with this branch -# checked out. Swaps a single file to origin/main, runs the probe, then restores. -# The cactus package is editable-installed against this tree, so a file swap + -# fresh interpreter is the way to compare; worktrees would all import this tree. -set -euo pipefail -cd "$(git rev-parse --show-toplevel)" -BR=audit/transpiler-lowering -F=python/cactus/transpile - -restore() { git checkout -q "$BR" -- "$1"; } - -echo "########## aten_ops: ORIGIN routing (expect add/min/max/slice/index) ##########" -git checkout -q origin/main -- "$F/aten_ops.py" -python -c " -from cactus.transpile.normalize import normalize_target -import torch -for op in ['addmm','minimum','maximum','slice_scatter','select_scatter']: - print(' ORIGIN', op, '->', normalize_target(getattr(torch.ops.aten, op).default))" -restore "$F/aten_ops.py" - -echo "########## rms_norm: WITH guard, non-last must NOT fuse / last must fuse ##########" -# this branch ships the guard; show it explicitly -python -c " -import torch, torch.nn as nn -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.optimize_graph import fuse_rms_norm -def rms(x,d,w,e=1e-6): return x*torch.rsqrt(x.pow(2).mean(d,keepdim=True)+e)*w -class Chan(nn.Module): - def __init__(s): super().__init__(); s.w=nn.Parameter(torch.ones(4,1,1,dtype=torch.float16)) - def forward(s,x): return rms(x,1,s.w) -class Last(nn.Module): - def __init__(s): super().__init__(); s.w=nn.Parameter(torch.ones(8,dtype=torch.float16)) - def forward(s,x): return rms(x,-1,s.w) -for name,m,ex in [('non-last',Chan(),torch.randn(2,4,5,5,dtype=torch.float16)),('last',Last(),torch.randn(2,3,8,dtype=torch.float16))]: - g=capture_model(m.half().eval(),(ex,)).ir_graph - ch=fuse_rms_norm(g); rn=sum(1 for n in g.order if g.nodes[n].op=='rms_norm') - print(' WITH-GUARD',name,'fused=',ch,'rms_norm=',rn)" - -echo "########## capture_jax: ORIGIN aliases strided slice (expect ops=[] shape=(6,)) ##########" -git checkout -q origin/main -- "$F/capture_jax.py" -python python/tests/transpile/tools/repro_706/repro_jax.py 2>/dev/null || echo " (needs jax[cpu])" -restore "$F/capture_jax.py" - -echo "########## batch_norm: run_decompositions turns it into the no_training op ##########" -python -c " -import torch, torch.nn as nn -m=nn.BatchNorm2d(4).eval() -ep=torch.export.export(m,(torch.randn(1,4,8,8),)) -print(' default export :', [str(n.target) for n in ep.graph.nodes if 'batch_norm' in str(n.target)]) -ep2=ep.run_decompositions() -print(' after decompose :', [str(n.target) for n in ep2.graph.nodes if 'batch_norm' in str(n.target)]) -print(' (decompose only runs on the NPU->CoreML path, never into import_captured_to_ir)')" diff --git a/python/tests/transpile/tools/repro_706/repro_jax.py b/python/tests/transpile/tools/repro_706/repro_jax.py deleted file mode 100644 index 0b4bc56b1..000000000 --- a/python/tests/transpile/tools/repro_706/repro_jax.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Reachability repro for PR #706 (JAX capture path). - - pip install "jax[cpu]" - python python/tests/transpile/tools/repro_706/repro_jax.py - -Shows the stride-aware slice-aliasing fix in capture_jax.py: a strided slice -(x[::2]) must not be treated as a no-op alias. -""" -from __future__ import annotations -import numpy as np -import jax, jax.numpy as jnp - -from cactus.transpile.capture_jax import capture_jax_function - - -def describe(tag, fn, args): - ir = capture_jax_function(fn, args) - ops = [ir.nodes[n].op for n in ir.order] - out_id, in_id = ir.outputs[0], ir.inputs[0] - out_shape = getattr(ir.values.get(out_id), "shape", None) - print(f"[{tag}] ops={ops} out_shape={out_shape} output_is_input_alias={out_id == in_id}") - - -# Strided slice over a length-6 vector -> expected length 3. -# origin: stride ignored in changed_axes -> start=0,limit=full -> aliased no-op (WRONG). -# this branch: stride!=1 -> a real slice node, length 3. -x = jnp.arange(6.0) -print("eager jax x[::2] =", np.asarray(jax.jit(lambda v: v[::2])(x))) -describe("strided-slice x[::2]", lambda v: v[::2], (x,)) diff --git a/python/tests/transpile/tools/repro_706/repro_pytorch.py b/python/tests/transpile/tools/repro_706/repro_pytorch.py deleted file mode 100644 index 1cd3c89a0..000000000 --- a/python/tests/transpile/tools/repro_706/repro_pytorch.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Reachability repros for PR #706 (PyTorch capture/import path). - -Run from the repo root with the branch checked out: - - python python/tests/transpile/tools/repro_706/repro_pytorch.py - -Each section reaches a code path the PR touches (or shows it is unreachable), -on the *current* checkout. Pair with origin_vs_pr.sh to see the origin behaviour. -""" -from __future__ import annotations -import traceback -from collections import Counter -import torch, torch.nn as nn, torch.nn.functional as F - -from cactus.transpile.capture_pytorch import capture_model -from cactus.transpile.optimize_graph import fuse_rms_norm -from cactus.transpile.normalize import normalize_target -from cactus.transpile.import_semantics import apply_import_semantics - - -def opc(g): - return Counter(g.nodes[n].op for n in g.order) - - -def line(t): - print(f"\n========== {t} ==========") - - -# ---- 1. aten_ops longest-prefix routing (KEEP: reached by natural ops) ---- -line("1. aten_ops routing (this branch)") -for op in ["addmm", "minimum", "maximum", "slice_scatter", "select_scatter"]: - t = getattr(torch.ops.aten, op).default - print(f" normalize_target(aten.{op}) -> {normalize_target(t)}") - -class MinMax(nn.Module): - def forward(self, a, b): - return torch.minimum(a, b) + torch.maximum(a, b) -c = capture_model(MinMax(), (torch.randn(3), torch.randn(3))) -print(" MinMax model IR ops:", dict(opc(c.ir_graph))) - - -# ---- 2. BatchNorm: which aten op does the IR import path actually see? ---- -# REMOVED from PR: the no_training (7-arg) branch is unreachable here because -# the capture path never runs run_decompositions (that only happens on the -# NPU->CoreML path). torch.export emits the 8-arg aten.batch_norm.default. -line("2. eval-mode BatchNorm export form") -class ConvBN(nn.Module): - def __init__(self): - super().__init__() - self.conv = nn.Conv2d(3, 4, 3, padding=1) - self.bn = nn.BatchNorm2d(4) - def forward(self, x): - return self.bn(self.conv(x)) -ep = torch.export.export(ConvBN().eval(), (torch.randn(1, 3, 8, 8),)) -print(" exported BN target(s):", [str(n.target) for n in ep.graph.nodes if "batch_norm" in str(n.target)]) -print(" (run_decompositions would instead emit _native_batch_norm_legit_no_training -> only feeds CoreML)") - - -# ---- 3. scaled addmm fail-closed (KEEP: reached) ---- -line("3. scaled addmm (beta=2.0)") -class ScaledAddmm(nn.Module): - def forward(self, bias, a, b): - return torch.addmm(bias, a, b, beta=2.0) -try: - capture_model(ScaledAddmm(), (torch.randn(4), torch.randn(4, 5), torch.randn(5, 4))) - print(" imported WITHOUT failing (origin behaviour: silently drops beta)") -except Exception as e: - print(" fail-closed OK:", type(e).__name__, "->", str(e).split(":")[-1].strip()[:80]) - - -# ---- 4. SDPA literal mask -> attrs['mask']? (REMOVED: never produced) ---- -line("4. SDPA mask -> attrs['mask']? (3 natural forms)") -def check(name, mod, args): - g = capture_model(mod, args).ir_graph - apply_import_semantics(g) - for n in g.order: - nd = g.nodes[n] - if nd.op == "attention": - print(f" [{name}] attrs={sorted(nd.attrs)} mask_in_attrs={'mask' in nd.attrs} n_inputs={len(nd.inputs)}") - return - print(f" [{name}] no attention node") - -class MaskTensorArg(nn.Module): - def forward(self, q, k, v, m): - return F.scaled_dot_product_attention(q, k, v, attn_mask=m) -check("tensor-mask-as-input", MaskTensorArg(), (torch.randn(1, 1, 4, 8),) * 3 + (torch.zeros(1, 1, 4, 4),)) - -class MaskInlineConst(nn.Module): - def forward(self, q, k, v): - return F.scaled_dot_product_attention(q, k, v, attn_mask=torch.zeros(4, 4)) -check("inline-const-mask", MaskInlineConst(), (torch.randn(1, 1, 4, 8),) * 3) - -class CausalFlag(nn.Module): - def forward(self, q, k, v): - return F.scaled_dot_product_attention(q, k, v, is_causal=True) -check("is_causal-flag", CausalFlag(), (torch.randn(1, 1, 4, 8),) * 3) - - -# ---- 6. rms_norm guard (KEEP: non-last-dim mis-fuses without it) ---- -# NOTE: fuse_rms_norm skips FP32 inputs (kept unfused on purpose), so the inputs -# must be FP16 to reach the fusion -- as they are in real transpilation. -def rms(x, dim, w, eps=1e-6): - return x * torch.rsqrt(x.pow(2).mean(dim, keepdim=True) + eps) * w - -line("6a. LAST-dim RMSNorm, FP16") -class RMSLast(nn.Module): - def __init__(self): - super().__init__() - self.w = nn.Parameter(torch.ones(8, dtype=torch.float16)) - def forward(self, x): - return rms(x, -1, self.w) -g = capture_model(RMSLast().half().eval(), (torch.randn(2, 3, 8, dtype=torch.float16),)).ir_graph -changed = fuse_rms_norm(g) -print(" fuse_rms_norm changed?", changed, "| rms_norm nodes:", opc(g).get("rms_norm", 0)) - -line("6b. NON-last-dim RMSNorm (channel dim=1), FP16") -class RMSChan(nn.Module): - def __init__(self): - super().__init__() - self.w = nn.Parameter(torch.ones(4, 1, 1, dtype=torch.float16)) - def forward(self, x): - return rms(x, 1, self.w) -g = capture_model(RMSChan().half().eval(), (torch.randn(2, 4, 5, 5, dtype=torch.float16),)).ir_graph -print(" mean axes:", [(n, g.nodes[n].attrs.get("axis")) for n in g.order if g.nodes[n].op == "mean"]) -changed = fuse_rms_norm(g) -for n in g.order: - if g.nodes[n].op == "rms_norm": - wv = g.values.get(g.nodes[n].inputs[1]) - print(" fused rms_norm weight shape:", getattr(wv, "shape", None), "(channel=4, last dim=5)") -print(" fuse_rms_norm changed?", changed, - "| this branch ships the guard so changed=False (correct);", - "checkout origin/main rms_norm.py to see changed=True with weight (5,)") - -print("\n[done]")