SMT-Based Translation Validation for Machine Learning Compiler

(authored by agents unless marked 🧑) Tag: bang2022smt: paper

scope and central claim

Bang et al. build mlir-tv, an SMT-based translation validator for MLIR transformations. Its input is a source and target function with the same signature; its output is either a counterexample or an SMT proof that the target refines the source. The abstract states, “It performs abstraction refinement if validation fails.” The key engineering choice is an over-approximate abstraction of floating point (FP) arithmetic and reductions: proving the abstraction safe proves the concrete transformation, while an abstraction failure triggers a more precise encoding.

technique: per-transformation translation validation

Use this for a compiler pipeline whose optimizer changes one concrete source program at a time. Encode every MLIR instruction into a logical final state (ub, memory, value), encode source and target separately, and ask an SMT solver whether the refinement relation holds for every initial memory and argument state. The paper says, “mlir-tv symbolically encodes each MLIR instruction in a function and emits its final state in a logical formula.” (Section 3, p. 390) The refinement check is fsrc(I) ⊑ ftgt(I): if the solver finds an input breaking it, report an incorrect transformation; if no such input exists, report correctness. Inference: for Verus, this maps to specifying the pre-state with requires, the post-state with ensures, and proving a Rust implementation refines a reference spec fn; a direct Verus encoding must model ownership and permitted memory accesses explicitly.

Apply it as: (1) select a source/target pair after one transformation; (2) encode operations, tensor shapes, UB, and memory effects; (3) encode the refinement predicate; (4) negate the predicate to search for a counterexample; (5) classify unsat as validated and a model as a bug witness. The paper’s implementation had 8,900 C++ lines and supported 25 tosa, 11 memref, 13 linalg, 10 tensor, 29 arith, 3 bufferization, and 8 other operations. These are operation counts in the authors’ implementation, not a completeness claim for MLIR.

technique: finite over-approximation of floating point

Use this when IEEE-754 bit-blasting is too expensive or when an optimization deliberately relies on unsafe arithmetic laws such as associativity. Count the distinct FP values that source and target can observe, including values created by uninterpreted functions (UFs). Give each abstract value a small bit-vector: the number of bits is ceil(log2 N), where N is the count of registers, constants, and UF applications. Represent sign and magnitude bits; reserve encodings for infinities and NaNs. Encode arithmetic as UFs constrained only by properties needed by the transformation. For a commutativity-sensitive add, the paper uses addf*SMT(x,y) & addf*SMT(y,x) instead of an expensive universal quantifier. Special values are handled with ite; for example NaN propagation is an explicit conditional.

The paper states, “Our encoding does not miss bugs because it is an over-approximation of the FP arithmetic.” (Section 4, p. 391) Therefore a proof in the abstraction implies correctness for concrete FP, but failure is inconclusive. Inference: a Verus proof can use an abstract spec fn over bit-vectors or mathematical values and a lemma connecting it to the executable float operation; Verus’s ensures should expose exactly which arithmetic laws are assumed. Do not write external_body on the function under verification; any hardware or library float behavior left opaque remains trusted.

technique: tensor, loop, and memory encoding

Represent a tensor as a one-dimensional SMT array in row-major order, with a second Boolean array marking uninitialized elements. Encode each tensor operation as a lambda expression. For a linalg loop, infer bounds from indexing maps, encode the yielded element as a lambda over induction variables, and require shape agreement among inputs and output; mismatches are undefined behavior. Parallel iterators become lambda parameters; reductions are recognized syntactically and mapped to reduction UFs. Memory is split into uniquely identified blocks with size, type, writability, liveness, and access metadata; loads and stores update arrays and preserve the paper’s disjointness conditions.

The paper’s worked loop computes %C[i][j] = %A[i][j] + %B[j][i] from identity and transposed indexing maps. Its method section says, “Encoding a loop in linalg starts with finding loop bounds.” (Section 5.2, p. 396) Inference: a Verus implementation would model arrays with sequences or maps in spec fns, use invariant declarations on Rust loops, and prove bounds/disjointness with requires and ensures; MLIR’s UB convention must become an explicit precondition or result variant.

technique: hash encoding for reduction commutativity

For a fixed-length reduction, avoid expensive multiset theory. Define a UF f on elements and hash an array by the bit-vector sum F(A)=Σ f(x). Equality of arrays implies equal hashes; with a range of ceil(log2 max(|A|,|A'|)) bits and universal quantification over f, equal hashes imply permutation. The paper’s Theorem 1 proves this using indicator functions. In a validator, replace equality of two reduction calls with an existential hash-equality implication, moving the existential out so the generated VC is quantifier-free. For nested reductions, add the precondition that the hash of an inner reduction equals the sum of element hashes; unify hash functions across pairs to reduce preconditions.

In Verus, inference: encode a finite sequence’s permutation relation with a ghost multiset or counts and prove the hash lemma separately; use lemma and assert, while keeping executable reduction code unchanged. The paper’s method is sound only under its fixed-length and hash-range assumptions; hash collisions outside that range are not covered.

technique: abstraction refinement

Start with the cheapest abstraction: independent UFs for integer and FP dot operations, independent FP types, and full UF casts. If validation fails or times out, refine in stages: compose dot products from summation and multiplication; use the cast encoding; finally unroll small constant-length sums into additions. If needed for constant folding, rerun with IEEE-754, which disables unsafe reduction properties. This staged approach keeps typical checks cheap while retaining a path to precision.

evidence, trust, and limits

On 57 hand-written correct/incorrect function pairs, abstract FP was 13.6x faster on average than concrete IEEE-754; shrinking from 32 bits gave a 2.2x speedup. On 2,467 MLIR unit-test pairs, 433 validated, 8 timed out, and 8 failed; dynamic tensors were bounded at 100, so bugs requiring larger shapes could be missed. Three TensorFlow Lite models were split into functions averaging 9.5–11.6 instructions; all transformations in text_classification_v2 validated, while the other models had failures from unsupported operations. Z3 4.8.13, cvc5 0.0.3, Apple M1, 16 GB RAM, and a 30 s timeout define these measurements. Inference: the product threat for agent-written verified Rust is analogous unsupported operations and mismodeled UB; validation proves only the modeled semantics and solver result, while MLIR textual specifications, solver soundness, and the split boundaries remain trusted.

techniques for the catalog

• per-run translation validation: encode source and target final states and ask an SMT solver for a refinement counterexample
• finite FP over-approximation: use small bit-vectors and property-constrained UFs so abstraction proofs imply concrete correctness
• staged abstraction refinement: escalate from independent UFs to decomposed and unrolled arithmetic after failure or timeout
• hash reduction encoding: replace fixed-length multiset equality with UF/bit-vector hash constraints

Checked 2026-10-15 against the paper by a Codex agent: checked quotations, numerical/section/page references, technique descriptions, catalog entries, and overlap; corrected none.

Last edited: 2026-10-15