Commit 2c515ff
committed
feat(track-c): z3 GEMM-rewrite prover (serial-reduction -> T.gemm)
Built-in z3 proof that a serial-reduction -> T.gemm (matmul2d) rewrite is
algebraically equivalent + race-free, NON-vacuously. RULE #1: the auto-GEMM
pass (track B) calls require_gemm_rewrite_proof() and the rewrite only applies
if z3 proves it; an unprovable contraction RAISES GemmRewriteNotProven (forced
mode) or keeps the serial path (prove_gemm_rewrite detect-and-prefer mode).
Six z3 obligations, each checking the NEGATION for unsat (mirrors
_async_barrier_plan / mamba3_path_c provers): operand-map match (no accidental
transpose), tiling injectivity over k, k-coverage/surjectivity, mask-predicate
equivalence, scale-fold equivalence (uninterpreted base fn), single-writer
output disjointness. Concrete F0 (cb=C@B^T, summary_states scale fold) + B2
(dinp lower-tri mask) contraction descriptors included.
Non-vacuity locked in by tests: a transposed operand / off-by-one k-tile /
overlapping output tiles / strict-vs-inclusive causal mask / off-by-one scale
fold each yield a SAT counter-witness and z3_proved=False. Structure-only on
purpose -- fp16/fp32 parity stays the separate numeric gate.
Files: cppmega_mlx/nn/_tilelang/_gemm_rewrite_proof.py (+ test).1 parent 283889b commit 2c515ff
2 files changed
Lines changed: 816 additions & 0 deletions
0 commit comments