SIMD-accelerate fp32 quantisation path and add q2_lee_distance

[Copilot]

Issue asked whether WASM SIMD suggestions from Google's AI improve on our kernel. Answer: the kernel was entirely scalar. The suggestions are real optimizations. Implemented them.

SIMD fp32 fast path (q2_quantise, dtype=0)

The hot path (transformers.js always sends fp32) now uses v128 throughout:

• Load/norm/normalize: v128.load, f32x4.mul/add, f32x4.splat — 4× fewer loop iterations
• Quantise: f32x4.le + v128.bitselect — branchless symbol classification, replacing 3 nested if/else per dimension. Uses native f32x4.le (not v128.not(f32x4.gt(...))) to match scalar path's NaN→D semantics.
• Gray encode + pack: v128.xor + i32x4.shr_u → i32x4.extract_lane — 4 Gray codes into one byte directly, eliminating the output pre-zeroing loop
• Horizontal reduction: i8x16.shuffle pairwise swap-and-add for L2 norm summation

Non-fp32 dtype paths (1–4) remain scalar — they're not on the hot path.

New q2_lee_distance export

SIMD Lee distance via the Gray-map isometry (DESIGN.md §2.7, Theorem 2.1: d_H(φ(u),φ(v)) = d_L(u,v)):

v128.xor          ;; XOR 16 packed bytes at once
i8x16.popcnt      ;; popcount per byte
i16x8.extadd_pairwise_i8x16_u  ;; hierarchical sum
i32x4.extadd_pairwise_i16x8_u  ;; reduction to 4 lanes
;; + shuffle horizontal sum to scalar

Scalar tail loop handles sub-16-byte remainders.

TypeScript (q2.ts)

• Q2Kernel.leeDistance() interface + WASM binding
• q2LeeDistanceDirect() re-exports leeDistancePacked from q2stats.ts (single canonical implementation, no duplication)

What we didn't need from the suggestion

i32.reinterpret_f32 / raw bit extraction for sign — unnecessary since f32x4.le against thresholds gives us the classification directly with correct boundary semantics.

Validation

WASM output verified byte-identical to TS reference across n=128..4096. All 257 existing tests pass; 6 new Lee distance tests added.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Does this math better?</issue_title>
<issue_description>I had an interesting discussion with Google's AI about our system.
There was a suggestion we do the following.

In WebAssembly (WASM), you can achieve "revercimal-style" Q2 quantization and ring arithmetic with extreme efficiency by using 128-bit SIMD (v128). This allows you to process four f32 values at once using raw bit manipulation rather than slow decimal math.
1. Fast Quantization (The Bitwise Shortcut)
Instead of converting to a string or doing multiple modulo operations, you can extract the Q2 symbols directly from the f32 bit patterns using a single comparison and a bit-shift.

* Sign Bit (High bit of Gray code): In WASM, you can extract the sign bits of four floats at once using i32x4.bitmask. This gives you a 4-bit integer where each bit represents the sign of one float.
* Magnitude Bit (Low bit of Gray code): Use f32x4.abs followed by f32x4.gt (greater than) against your threshold $\tau^*$. This returns a mask of 0xFFFFFFFF (true) or 0x00000000 (false). [1, 2, 3] 

The WASM Logic Flow:

   1. Load four f32 values into a v128 register.
   2. Compare them to a v128.const threshold to get a magnitude mask.
   3. Extract sign bits using i32x4.shr_u (shift right) to move the sign bit to the target position.
   4. Combine the results using v128.and and v128.or to pack the 2-bit Gray codes into a single integer. [2] 

2. Math in the $\mathbb{Z}_4$ Ring (Lee Distance)
For the "math" part of your project (comparing vectors), do not convert back to floats. The Q2 design is optimized for the Lee Distance, which in WASM is just a popcnt on the XORed bits. [4] 

* XOR the bits: v128.xor(vector_a, vector_b)
* Count the 1s: Use i8x16.popcnt. This gives you the distance between the two $\mathbb{Z}_4$ vectors in a single operation.
* This is roughly 10–20x faster than standard floating-point distance calculations in WASM. [4] 

3. Implementation Tip: bit_cast
In WASM, the i32.reinterpret_f32 instruction is your best friend. It allows you to look at the "bits" of a float as an integer without any performance penalty or numeric rounding. [5] 

| Goal [6] | WASM Instruction |
|---|---|
| Get raw bits | i32.reinterpret_f32 |
| Vector Abs | f32x4.abs |
| Thresholding | f32x4.gt |
| Bit-counting | i8x16.popcnt |

