![Synthesizing and formally verifying a SWAR bit-hack for INT4 dot products using Z3 and Lean 4 [P]](https://external-preview.redd.it/yzu-Be9zNgnjsjw34ElFlU69EcyXk5je11gADNo48BQ.png?width=1080&crop=smart&auto=webp&s=66dec8615eeb6fad9680699b89269e71a8c0960a)
Synthesizing and formally verifying a SWAR bit-hack for INT4 dot products using Z3 and Lean 4 [P]
INT4 quantization is ubiquitous in ML right now, but evaluating dot products on hardware without native SIMD/vector instructions (like WebAssembly or older ARM chips) usually requires slow sequential loops. A classic workaround is SWAR (SIMD Within A Register), but deriving the bitwise operations by hand to unpack, multiply, and sum eight 4-bit integers packed into a single 32-bit register is tedious and error-prone.
Instead of writing the bit-hack manually, I wrote a pipeline that uses an SMT solver to discover the exact bitwise formula from scratch, and then uses a theorem prover to mathematically guarantee its correctness.
Here is a breakdown of the technical process:
1. Synthesis via CEGIS Loop (Z3) I set up a Counter-Example Guided Inductive Synthesis (CEGIS) loop in Python using the Z3 SMT solver. The solver is given a ground-truth specification (naive loop: extract nibbles, sign-extend, multiply, sum) and a bounded set of allowed instructions (AND, OR, XOR, ADD, SUB, MUL, shifts). Z3 searches the space of possible instruction sequences. If it finds a candidate, we test it against random inputs. If it fails, the failing input is added to Z3's constraints, and it tries again. Eventually, it converges on a pure, branchless sequence of operations.
2. The Generated Math The resulting algorithm utilizes a known multiplier trick for byte-reversals, but Z3 managed to perfectly interleave the even/odd nibble extraction. For example, part of the generated code handles even/odd multiplications by exploiting 32-bit hardware multiplications: (ea_low * eb_low_rev) >>> 16 This evaluates two 4-bit multiplications at opposite ends of the register simultaneously without cross-talk.
3. Formal Proof in Lean 4 Passing a million random tests is nice, but it’s not a mathematical guarantee for bit-hacks. To prove there are absolutely no edge cases or overflow bugs among the 2^(64) possible input combinations (two 32-bit registers), I ported the synthesized SWAR function to Lean 4. By leveraging Lean's bv_decide (BitVec SAT solver) and omega for modular arithmetic, Lean compiles the equivalence check into a boolean satisfiability problem. The proof successfully verifies that swar_dot_product a b = ground_truth_dot_product a b for all possible inputs.
If anyone is interested in how the Z3 synthesis script is structured or wants to see the Lean 4 proof, I've put the source code:
https://github.com/Peloxerat/int4-swar-dotprod
I'd be curious to hear if there are ways to constrain Z3 to find an even shorter instruction path.