Claude/phase 8 4 4 fp16 bf16 - #10
Merged
Merged
Conversation
Extends the apple_gpu matmul runtime path with fp16 and bf16 dtype
variants. Mirrors the Phase 8.2 BNNS bf16 follow-up on the CPU side:
- fp16: native MPSDataTypeFloat16 (Apple Silicon GPUs run fp16 at
higher throughput than fp32 on most ops).
- bf16: fp32-conversion path inside the runtime shim because MPS does
NOT natively support bf16 matrix descriptors as of macOS 14.
Same pattern as the CPU bf16 cblas_sgemm fallback.
Scope is intentionally narrow — only the matmul kernel gets dtype
variants this phase. The other custom MSL kernels (rope, softmax, gelu,
flash_attn, matmul_softmax_fusion) remain f32-only; their dtype variants
are 8.4.4.x followups.
MLIR / runtime
- Two new C symbols in apple_gpu_runtime.mm:
* tessera_apple_gpu_mps_matmul_f16 — native MPSDataTypeFloat16. ABI
is uint16_t* for fp16 bit-pattern transmission (no _Float16 dep).
* tessera_apple_gpu_mps_matmul_bf16 — fp32 conversion path. Decodes
bf16 bit-pattern via shift, runs MPSDataTypeFloat32 matmul, encodes
back with round-to-nearest-even.
- apple_gpu_runtime_stub.cpp gets matching reference fallbacks
(fp32-via-conversion) for non-Darwin builds.
- MatmulToAppleGPU.cpp picks the runtime symbol by input element type:
f32 / f16 / bf16. Same i64×3 + i32×3 ABI shape across all three —
the element type is encoded in the symbol name only.
Python
- driver.py: _apple_gpu_matmul_dtype_suffix extracts the dtype from
the Graph IR operand types (tensor<*xf16>, tensor<*xbf16>) and routes
the backend artifact's runtime symbol selection accordingly.
- schedule_ir.py: _base_attrs now surfaces dtype on every schedule op
by parsing the IROp's operand_types. The attr propagates through
Schedule -> Tile -> Target IR layers so target_ir's mps_matmul
emission carries the right dtype attr.
- runtime.py: _apple_gpu_dispatch_matmul detects input array dtype at
launch time (call-site dtypes are runtime-only since the @jit
function signatures are type-polymorphic) and routes to the matching
ctypes wrapper. fp16 and bf16 paths both use uint16_t* ABI via
numpy's .view(np.uint16); ml_dtypes.bfloat16 is byte-compatible.
- New ctypes wrappers _apple_gpu_mps_matmul_f16 / _bf16. Loader gate
now requires both new symbols (forces rebuild after Phase 8.4.4).
Tests
- New lit fixture apple_gpu_matmul_dtypes.mlir — three positive cases
(f32, f16, bf16 matmul lower to the right runtime symbol with the
shared i64×3 + i32×3 ABI) and one negative case (mixed-dtype operands
fall back to the artifact path).
- 4 new unit tests in test_apple_backend_roadmap.py:
* f32 default artifact contract (compile-time symbol selection)
* fp16 end-to-end matches fp32-converted reference at fp16 tolerance
* bf16 end-to-end matches fp32-converted reference at bf16 tolerance
(gated on ml_dtypes presence, mirrors the CPU bf16 soft-dep)
* fp16 + bf16 ABI shim correctness (direct ctypes invocation against
a freshly-compiled shim)
Verified on Apple Silicon (LLVM/MLIR 21, Metal active):
1994 unit tests passing (1991 + 3 net new fp16/bf16 tests);
12/12 Phase 8 lit fixtures passing against the in-tree tessera-opt.
fp16 matmul matches fp32-converted reference at rtol=5e-2 (MPS does
fp16 internal accumulation; minor drift from the per-element
reference is expected). bf16 matches at rtol=2e-2.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Extends each of the simple custom MSL kernels — rope, softmax, gelu —
with fp16 + bf16 dtype variants. Mirrors the Phase 8.4.4 matmul pattern:
- fp16: native MSL `half` kernel with `float` internal compute for
accuracy. Apple Silicon GPUs run `half` at higher throughput
than `float` for the elementwise math involved.
- bf16: fp32-conversion path inside the runtime shim. MSL has no
stable `bfloat` type pre-Metal-3.1, and the cost of decode +
re-encode is negligible relative to the actual GPU compute.
Six new C-ABI symbols (3 kernels × 2 dtypes), all with `uint16_t*`
boundary types. `numpy.view(np.uint16)` and `ml_dtypes.bfloat16` are
both byte-compatible.
MLIR / runtime
- Three pairs of new C symbols in apple_gpu_runtime.mm:
* tessera_apple_gpu_rope_f16/bf16
* tessera_apple_gpu_softmax_f16/bf16
* tessera_apple_gpu_gelu_f16/bf16
- New native MSL kernels rope_f16, softmax_f16, gelu_f16 use `half` I/O
with `float` internal compute (cos/sin, exp/sum, tanh).
- bf16 paths convert each operand to fp32 at the boundary, run the
existing f32 MSL kernel, encode back with round-to-nearest-even.
- apple_gpu_runtime_stub.cpp gets matching reference fallbacks (all
fp32-conversion) for non-Darwin builds.
- RopeToAppleGPU / SoftmaxToAppleGPU / GeluToAppleGPU passes pick the
runtime symbol by input element type. Same i64 + i32 ABI shape across
all three dtypes per kernel; the element type is encoded in the
symbol name only.
Python
- target_ir.py: three new fp16 MSL source constants
(_APPLE_GPU_{ROPE,SOFTMAX,GELU}_MSL_SOURCE_F16) + sha256 cache_keys.
bf16 reuses the f32 source (the runtime does the conversion); the
IR-level marker just flips entry_point + cache_key + dtype attr.
New helper _apple_gpu_kernel_msl_for_dtype maps (kernel, dtype)
pairs to (msl_source, entry_point, cache_key, dtype_attr) tuples
so the rope/softmax/gelu emission blocks share the dtype dispatch
logic.
- runtime.py: each of _apple_gpu_dispatch_{rope,softmax,gelu} now
detects input array dtype at launch time and routes to the matching
ctypes wrapper. fp16 + bf16 use uint16_t* ABI via numpy.view.
Six new wrappers _apple_gpu_{rope,softmax,gelu}_{f16,bf16}.
Loader gate now requires all six new symbols (forces rebuild after
Phase 8.4.4.1).
Tests
- New lit fixture apple_gpu_msl_dtypes.mlir — verifies dtype-aware
symbol selection for softmax + gelu with f32/f16/bf16 input tensors.
rope is omitted because tessera.rope is not a registered dialect op
(it's covered by Python tests instead).
- Seven new unit tests in test_apple_backend_roadmap.py:
* rope/softmax/gelu fp16 end-to-end (3 tests, native MSL path)
* rope/softmax/gelu bf16 end-to-end (3 tests, fp32-conversion path,
gated on ml_dtypes presence)
* runtime shim ABI exposure for all 6 new symbols (1 test)
Test bug fix in 8.4.4.1 contract tests
- The bf16 input fixtures had `(rng.randn(...).astype(bf16)) * 0.5`
— multiplying a bf16 array by a Python float promotes back to fp32
because numpy's bf16 (via ml_dtypes) doesn't special-case Python
scalar mixing the way native fp16 does. Fixed by applying the
multiplication in fp32 BEFORE the .astype(bf16) cast.
Verified on Apple Silicon (LLVM/MLIR 21, Metal active):
2001 unit tests passing (1994 + 7 net new fp16/bf16 tests);
13/13 Phase 8 lit fixtures passing against the in-tree tessera-opt.
fp16 paths match fp32-converted reference at rtol=5e-3; bf16 at
rtol=2e-2 across rope/softmax/gelu.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Phase 8.4.4.1 — fp16 / bf16 for simple MSL kernels (rope, softmax, gelu)
4 tasks
gstoner
added a commit
that referenced
this pull request
Jun 2, 2026
A + B + C
A. f16 / bf16 conv2d encode lanes ✅
New runtime symbols tessera_apple_gpu_conv2d_dev_f16_enc + _bf16_enc (reuse mpsg_encode_conv2d_dev helper with MPSDataTypeFloat16/MPSDataTypeBFloat16)
Python conv2d_enc_f16 + conv2d_enc_bf16 (full bias surface) + conv2d_enc_no_bias_f16 + conv2d_enc_no_bias_bf16 (registry-friendly)
Internal _conv2d_enc_dispatch parameterizes by dtype — clean shared codepath
Registered in ENCODE_OP_REGISTRY as ("conv2d", "f16") + ("conv2d", "bf16")
Manifest entry's dtypes lifted to full _APPLE_GPU_FUSED matrix
6 new tests in test_apple_gpu_conv2d_f16_bf16_encode.py: symbol resolution, registry completeness, f16 numerical-equivalence vs legacy host path, bf16 numerical-equivalence vs f32 reference (skips honestly when bf16 unsupported), single-cb chain composition. Updated test_bf16_registry_covers_full_op_envelope to assert the symmetric 9×3 matrix; updated the "not eligible" test to "is eligible"
B. MPSGraph cache eviction (LRU) ✅
Replaced the unbounded NSMutableDictionary cache with an LRU keyed alongside an NSMutableOrderedSet for MRU tracking — O(1) operations, single mutex
TESSERA_MPSGRAPH_CACHE_CAPACITY env knob (default 1024; 0 = unbounded for backcompat)
New C ABI: tessera_apple_gpu_mpsgraph_cache_evictions() + _cache_capacity() for introspection
Existing _cache_size() keeps working
8 new tests in test_apple_gpu_mpsgraph_cache_lru.py: symbol availability, default capacity = 1024, env-var override to N (subprocess-isolated since capacity is call_once), env-var = 0 means unbounded, eviction triggers when capacity exceeded, hot entries survive under pressure (LRU correctness), no spurious evictions in normal default-capacity runs
C. Lit fixture coverage for conv2d encode lane ✅
Added tessera.conv2d (+ tessera.conv3d for symmetry) to kRuntimeOps in TileToApple.cpp
Added a tile.mock for tessera.conv2d in apple_gpu_lowering.mlir + CHECK assertion that it lowers to tessera_apple.gpu.metal_kernel with status = "metal_runtime"
Extended _runtime_envelope() in the existing drift test to also walk _APPLE_GPU_CONV_OPS — closes glass-jaw #10 from the audit (conv2d was previously runtime-executable on the Python side but artifact_only on the C++ side; the drift gate now catches this)
Honest comment about the remaining gap families (_APPLE_GPU_PROJECTION_OPS, _APPLE_GPU_REDUCTION_OPS, _APPLE_GPU_LINALG_OPS) left for a follow-up sprint
gstoner
added a commit
that referenced
this pull request
Jun 3, 2026
…nd as compiler IR.
L3 + L4 complete — cholesky now traverses all four IR layers end-to-end as compiler IR.
L3 — Schedule→Tile
Added a TileCholesky pattern to TilingPass: since cholesky isn't a parallel loop nest, it lowers 1:1 to an opaque tile.cholesky Tile-IR op (carrying lower + a source = "tessera.cholesky" tag), making the Tile layer explicit and distinct.
Fixture apple_cholesky_tile_ir.mlir proves both the full spine (distribution-lowering → tiling) and bare tiling produce tile.cholesky.
L4 — Tile→Target (both Apple targets)
TileToApple now lowers tile.cholesky to registered tessera_apple ops carrying the runtime C ABI symbol (what L6 will execute):
CPU → tessera_apple.cpu.vector_op {abi="lapack_spotrf", op_kind="cholesky", symbol="tessera_apple_cpu_cholesky_f32"}
GPU → tessera_apple.gpu.metal_kernel {kernel="cholesky_contract", status="metal_runtime", symbol="tessera_apple_gpu_cholesky_f32"}
Added tessera.cholesky to isAppleGpuRuntimeOp and extended the drift gate to walk it — closing APPLE_AUDIT glass-jaw #10 for cholesky (the linalg family was previously excluded from the C++↔Python envelope agreement; tri_solve is the one remaining LINALG member, same template).
Hit two real constraints and handled them honestly: the tessera_apple dialect rejects unregistered ops (so CPU reuses the registered vector_op with identity attrs, like rope does), and the GPU pass crashes on mesh-wrapped input — confirming L4's correct input is bare Tile IR (the full-spine chaining is L7's job).
Verification
Drift test (now walking cholesky): 2/2 pass.
New L4 fixture apple_cholesky_target_ir.mlir: CPU + GPU pass.
Full sweep: 45 passed / 7 XFAIL / 0 unexpected.
The compiler path is now real for cholesky: tessera.cholesky → schedule.mesh.region{tile.cholesky} → tessera_apple.{cpu,gpu} with the runtime symbol named in the IR. That symbol-in-IR is exactly what L6's seam-closure executor will consume.
gstoner
added a commit
that referenced
this pull request
Jun 3, 2026
L7 complete — the full-spine alias + the linalg pilot is done
What landed (L7)
Two single-alias pipelines registered in tessera-opt.cpp (they span Transforms + Apple-backend passes that no single backend library links):
tessera-lower-to-apple_cpu-full and tessera-lower-to-apple_gpu-full = effect-annotation → distribution-lowering → tiling → tile-to-apple_{cpu,gpu}.
One invocation now drives tessera.cholesky (Graph IR) all the way to tessera_apple.{cpu,gpu} Target IR with the runtime symbol.
Crash-safe TileToApple: the artifact pass op->erase()'d originals blindly and segfaulted on real SSA Tile IR (the SSA-vs-artifact seam that keeps full_pipeline.mlir XFAIL). Added safeEraseLowered() — rebinds used results to a same-typed operand before erase, leaves the op in place when it can't (never crashes). A genuine robustness fix beyond the pilot.
Registered in Python pipeline_registry.py (alphabetized, correctly placed) + a dedicated lit fixture apple_cholesky_full_spine.mlir, and added tessera-opt.cpp to the registry drift-scan so the Python↔C++ pipeline gate covers the new aliases.
Verification
-full lit fixture (CPU + GPU): pass.
Pipeline-registry guards: 33/33 (alphabetization + C++-registration drift now both green).
Drift gate: 5/5 ok. L-series + apple + opt-build suites: 119 passed. Lit phase2/3/8: 46 / 7-XFAIL / 0-unexpected.
The pilot, end to end (L1–L7)
The complete executable spine for cholesky on Apple, vetted at every layer:
tessera.cholesky (L1 Graph IR op + verifier)
→ schedule.mesh.region (L2 Graph→Schedule; fixed a real dominance bug, un-XFAIL'd matmul)
→ tile.cholesky (L3 Schedule→Tile)
→ tessera_apple.{cpu,gpu} (L4 Tile→Target, symbol-in-IR; closed glass-jaw #10 for cholesky)
→ LAPACK spotrf / MSL (L5 CPU runtime; exact vs numpy)
→ executed via IR-named symbol (L6 seam-closure; the number comes from the IR)
→ one alias does all of it (L7 -full pipeline)
This is exactly the foundation you asked for: front end and every IR pass to backend, vetted end-to-end on real Apple CPU+GPU hardware, seam-closed so a regression in any pass surfaces as a numeric failure. And it's a template — tri_solve and svd now follow the identical 7-step rails (they already have GPU runtime symbols; tri_solve's CPU LAPACK is strsm, svd's is gesvd).
Two real bugs surfaced and fixed by forcing one op through the whole stack (the dominance violation in DistributionLoweringPass, the SSA segfault in TileToApple) — the payoff of the pilot-first strategy.
gstoner
added a commit
that referenced
this pull request
Jun 10, 2026
…rf ratchets + auto_batch polish (#61) * Apple GPU: descriptor-driven dispatch, feature-table selection, perf ratchets Closes the three remaining Apple-compiler "Still Open" themes (P1 + 2×P2). P1 — descriptor-driven dispatch (single source of truth): - New compiler/apple_gpu_envelope.py holds the 21 lane sets + opcode dicts (166 runtime ops). driver.py and runtime.py import them; their literal duplicate tables are deleted. - runtime per-op dispatcher is now a lane->handler table built from APPLE_GPU_LANE_BY_OP (the ~200-line elif chain is gone); AppleKernelDescriptor gains a `lane` field and classifies from the envelope module. - C++ TileToApple kRuntimeOps is generated from the registry (scripts/generate_apple_runtime_ops_table.py -> apple_runtime_ops.inc), closing glass-jaw #10 (projection/reduction now tag metal_runtime). The drift gate covers the full 166-op envelope. - Oracles: test_apple_gpu_envelope_dispatch.py (lane vs legacy elif routing), test_apple_runtime_ops_table_in_sync.py (.inc), and the existing tessera-opt status drift test (now all 166 ops). P2 — feature-table-driven selection: bf16 native-vs-upcast (apple_supports_native_bf16, live MTLGPUFamily probe), fused-chain/flash-attn head_dim caps (apple_fused_chain_score_cap derives 256 from the 1 KiB fp32 stack budget; all 12 runtime cap checks consult it), threads-per-row (apple_threadgroup_threads_per_row = simdgroup_size). +7 tests in test_apple_feature_limit_lowering.py. P2 — perf ratchets: hot-path manifest rows (matmul/softmax/rmsnorm/flash_attn/ bmm/conv2d) + all 7 packaged rows carry benchmark_json; recorded live baseline (benchmarks/baselines/apple_gpu_hot_paths.json via record_hot_path_baseline.py); perf_gate.py gains --ratchet (regression + coverage). Locked by test_apple_gpu_perf_ratchet.py. Also de-flaked test_apple_gpu_predicate_logical_opcodes (randomized hash() could seed a negative RNG value). Audit docs (APPLE_AUDIT, MASTER_AUDIT) updated per Decision #26; generated dashboards regenerated. Full unit sweep 8306 passed; phase8 lit 57/57; mypy clean on touched modules. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Apple GPU auto_batch: auto-detect decode chains + skip unused Graph-IR emission Closes the last open Apple-compiler item (APPLE_AUDIT P3 "auto_batch polish"). auto_batch now defaults to None (auto-detect): - _recognized_decode_chain AST-scans the body and returns True only for a pure chain of >=2 encode-eligible apple_gpu ops and nothing else. A node whitelist rejects arithmetic on op results (silu(x)*2), subscripts, comparisons, control flow, tuple returns, and any non-encode call — so detection never silently changes the semantics of a non-decode function. Explicit auto_batch=True/False override detection. - When the route is on (apple_gpu, not emit_package), the AST Graph IR the tracer never reads is no longer emitted: an _AutoBatchSkipEmission sentinel at the top of the Step 6 try lands the deferred state (empty module, no plan/ bundle) so __call__ falls through to the auto_batch wrapper. Crucially this does NOT set _trace_deferred (that would route to the surgical tracer and bypass auto_batch). Diagnostics: new JIT_APPLE_GPU_AUTO_BATCH code (JitDiagnosticCode enum + diagnostic_codes registry). Encode-op-name set is drift-gated against apple_gpu_chain.ENCODE_OP_REGISTRY. Tests: test_apple_gpu_jit_auto_batch_autodetect.py (19 — detection truth-table, emission-skip introspection, registry drift, misuse guards). Existing auto_batch/interception suites (34) unchanged; updated a stale comment in test_apple_gpu_jit_auto_batch_canonical.py. Audit docs updated (APPLE_AUDIT Still Open now empty; MASTER_AUDIT); generated dashboards regenerated. Full unit sweep green (the one differential-generator failure is a pre-existing _jit_boundary bug owned by a separate session). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Update _jit_boundary.py --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 10, 2026
CODE_AUDIT_2026_06_10 finding #10: SwigluFusionPass and MLAFusionPass built tessera.swiglu_fused / tessera.mla_decode_fused while dropping the numeric_policy attribute (storage/accum coupling, Decision #15a) carried by the constituent matmuls / flash_attn. - TesseraOps.td: add OptionalAttr<Tessera_NumericPolicyAttr> to Tessera_SwigluFusedOp and Tessera_MLADecodeFusedOp. - SwigluFusionPass: propagate when the three matmuls agree; DECLINE to fuse on conflict (one fused op cannot express per-stage policies). - MLAFusionPass: carry the flash_attn's policy (the attention step dominates the fused kernel's numerics; compress/expand GEMMs inherit). - Lit: propagation + conflict-no-fuse cases in swiglu_fusion.mlir; propagation case in mla_decode_fusion.mlir. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 10, 2026
Update CODE_AUDIT_2026_06_10 findings #5/#8/#10 + P2 to reflect the landed follow-ons: numeric_policy propagation (C++), SwiGLU fusion-group derivation + executor consumption, strict-dispatch CI lane wiring, and the matmul/unary dtype-table refactor (with the consciously-skipped symbol-getter memoization noted). Sync the matching COMPILER_AUDIT line. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.