This approach treats the float as a "word" (like Wilderberg's revercimals) but stays in the high-speed binary domain required for WASM performance.

[1] [https://developer.arm.com](https://developer.arm.com/community/arm-community-blogs/b/architectures-and-processors-blog/posts/webassembly-bitmask-operations)
[2] [https://tty4.dev](https://tty4.dev/development/wasm-simd-operations/)
[3] [https://github.com](https://github.com/WebAssembly/simd/issues/192)
[4] [https://developer.mozilla.org](https://developer.mozilla.org/en-US/docs/WebAssembly/Reference/Numeric/popcnt)
[5] [https://github.com](https://github.com/WebAssembly/spec/blob/main/proposals/simd/SIMD.md)
[6] [https://www.youtube.com](https://www.youtube.com/watch?v=VYJZGa9m34w&t=280)

Is this better, faster etc than our current wasm .wat kernel or does it describe what we're doing already?

Don't change anything, just compare and contrast. Because that kernel is a hotpath that needs to be as optimized as possible.</issue_description>

<agent_instructions>You're not being asked to fix anything, just to compare, contrast and elaborate. However, if you do see a clear speed up that produces the same thing as the DESIGN.md then by all means please optimize the .wat file.

In fact if you see other optimizations while you're in there go ahead and just do them and then document them. Keep in mind transformers.js only ever gives us f32 encoded versions of whatever the model is giving us.</agent_instructions>

Comments on the Issue (you are @copilot in this section)

• Fixes Does this math better? #71

---

📱 Kick off Copilot coding agent tasks wherever you are with GitHub Mobile, available on iOS and Android.

[Copilot]

Pull request overview

This PR upgrades the WASM Q² kernel to use SIMD for the fp32 (dtype=0) quantisation hot path and adds a new q2_lee_distance export, along with TypeScript bindings and a TS reference implementation for Lee distance on packed Gray vectors.

Changes:

• SIMD-accelerate q2_quantise for fp32 with v128 loads, vectorized norm/reduction, branchless classification, and direct byte packing.
• Add a new WASM export q2_lee_distance (SIMD XOR + popcnt) and bind it in src/q2.ts.
• Add q2LeeDistanceDirect() TS fallback/reference + unit tests.

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 4 comments.

File	Description
test/q2.test.ts	Adds unit tests for q2LeeDistanceDirect() behavior on packed vectors.
src/q2.wat	Implements SIMD fp32 quantisation path and adds q2_lee_distance SIMD export + documentation updates.
src/q2.wasm	Updates the compiled WASM artifact to match the new WAT implementation.
src/q2.ts	Updates embedded WASM, adds Q2Kernel.leeDistance() binding, and adds q2LeeDistanceDirect() TS fallback.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This adds a second packed-bytes Lee distance implementation (q2LeeDistanceDirect) even though src/q2stats.ts already has leeDistancePacked() (also popcount(a XOR b)). To avoid divergence and duplicated micro-optimizations, consider reusing a single shared helper (or aligning semantics, e.g. both use min(length) vs n-based length).

[Copilot]

The test case name says “returns maximum distance (2) for complement pairs A↔C”, but the expectation is the total distance (16 for n=8). Consider renaming the test to reflect that the function returns the sum over dimensions (e.g., “returns 2 per dimension for …” or “returns 16 for n=8 …”) to avoid confusion.

[Copilot]

The SIMD masks implement “v ≤ threshold” as v128.not(f32x4.gt(v, threshold)). That differs from the scalar fallback’s f32.le semantics for NaNs: for NaN lanes, gt is false so not(false) becomes true, causing NaNs to be classified as A/B/C instead of the scalar path’s default D. Use f32x4.le (or explicitly mask out NaNs) to keep classification identical across paths.

Suggested change

	;; Cascaded ≤ masks (≤ is NOT >)
	(local.set $mask_a
	(v128.not (f32x4.gt (local.get $v4) (local.get $neg_tau4))))
	(local.set $mask_b
	(v128.not (f32x4.gt (local.get $v4) (local.get $zero4))))
	(local.set $mask_c
	(v128.not (f32x4.gt (local.get $v4) (local.get $tau4))))
	;; Cascaded ≤ masks
	(local.set $mask_a
	(f32x4.le (local.get $v4) (local.get $neg_tau4)))
	(local.set $mask_b
	(f32x4.le (local.get $v4) (local.get $zero4)))
	(local.set $mask_c
	(f32x4.le (local.get $v4) (local.get $tau4)))

[Copilot]

q2LeeDistanceDirect’s doc says a and b are “(n/4 bytes)”, but the implementation pads out-of-range bytes with 0 via (a[i] ?? 0) / (b[i] ?? 0) and also derives nBytes as n >> 2 (floor). Either tighten the implementation to enforce the documented preconditions (length and n % 4 === 0), or update the doc to explicitly describe the padding/flooring behavior.

[devlux76]

@copilot apply changes based on the comments in this thread

Use your best judgement on any tradeoffs, maximize speed and accuracy

[Copilot]

@copilot apply changes based on the comments in this thread

Use your best judgement on any tradeoffs, maximize speed and accur...

All 4 review comments addressed in 894dd62:

1. NaN semantics — Replaced v128.not(f32x4.gt(...)) with native f32x4.le(...). Both scalar f32.le and SIMD f32x4.le return false for NaN, so NaN lanes now consistently classify as D (sym=3) across both paths. Also saves 3 instructions (no v128.not wrapper), net −6 bytes in the compiled WASM.

2. Deduplicate Lee distance — Removed the standalone q2LeeDistanceDirect implementation; it's now a re-export of leeDistancePacked from q2stats.ts. Single canonical implementation, no divergence risk.

3. Test name — Renamed to "returns 2 per dimension (16 total for n=8) for complement pairs A↔C".

4. Doc/precondition alignment — By re-exporting leeDistancePacked, the function now uses min(a.length, b.length) semantics (already documented in q2stats.ts) instead of the n >> 2 floor approach.

All 257 tests pass, typecheck clean, WASM validated across n=128..4096.