Skip to content

Commit 2c515ff

Browse files
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

File tree

0 commit comments

Comments
 (0)