diff --git a/docs/audit/compiler/COMPILER_AUDIT.md b/docs/audit/compiler/COMPILER_AUDIT.md index 6d2984cad..29c480ce8 100644 --- a/docs/audit/compiler/COMPILER_AUDIT.md +++ b/docs/audit/compiler/COMPILER_AUDIT.md @@ -77,8 +77,27 @@ multiple root audit documents and compiler archive files. locked by `tests/unit/test_strict_dispatch.py` (short-circuit + legacy-path tests). **SwiGLU is now derived too** (`_match_swiglu_at` handles the DAG — gate/up share %x — inside the known-chain scan) and consumed by the - executor. Still open: Target IR (C++ fusion passes) re-discovers the same - chains. + executor. **Target IR descriptor consume/emit — landed 2026-06-11.** All 7 + Apple fusion passes now *emit* a first-class fusion descriptor on the fused + call (`tessera.fusion.kernel` + `tessera.fusion.source`): the 4 chain passes + (matmul→softmax→matmul / matmul→softmax / matmul→gelu / matmul→rmsnorm) also + *consume* an upstream `tessera.fusion.intent` (source `"descriptor"` vs + `"rediscovered"`, with a Decision-#21 warning on descriptor/IR disagreement); + the 3 composite passes (swiglu / mla_decode / native_sparse_attn) emit + `source = "composite_op"` (the pre-fused op *is* the descriptor). The Python + emit-half `canonical_compile.stamp_fusion_intents(module)` stamps the intent on + each recognized chain's terminal op from the canonical `_KNOWN_FUSION_CHAINS`, + so the frontend produces descriptor-annotated IR. Lit: + `tests/tessera-ir/phase8/apple_gpu_fusion_descriptor.mlir`; Python: + `tests/unit/test_apple_fusion_descriptor.py` + `test_fusion_intent_emitter.py` + (incl. an emit↔consume contract guard). **Auto-wired 2026-06-11:** + `driver.compile_graph_module` calls `stamp_fusion_intents(module)` before + rendering the Graph IR for Apple targets (gated to `apple_gpu`/`apple_cpu`; + the descriptor is backend-agnostic so it extends when other backends consume + it), so every Apple compile now produces descriptor-annotated Graph IR that + the Target IR passes consume. The intent is stamped into the op's MLIR `attrs` + (not `kwargs`, which are the op's real call arguments in the reference/runtime + path). Loop closed end-to-end. - **Layout and binding contracts are uneven.** Graph/Schedule/Tile/Target IR need stronger dtype, layout, aliasing, and buffer-binding contracts. - **Complete claims need fixtures.** A completed backend claim should resolve to diff --git a/docs/audit/generated/test_coverage.csv b/docs/audit/generated/test_coverage.csv index 663fdb6ba..c6ad38da8 100644 --- a/docs/audit/generated/test_coverage.csv +++ b/docs/audit/generated/test_coverage.csv @@ -184,7 +184,7 @@ gated_attention,6,0,0,6,0,,needs_direct_test,category default for 'attention' gated_deltanet,7,2,0,9,0,,needs_direct_test,category default for 'attention' gather,9,3,0,12,0,fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,needs_direct_test,category default for 'layout_transform' ge,1,0,0,1,1,,covered_by_family,category default for 'comparison' -gelu,24,16,1,40,0,bf16 f16 f32 fp16 fp32,covered_by_family,category default for 'elementwise' +gelu,26,17,1,43,0,bf16 f16 f32 fp16 fp32,covered_by_family,category default for 'elementwise' gemm,69,2,6,71,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,structural_only,unclassified — defaults to structural_only gguf_export,1,0,0,1,1,,structural_only,category default for 'aot' gqa_attention,5,0,0,5,0,,needs_direct_test,category default for 'attention' @@ -252,7 +252,7 @@ mae_loss,1,0,0,1,1,,covered_by_family,category default for 'loss' map,0,0,0,0,1,,structural_only,category default for 'control_flow' masked_categorical,10,0,0,10,0,,structural_only,unclassified — defaults to structural_only masked_fill,4,0,0,4,0,fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,needs_direct_test,category default for 'layout_transform' -matmul,296,152,15,448,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,structural_only,unclassified — defaults to structural_only +matmul,309,158,15,467,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,structural_only,unclassified — defaults to structural_only max,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only max_pool,4,0,0,4,0,,needs_direct_test,category default for 'pooling' maximum,2,0,0,2,0,,structural_only,unclassified — defaults to structural_only @@ -332,7 +332,7 @@ repeat,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only reshape,4,5,0,9,0,,structural_only,unclassified — defaults to structural_only retention,6,0,1,6,0,,needs_direct_test,category default for 'attention' rfft,3,0,0,3,0,fp32,needs_direct_test,category default for 'spectral' -rmsnorm,53,9,2,62,0,bf16 fp32,needs_direct_test,category default for 'normalization' +rmsnorm,54,10,2,64,0,bf16 fp32,needs_direct_test,category default for 'normalization' rmsnorm_safe,3,5,0,8,0,,needs_direct_test,category default for 'normalization' rng_bernoulli,0,0,0,0,1,,covered_by_family,category default for 'rng' rng_beta,0,0,0,0,1,,covered_by_family,category default for 'rng' @@ -386,7 +386,7 @@ sin,1,0,0,1,1,,covered_by_family,category default for 'elementwise' sinh,1,0,0,1,1,,covered_by_family,category default for 'elementwise' slice,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only smooth_l1_loss,1,0,0,1,1,,covered_by_family,category default for 'loss' -softmax,70,35,4,105,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,needs_direct_test,category default for 'stable_reduction' +softmax,75,37,4,112,0,bf16 f16 f32 fp16 fp32 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,needs_direct_test,category default for 'stable_reduction' softmax_safe,3,4,0,7,0,,needs_direct_test,category default for 'stable_reduction' softplus,2,0,0,2,0,,covered_by_family,category default for 'elementwise' sort,3,0,0,3,0,,structural_only,unclassified — defaults to structural_only diff --git a/docs/audit/generated/test_coverage.md b/docs/audit/generated/test_coverage.md index b697ed3ce..110193114 100644 --- a/docs/audit/generated/test_coverage.md +++ b/docs/audit/generated/test_coverage.md @@ -7,7 +7,7 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by ## Headline - **449** ops in `primitive_coverage` registry. -- **2242** total Python-test references, **571** total lit-fixture references. +- **2263** total Python-test references, **581** total lit-fixture references. - **103** ops have **zero** references in either test surface. - **201** ops have ≤1 reference ("thinly tested"). - **48** ops have ≥10 references ("well tested"). @@ -17,19 +17,19 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by | Op | py refs | lit refs | total | neg | dtypes | |----|--------:|---------:|------:|----:|--------| -| `matmul` | 296 | 152 | 448 | 15 | `bf16`, `f16`, `f32`, `fp16` … | +| `matmul` | 309 | 158 | 467 | 15 | `bf16`, `f16`, `f32`, `fp16` … | | `flash_attn` | 79 | 42 | 121 | 4 | `bf16`, `f32`, `fp16`, `fp32` … | -| `softmax` | 70 | 35 | 105 | 4 | `bf16`, `f16`, `f32`, `fp16` … | +| `softmax` | 75 | 37 | 112 | 4 | `bf16`, `f16`, `f32`, `fp16` … | | `relu` | 92 | 7 | 99 | 9 | `bf16`, `f16`, `f32`, `fp32` | | `silu` | 76 | 2 | 78 | 5 | `bf16`, `f16`, `f32`, `fp16` … | | `gemm` | 69 | 2 | 71 | 6 | `bf16`, `f16`, `f32`, `fp16` … | | `add` | 62 | 6 | 68 | 8 | `bf16`, `f16`, `f32`, `fp32` | | `reduce` | 66 | 0 | 66 | 6 | `f32`, `fp16`, `fp32`, `fp4_e2m1` … | +| `rmsnorm` | 54 | 10 | 64 | 2 | `bf16`, `fp32` | | `mul` | 62 | 0 | 62 | 6 | `fp16`, `fp32`, `fp4_e2m1`, `fp6_e2m3` … | -| `rmsnorm` | 53 | 9 | 62 | 2 | `bf16`, `fp32` | | `attn_local_window_2d` | 34 | 25 | 59 | 1 | `fp32` | | `cholesky` | 15 | 30 | 45 | 0 | `bf16`, `f16`, `f32`, `fp16` … | -| `gelu` | 24 | 16 | 40 | 1 | `bf16`, `f16`, `f32`, `fp16` … | +| `gelu` | 26 | 17 | 43 | 1 | `bf16`, `f16`, `f32`, `fp16` … | | `linear_attn` | 31 | 8 | 39 | 1 | | | `grouped_gemm` | 27 | 10 | 37 | 2 | `fp32`, `fp4_e2m1`, `fp8_e4m3`, `fp8_e5m2` … | | `selective_ssm` | 36 | 0 | 36 | 2 | | diff --git a/python/tessera/compiler/canonical_compile.py b/python/tessera/compiler/canonical_compile.py index 0a586db3b..41be64e19 100644 --- a/python/tessera/compiler/canonical_compile.py +++ b/python/tessera/compiler/canonical_compile.py @@ -746,6 +746,57 @@ def _derive_compile_metadata(module: GraphIRModule) -> dict[str, Any]: } +# Linear fusion chains the Apple Target IR passes re-discover and now also +# consume from a descriptor. (swiglu / mla / nsa lower a pre-fused op — the op +# *is* the descriptor — so they don't need an intent stamp.) +_INTENT_KERNELS: frozenset[str] = frozenset({ + "matmul_softmax_matmul", "matmul_softmax", "matmul_gelu", "matmul_rmsnorm", +}) + + +def stamp_fusion_intents(module: GraphIRModule) -> int: + """Decision #19 emit-half — stamp ``tessera.fusion.intent`` on the terminal + op of each recognized linear fusion chain so the Apple Target IR fusion + passes *consume* the compiler's fusion decision (the emitted call is tagged + ``source = "descriptor"``) instead of re-discovering the chain. Returns the + number of chains stamped. Idempotent. + + The terminal op (the highest-index op of the chain — the tail matmul / + softmax / gelu / rmsnorm) is exactly where each C++ pass reads the intent. + """ + groups = _derive_fusion_groups(module) + by_fn = {fn.name: fn for fn in module.functions} + stamped = 0 + for group in groups: + if group.get("kind") != "known_chain": + continue + kernel = group.get("fused_kernel") + if kernel not in _INTENT_KERNELS: + continue + fn_name = group.get("function") + fn = by_fn.get(fn_name) if isinstance(fn_name, str) else None + if fn is None: + continue + indices = [int(e["index"]) for e in group.get("ops", ()) + if isinstance(e, Mapping) and isinstance(e.get("index"), int)] + if not indices: + continue + terminal = max(indices) + if 0 <= terminal < len(fn.body): + op = fn.body[terminal] + # Stamp into `attrs` (the MLIR-attribute-only field), NOT `kwargs`: + # kwargs are forwarded as the op's real call arguments in the + # reference/runtime execution path, so a descriptor placed there + # would leak into the numpy op call (e.g. gelu(**kwargs)). + intent_attr = f'tessera.fusion.intent = "{kernel}"' + if not op.attrs: + op.attrs = intent_attr + elif "tessera.fusion.intent" not in op.attrs: + op.attrs = f"{op.attrs}, {intent_attr}" + stamped += 1 + return stamped + + # --- Synthesizers -------------------------------------------------------- def _result_from_bundle( diff --git a/python/tessera/compiler/driver.py b/python/tessera/compiler/driver.py index cc786b42c..0add43c55 100644 --- a/python/tessera/compiler/driver.py +++ b/python/tessera/compiler/driver.py @@ -239,6 +239,17 @@ def compile_graph_module( enable_tool_validation: bool = True, ) -> CompileArtifactBundle: target_kind = normalize_target_kind(target) + # Decision #19 — auto-emit fusion descriptors. Before the Graph IR is + # rendered, stamp `tessera.fusion.intent` on each recognized fusion chain's + # terminal op so the Apple Target IR fusion passes *consume* the compiler's + # intent (the fused call is tagged source="descriptor") instead of + # re-discovering the chain. Apple-gated (the only Target IR consumer today); + # the descriptor is backend-agnostic, so this extends to other backends when + # their Target IR consumes it. Lazy import avoids the canonical_compile ↔ + # driver cycle. + if target_kind in ("apple_gpu", "apple_cpu"): + from .canonical_compile import stamp_fusion_intents + stamp_fusion_intents(module) graph_text = module.to_mlir() function_name = module.functions[0].name if module.functions else "" request = CompileRequest( diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MLADecodeFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MLADecodeFusionToAppleGPU.cpp index 8e795f77d..a9fae6694 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MLADecodeFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MLADecodeFusionToAppleGPU.cpp @@ -148,10 +148,14 @@ struct LowerMLADecodeFusionToAppleGPU : public RewritePattern { i32Ty, i32Ty, i32Ty, i32Ty, i32Ty, i32Ty}, {}); ensureExternalDecl(mod, kMLADecodeF32Symbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, kMLADecodeF32Symbol, TypeRange{}, ValueRange{xPtr, wDkvPtr, wUkPtr, wUvPtr, qPtr, oPtr, Bv, Skv, Dx, Dlat, Sq, Dh}); + // Decision #19 — emit the fusion descriptor. MLA decode lowers a pre-fused + // tessera.mla_decode_fused op (the op is the descriptor): source="composite_op". + callOp->setAttr("tessera.fusion.kernel", rewriter.getStringAttr("mla_decode")); + callOp->setAttr("tessera.fusion.source", rewriter.getStringAttr("composite_op")); auto outTensorTy = RankedTensorType::get({B, S_q, D_h}, elem); Value result = diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulGeluFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulGeluFusionToAppleGPU.cpp index 0d1b6a207..67509a49b 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulGeluFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulGeluFusionToAppleGPU.cpp @@ -54,6 +54,12 @@ struct LowerMatmulGeluFusionToAppleGPU : public RewritePattern { if (geluOp->getNumOperands() < 1) return failure(); Value geluIn = geluOp->getOperand(0); + // Decision #19 — consume the compiler's fusion descriptor when present. + StringRef intent; + if (auto a = geluOp->getAttrOfType("tessera.fusion.intent")) + intent = a.getValue(); + bool descriptorDriven = (intent == "matmul_gelu"); + auto gTy = dyn_cast(geluIn.getType()); if (!gTy || gTy.getRank() != 2) return rewriter.notifyMatchFailure(geluOp, "matmul_gelu fusion: rank-2 only"); @@ -63,8 +69,13 @@ struct LowerMatmulGeluFusionToAppleGPU : public RewritePattern { Operation *defOp = geluIn.getDefiningOp(); if (!defOp) return rewriter.notifyMatchFailure(geluOp, "matmul_gelu fusion: no defining op"); - if (defOp->getName().getStringRef() != "tessera.matmul") + if (defOp->getName().getStringRef() != "tessera.matmul") { + if (descriptorDriven) // Decision #21 — descriptor/IR disagreement. + geluOp->emitWarning( + "tessera.fusion.intent = \"matmul_gelu\" but gelu operand is not " + "from tessera.matmul — descriptor/IR mismatch; falling back"); return rewriter.notifyMatchFailure(geluOp, "matmul_gelu fusion: defining op is not tessera.matmul"); + } if (!geluIn.hasOneUse()) return rewriter.notifyMatchFailure(geluOp, "matmul_gelu fusion: matmul result has multiple uses"); @@ -123,9 +134,12 @@ struct LowerMatmulGeluFusionToAppleGPU : public RewritePattern { ctx, {i64Ty, i64Ty, i64Ty, i32Ty, i32Ty, i32Ty}, {}); ensureExternalDecl(mod, kMatmulGeluF32Symbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, kMatmulGeluF32Symbol, TypeRange{}, ValueRange{aPtr, bPtr, oPtr, Mv, Nv, Kv}); + callOp->setAttr("tessera.fusion.kernel", rewriter.getStringAttr("matmul_gelu")); + callOp->setAttr("tessera.fusion.source", + rewriter.getStringAttr(descriptorDriven ? "descriptor" : "rediscovered")); auto outTensorTy = RankedTensorType::get({M, N}, f32Ty); Value result = diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulRMSNormFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulRMSNormFusionToAppleGPU.cpp index c75ea0941..0ca4090de 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulRMSNormFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulRMSNormFusionToAppleGPU.cpp @@ -64,6 +64,12 @@ struct LowerMatmulRMSNormPatternBase : public RewritePattern { if (normOp->getNumOperands() < 1) return failure(); Value normIn = normOp->getOperand(0); + // Decision #19 — consume the compiler's fusion descriptor when present. + StringRef intent; + if (auto a = normOp->getAttrOfType("tessera.fusion.intent")) + intent = a.getValue(); + bool descriptorDriven = (intent == "matmul_rmsnorm"); + auto nTy = dyn_cast(normIn.getType()); if (!nTy || nTy.getRank() != 2) return rewriter.notifyMatchFailure(normOp, "matmul_rmsnorm fusion: rank-2 only"); @@ -73,8 +79,13 @@ struct LowerMatmulRMSNormPatternBase : public RewritePattern { Operation *defOp = normIn.getDefiningOp(); if (!defOp) return rewriter.notifyMatchFailure(normOp, "matmul_rmsnorm fusion: no defining op"); - if (defOp->getName().getStringRef() != "tessera.matmul") + if (defOp->getName().getStringRef() != "tessera.matmul") { + if (descriptorDriven) // Decision #21 — descriptor/IR disagreement. + normOp->emitWarning( + "tessera.fusion.intent = \"matmul_rmsnorm\" but norm operand is not " + "from tessera.matmul — descriptor/IR mismatch; falling back"); return rewriter.notifyMatchFailure(normOp, "matmul_rmsnorm fusion: defining op is not tessera.matmul"); + } if (!normIn.hasOneUse()) return rewriter.notifyMatchFailure(normOp, "matmul_rmsnorm fusion: matmul result has multiple uses"); @@ -140,9 +151,12 @@ struct LowerMatmulRMSNormPatternBase : public RewritePattern { ctx, {i64Ty, i64Ty, i64Ty, i32Ty, i32Ty, i32Ty, f32Ty}, {}); ensureExternalDecl(mod, kMatmulRMSNormF32Symbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, kMatmulRMSNormF32Symbol, TypeRange{}, ValueRange{aPtr, bPtr, oPtr, Mv, Nv, Kv, epsV}); + callOp->setAttr("tessera.fusion.kernel", rewriter.getStringAttr("matmul_rmsnorm")); + callOp->setAttr("tessera.fusion.source", + rewriter.getStringAttr(descriptorDriven ? "descriptor" : "rediscovered")); auto outTensorTy = RankedTensorType::get({M, N}, f32Ty); Value result = diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxFusionToAppleGPU.cpp index bc7b232c8..c47300735 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxFusionToAppleGPU.cpp @@ -73,6 +73,12 @@ struct LowerMatmulSoftmaxFusionToAppleGPU : public RewritePattern { return failure(); Value softmaxIn = softmaxOp->getOperand(0); + // Decision #19 — consume the compiler's fusion descriptor when present. + StringRef intent; + if (auto a = softmaxOp->getAttrOfType("tessera.fusion.intent")) + intent = a.getValue(); + bool descriptorDriven = (intent == "matmul_softmax"); + // axis: defaults to -1. Anything else falls out of fusion. int64_t axis = -1; if (auto attr = softmaxOp->getAttrOfType("axis")) @@ -100,8 +106,13 @@ struct LowerMatmulSoftmaxFusionToAppleGPU : public RewritePattern { Operation *defOp = softmaxIn.getDefiningOp(); if (!defOp) return rewriter.notifyMatchFailure(softmaxOp, "fusion: softmax operand has no defining op"); - if (defOp->getName().getStringRef() != "tessera.matmul") + if (defOp->getName().getStringRef() != "tessera.matmul") { + if (descriptorDriven) // Decision #21 — descriptor/IR disagreement. + softmaxOp->emitWarning( + "tessera.fusion.intent = \"matmul_softmax\" but softmax operand is " + "not from tessera.matmul — descriptor/IR mismatch; falling back"); return rewriter.notifyMatchFailure(softmaxOp, "fusion: defining op is not tessera.matmul"); + } if (!softmaxIn.hasOneUse()) return rewriter.notifyMatchFailure(softmaxOp, "fusion: matmul result has multiple uses"); @@ -163,9 +174,12 @@ struct LowerMatmulSoftmaxFusionToAppleGPU : public RewritePattern { ctx, {i64Ty, i64Ty, i64Ty, i32Ty, i32Ty, i32Ty}, {}); ensureExternalDecl(mod, symbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, symbol, TypeRange{}, ValueRange{aPtr, bPtr, oPtr, Mv, Nv, Kv}); + callOp->setAttr("tessera.fusion.kernel", rewriter.getStringAttr("matmul_softmax")); + callOp->setAttr("tessera.fusion.source", + rewriter.getStringAttr(descriptorDriven ? "descriptor" : "rediscovered")); auto outTensorTy = RankedTensorType::get({M, N}, smElem); Value result = diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxMatmulFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxMatmulFusionToAppleGPU.cpp index b5b6d4485..ab44062ee 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxMatmulFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulSoftmaxMatmulFusionToAppleGPU.cpp @@ -76,12 +76,31 @@ struct LowerMatmulSoftmaxMatmulFusionToAppleGPU : public RewritePattern { Value secondLhs = secondMatmulOp->getOperand(0); // softmax output Value secondRhs = secondMatmulOp->getOperand(1); // C + // Decision #19 — consume the compiler's fusion descriptor. When the + // canonical compile recognized this chain it stamps the tail op with + // `tessera.fusion.intent = "matmul_softmax_matmul"`; we take that as + // authoritative (the emitted call is tagged `source = "descriptor"`), + // rather than re-discovering the fusion purely structurally. Absent the + // intent, the structural walk below still recognizes the chain and the + // call is tagged `source = "rediscovered"` (back-compat). + StringRef intent; + if (auto a = secondMatmulOp->getAttrOfType("tessera.fusion.intent")) + intent = a.getValue(); + bool descriptorDriven = (intent == "matmul_softmax_matmul"); + // Walk: secondLhs must be the result of a tessera.softmax with one use. Operation *softmaxOp = secondLhs.getDefiningOp(); if (!softmaxOp) return rewriter.notifyMatchFailure(secondMatmulOp, "fusion3: tail matmul lhs has no defining op"); - if (softmaxOp->getName().getStringRef() != "tessera.softmax") + if (softmaxOp->getName().getStringRef() != "tessera.softmax") { + // Decision #21 — a fusion intent that the IR structure contradicts is a + // real inconsistency; surface it by name instead of silently falling back. + if (descriptorDriven) + secondMatmulOp->emitWarning( + "tessera.fusion.intent = \"matmul_softmax_matmul\" but tail matmul " + "lhs is not a tessera.softmax — descriptor/IR mismatch; falling back"); return rewriter.notifyMatchFailure(secondMatmulOp, "fusion3: tail matmul lhs is not a softmax"); + } if (!secondLhs.hasOneUse()) return rewriter.notifyMatchFailure(secondMatmulOp, "fusion3: softmax result has multiple uses"); @@ -187,9 +206,17 @@ struct LowerMatmulSoftmaxMatmulFusionToAppleGPU : public RewritePattern { ctx, {i64Ty, i64Ty, i64Ty, i64Ty, i32Ty, i32Ty, i32Ty, i32Ty}, {}); ensureExternalDecl(mod, symbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, symbol, TypeRange{}, ValueRange{aPtr, bPtr, cPtr, oPtr, Mv, Kv, Nv, Pv}); + // Decision #19 — emit the fusion descriptor into the Target IR so the + // fusion decision is first-class/auditable (which kernel, and whether the + // compiler's intent drove it or it was re-discovered structurally). + callOp->setAttr("tessera.fusion.kernel", + rewriter.getStringAttr("matmul_softmax_matmul")); + callOp->setAttr("tessera.fusion.source", + rewriter.getStringAttr(descriptorDriven ? "descriptor" + : "rediscovered")); auto outTensorTy = RankedTensorType::get({M, P}, elem); Value result = diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/NativeSparseAttnFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/NativeSparseAttnFusionToAppleGPU.cpp index b171ce80e..8d51d40bf 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/NativeSparseAttnFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/NativeSparseAttnFusionToAppleGPU.cpp @@ -133,10 +133,16 @@ struct LowerNSAFusionToAppleGPU : public RewritePattern { i32Ty, i32Ty, i32Ty, i32Ty}, {}); ensureExternalDecl(mod, kNSASymbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, kNSASymbol, TypeRange{}, ValueRange{qPtr, kPtr, vPtr, gPtr, oPtr, Bv, Hv, Sv, Dv, Wv, Bkv, Tkv, Cv}); + // Decision #19 — emit the fusion descriptor. NSA lowers a pre-fused + // tessera.native_sparse_attn_fused op (the op is the descriptor): + // source = "composite_op". + callOp->setAttr("tessera.fusion.kernel", + rewriter.getStringAttr("native_sparse_attn")); + callOp->setAttr("tessera.fusion.source", rewriter.getStringAttr("composite_op")); auto outTensorTy = RankedTensorType::get({B, H, S, D}, elem); Value result = diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SwigluFusionToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SwigluFusionToAppleGPU.cpp index c6c057e89..c90de0254 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SwigluFusionToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SwigluFusionToAppleGPU.cpp @@ -164,9 +164,14 @@ struct LowerSwigluFusedToAppleGPU : public RewritePattern { {}); ensureExternalDecl(mod, symbol, fnTy); - rewriter.create( + auto callOp = rewriter.create( loc, symbol, TypeRange{}, ValueRange{xPtr, wgPtr, wuPtr, wdPtr, oPtr, Mv, Kv, Hv, Kov}); + // Decision #19 — emit the fusion descriptor. swiglu lowers a pre-fused + // tessera.swiglu_fused op, so the op itself is the descriptor (no chain + // re-discovery): source = "composite_op". + callOp->setAttr("tessera.fusion.kernel", rewriter.getStringAttr("swiglu")); + callOp->setAttr("tessera.fusion.source", rewriter.getStringAttr("composite_op")); auto outTensorTy = RankedTensorType::get({M, Kout}, elem); Value result = diff --git a/tests/tessera-ir/phase8/apple_gpu_fusion_descriptor.mlir b/tests/tessera-ir/phase8/apple_gpu_fusion_descriptor.mlir new file mode 100644 index 000000000..e712053df --- /dev/null +++ b/tests/tessera-ir/phase8/apple_gpu_fusion_descriptor.mlir @@ -0,0 +1,54 @@ +// RUN: tessera-opt %s --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck %s + +// Decision #19 — the Apple Target IR fusion pass emits a first-class fusion +// descriptor on the fused call (`tessera.fusion.kernel` + `tessera.fusion.source`) +// and *consumes* the compiler's `tessera.fusion.intent` when present: +// * intent stamped (the canonical compile recognized the chain) → source="descriptor" +// * no intent (legacy IR) → the structural walk re-discovers it → source="rediscovered" +// Both fuse to the same matmul_softmax_matmul kernel. + +// CHECK-LABEL: func.func @descriptor_driven +func.func @descriptor_driven(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>, + %C: tensor<32x8xf32>) -> tensor<8x8xf32> { + // CHECK: call @tessera_apple_gpu_matmul_softmax_matmul_f32 + // CHECK-SAME: tessera.fusion.kernel = "matmul_softmax_matmul" + // CHECK-SAME: tessera.fusion.source = "descriptor" + %m1 = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + %p = "tessera.softmax"(%m1) : (tensor<8x32xf32>) -> tensor<8x32xf32> + %o = "tessera.matmul"(%p, %C) {tessera.fusion.intent = "matmul_softmax_matmul"} + : (tensor<8x32xf32>, tensor<32x8xf32>) -> tensor<8x8xf32> + return %o : tensor<8x8xf32> +} + +// CHECK-LABEL: func.func @rediscovered +func.func @rediscovered(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>, + %C: tensor<32x8xf32>) -> tensor<8x8xf32> { + // CHECK: call @tessera_apple_gpu_matmul_softmax_matmul_f32 + // CHECK-SAME: tessera.fusion.source = "rediscovered" + %m1 = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + %p = "tessera.softmax"(%m1) : (tensor<8x32xf32>) -> tensor<8x32xf32> + %o = "tessera.matmul"(%p, %C) : (tensor<8x32xf32>, tensor<32x8xf32>) -> tensor<8x8xf32> + return %o : tensor<8x8xf32> +} + +// matmul→gelu: descriptor-driven. +// CHECK-LABEL: func.func @gelu_descriptor +func.func @gelu_descriptor(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>) -> tensor<8x32xf32> { + // CHECK: call @tessera_apple_gpu_matmul_gelu_f32 + // CHECK-SAME: tessera.fusion.kernel = "matmul_gelu" + // CHECK-SAME: tessera.fusion.source = "descriptor" + %m = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + %g = "tessera.gelu"(%m) {tessera.fusion.intent = "matmul_gelu"} : (tensor<8x32xf32>) -> tensor<8x32xf32> + return %g : tensor<8x32xf32> +} + +// matmul→rmsnorm: re-discovered (no intent). +// CHECK-LABEL: func.func @rmsnorm_rediscovered +func.func @rmsnorm_rediscovered(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>) -> tensor<8x32xf32> { + // CHECK: call @tessera_apple_gpu_matmul_rmsnorm_f32 + // CHECK-SAME: tessera.fusion.kernel = "matmul_rmsnorm" + // CHECK-SAME: tessera.fusion.source = "rediscovered" + %m = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + %r = "tessera.rmsnorm"(%m) : (tensor<8x32xf32>) -> tensor<8x32xf32> + return %r : tensor<8x32xf32> +} diff --git a/tests/unit/test_apple_fusion_descriptor.py b/tests/unit/test_apple_fusion_descriptor.py new file mode 100644 index 000000000..c429704ab --- /dev/null +++ b/tests/unit/test_apple_fusion_descriptor.py @@ -0,0 +1,83 @@ +"""Apple Target IR fusion-descriptor emit + consume (Decision #19, slice 1). + +The matmul→softmax→matmul Apple fusion pass now emits a first-class fusion +descriptor on the fused call (`tessera.fusion.kernel` + `tessera.fusion.source`) +and consumes the compiler's `tessera.fusion.intent` when present, instead of +only re-discovering the chain structurally. See COMPILER_AUDIT "fusion intent is +too late". +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +REPO = Path(__file__).resolve().parents[2] +PASS_SRC = (REPO / "src" / "compiler" / "codegen" / "Tessera_Apple_Backend" / + "lib" / "Target" / "Apple" / "Lowering" / + "MatmulSoftmaxMatmulFusionToAppleGPU.cpp") +_CANDIDATES = ( + REPO / "build" / "tools" / "tessera-opt" / "tessera-opt", + REPO / "build-llvm22" / "tools" / "tessera-opt" / "tessera-opt", +) + + +def _find_opt(): + if explicit := os.environ.get("TESSERA_OPT_PATH"): + if Path(explicit).is_file(): + return explicit + for c in _CANDIDATES: + if c.is_file() and os.access(c, os.X_OK): + return str(c) + return shutil.which("tessera-opt") + + +_OPT = _find_opt() +_needs_opt = pytest.mark.skipif(_OPT is None, reason="tessera-opt not built") + +_FIXTURE = ''' +func.func @descriptor_driven(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>, %C: tensor<32x8xf32>) -> tensor<8x8xf32> { + %m1 = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + %p = "tessera.softmax"(%m1) : (tensor<8x32xf32>) -> tensor<8x32xf32> + %o = "tessera.matmul"(%p, %C) {tessera.fusion.intent = "matmul_softmax_matmul"} : (tensor<8x32xf32>, tensor<32x8xf32>) -> tensor<8x8xf32> + return %o : tensor<8x8xf32> +} +func.func @rediscovered(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>, %C: tensor<32x8xf32>) -> tensor<8x8xf32> { + %m1 = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + %p = "tessera.softmax"(%m1) : (tensor<8x32xf32>) -> tensor<8x32xf32> + %o = "tessera.matmul"(%p, %C) : (tensor<8x32xf32>, tensor<32x8xf32>) -> tensor<8x8xf32> + return %o : tensor<8x8xf32> +} +''' + + +def test_pass_emits_and_consumes_the_descriptor_in_source(): + src = PASS_SRC.read_text() + assert 'tessera.fusion.intent' in src # consumes the compiler intent + assert 'tessera.fusion.kernel' in src # emits the descriptor + assert 'tessera.fusion.source' in src + assert '"descriptor"' in src and '"rediscovered"' in src + assert 'descriptor/IR mismatch' in src # Decision #21 diagnostic + + +@_needs_opt +def test_descriptor_vs_rediscovered_source_tag(tmp_path): + f = tmp_path / "fd.mlir" + f.write_text(_FIXTURE) + out = subprocess.run( + [_OPT, str(f), + "--pass-pipeline=builtin.module(tessera-lower-to-apple_gpu-runtime)", + "--allow-unregistered-dialect"], + capture_output=True, text=True, timeout=60) + assert out.returncode == 0, out.stderr + text = out.stdout + # Both functions fuse to the same kernel ... + assert text.count('tessera.fusion.kernel = "matmul_softmax_matmul"') == 2 + # ... but the descriptor-annotated chain is tagged source="descriptor" and + # the bare chain source="rediscovered". + assert 'tessera.fusion.source = "descriptor"' in text + assert 'tessera.fusion.source = "rediscovered"' in text diff --git a/tests/unit/test_fusion_intent_emitter.py b/tests/unit/test_fusion_intent_emitter.py new file mode 100644 index 000000000..2159a67ca --- /dev/null +++ b/tests/unit/test_fusion_intent_emitter.py @@ -0,0 +1,98 @@ +"""Python fusion-intent emitter (Decision #19 emit-half). + +`stamp_fusion_intents(module)` tags the terminal op of each recognized linear +fusion chain with `tessera.fusion.intent = ""`, so the Apple Target IR +passes consume the compiler's fusion decision (source="descriptor") instead of +re-discovering it. This closes the Python→IR half of the loop; the C++ +consume-half is locked by tests/tessera-ir/phase8/apple_gpu_fusion_descriptor.mlir +and tests/unit/test_apple_fusion_descriptor.py. +""" + +from __future__ import annotations + +from pathlib import Path + +import tessera as ts +from tessera.compiler.graph_ir import GraphIRBuilder +from tessera.compiler.canonical_compile import stamp_fusion_intents, _INTENT_KERNELS + + +def _module(fn): + # The AST frontend lowers `def` functions (not lambdas), so callers pass a + # real def. + b = GraphIRBuilder() + b.lower(fn) + return b.module() + + +def test_stamps_matmul_softmax_matmul_terminal(): + def f(a, b, c): + return ts.ops.matmul(ts.ops.softmax(ts.ops.matmul(a, b)), c) + m = _module(f) + assert stamp_fusion_intents(m) == 1 + body = m.functions[-1].body + # the tail matmul (highest index) is where the C++ pass reads the intent. + # Stamped into `attrs` (MLIR-only), never kwargs (the op's call arguments). + assert 'tessera.fusion.intent = "matmul_softmax_matmul"' in (body[-1].attrs or "") + assert "tessera.fusion.intent" not in (body[0].attrs or "") + assert "tessera.fusion.intent" not in body[-1].kwargs + + +def test_stamps_two_op_chains(): + def gelu_chain(a, b): + return ts.ops.gelu(ts.ops.matmul(a, b)) + + def rmsnorm_chain(a, b): + return ts.ops.rmsnorm(ts.ops.matmul(a, b)) + + def softmax_chain(a, b): + return ts.ops.softmax(ts.ops.matmul(a, b)) + + for kernel, fn in (("matmul_gelu", gelu_chain), + ("matmul_rmsnorm", rmsnorm_chain), + ("matmul_softmax", softmax_chain)): + m = _module(fn) + assert stamp_fusion_intents(m) == 1, kernel + assert f'tessera.fusion.intent = "{kernel}"' in (m.functions[-1].body[-1].attrs or "") + + +def test_rendered_mlir_carries_intent(): + def gelu_chain(a, b): + return ts.ops.gelu(ts.ops.matmul(a, b)) + m = _module(gelu_chain) + stamp_fusion_intents(m) + mlir = m.functions[-1].to_mlir() + assert 'tessera.fusion.intent = "matmul_gelu"' in mlir + + +def test_idempotent(): + def f(a, b, c): + return ts.ops.matmul(ts.ops.softmax(ts.ops.matmul(a, b)), c) + m = _module(f) + assert stamp_fusion_intents(m) == 1 + assert stamp_fusion_intents(m) == 1 # re-stamping doesn't double-count + + +def test_no_chain_no_stamp(): + def single(a, b): + return ts.ops.matmul(a, b) + m = _module(single) + assert stamp_fusion_intents(m) == 0 + + +def test_emitter_intents_match_cpp_consumers(): + # Loop contract: every intent the emitter produces is an intent some Apple + # C++ fusion pass reads. Guards against the two halves drifting apart. + lowering = (Path(__file__).resolve().parents[2] / "src" / "compiler" / + "codegen" / "Tessera_Apple_Backend" / "lib" / "Target" / "Apple" / + "Lowering") + consumed = set() + for cpp in lowering.glob("Matmul*FusionToAppleGPU.cpp"): + src = cpp.read_text() + if "tessera.fusion.intent" not in src: + continue + for kernel in _INTENT_KERNELS: + if f'== "{kernel}"' in src or f'(intent == "{kernel}")' in src: + consumed.add(kernel) + assert _INTENT_KERNELS <= consumed, ( + f"emitter produces intents the C++ doesn't consume: {_INTENT_KERNELS - consumed}")