From 8185a767a09dd375332479c7de27b06f7d363b75 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 21 Mar 2026 01:37:00 +0000 Subject: [PATCH 1/4] Initial plan From c3e3c7513a8b6ff8ef6029531283d453c114996c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 21 Mar 2026 01:52:12 +0000 Subject: [PATCH 2/4] SIMD-accelerate fp32 quantisation path and add q2_lee_distance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add 128-bit SIMD (v128) fast path for dtype=0 (fp32) in q2_quantise: • v128.load for 4× f32 bulk loads • f32x4.mul/add for vectorised L2 norm accumulation • i8x16.shuffle for horizontal f32x4 reduction • f32x4.gt + v128.not + v128.bitselect for branchless symbol classification • v128.xor + i32x4.shr_u for vectorised Gray encoding • i32x4.extract_lane for direct byte packing (no pre-zeroing needed) - Add new exported q2_lee_distance function (DESIGN.md §2.6–§2.7): • v128.xor + i8x16.popcnt for SIMD Hamming distance on Gray codes • i16x8/i32x4.extadd_pairwise for hierarchical sum reduction • Scalar tail loop for sub-16-byte remainders - Update q2.ts: Q2Kernel.leeDistance interface + q2LeeDistanceDirect fallback - Add 6 new tests for q2LeeDistanceDirect (identical, complement, adjacent, cyclic-wrap, symmetry, empty vectors) - All 257 tests pass, typecheck clean, WASM validated across n=128..4096 Co-authored-by: devlux76 <86517969+devlux76@users.noreply.github.com> Agent-Logs-Url: https://github.com/devlux76/q2/sessions/7c4e359f-3a13-46a7-8a66-a9e387936697 --- src/q2.ts | 93 ++++++++++--- src/q2.wasm | Bin 920 -> 1664 bytes src/q2.wat | 342 +++++++++++++++++++++++++++++++++++++++++++++++- test/q2.test.ts | 44 +++++++ 4 files changed, 458 insertions(+), 21 deletions(-) diff --git a/src/q2.ts b/src/q2.ts index 5991ef2..186b47e 100644 --- a/src/q2.ts +++ b/src/q2.ts @@ -83,22 +83,34 @@ export const DTYPE_TO_Q2: Record = { * Regenerate with: bun run build:wat */ const WASM_B64 = - 'AGFzbQEAAAABHARgAX8BfWADf39/AX1gBX9/f39/AX9gAn9/AX4DBQQAAQIDBQMBAAgGBgF/AEEACwce' + - 'AwNtZW0CAAtxMl9xdWFudGlzZQACBnEyX2tleQADCrsGBFgBA38gAEGAgAJxQRB0IQEgAEEKdkEfcSEC' + - 'IABB/wdxIQMgAkUEQCABvg8LIAJBH0YEQCABQYCAgPwHIANBDXRycr4PCyABIAJB8ABqQRd0IANBDXRy' + - 'cr4LigEBAX8CQAJAAkACQAJAIAIOBAABAgMECyAAIAFBAnRqKgIADwsgACABQQF0ai8BABAADwsgACAB' + - 'aiwAALIPCyAAIAFBAXZqLQAAIQMgAUEBcQRAIANBD3EhAwUgA0EEdiEDCyADQQhrsg8LIAAgAUECdmot' + - 'AAAhAyADQQMgAUEDcWtBAXR2QQNxswu8AwMEfwR9BH8gAkECdiEFIANBBEYEQEEAIQYCQANAIAYgBU8N' + - 'ASAEIAZqIAAgBmotAAA6AAAgBkEBaiEGDAALCyAFDwtBACEGAkADQCAGIAJPDQEgAUEBayACbCAGaiEH' + - 'IwAgBkECdGohCCAIIAAgByADEAE4AgAgBkEBaiEGDAALC0MAAAAAIQpBACEGAkADQCAGIAJPDQEjACAG' + - 'QQJ0aioCACEJIAogCSAJlJIhCiAGQQFqIQYMAAsLIApDlZXmJF4EQEMAAIA/IAqRlSELQQAhBgJAA0Ag' + - 'BiACTw0BIwAgBkECdGohCCAIIAgqAgAgC5Q4AgAgBkEBaiEGDAALCwtDCKwsPyACs5GVIQxBACEGAkAD' + - 'QCAGIAVPDQEgBCAGakEAOgAAIAZBAWohBgwACwtBACEGAkADQCAGIAJPDQEjACAGQQJ0aioCACEJQQMh' + - 'DSAJIAyMXwRAQQAhDQUgCUMAAAAAXwRAQQEhDQUgCSAMXwRAQQIhDQsLCyANIA1BAXZzIQ4gBkECdiEP' + - 'QQMgBkEDcWtBAXQhECAEIA9qIAQgD2otAAAgDiAQdHI6AAAgBkEBaiEGDAALCyAFC5UBAwZ/AX4Bf0H/' + - 'ASEHQgAhCEEAIQlBACECAkADQCACIAFPDQEgAkECdiEDQQMgAkEDcWtBAXQhBCAAIANqLQAAIAR2QQNx' + - 'IQUgBUECcSAFQQF2IAVBAXFzciEGIAYgB0cEQCAGIQcgCUEgSQRAIAggBq1BPiAJQQF0a62GhCEIIAlB' + - 'AWohCQsLIAJBAWohAgwACwsgCAs='; + 'AGFzbQEAAAABIwVgAX8BfWADf39/AX1gBX9/f39/AX9gAn9/AX5gA39/fwF/AwYFAAECAwQFAwEACAYG' + + 'AX8AQQALBzAEA21lbQIAC3EyX3F1YW50aXNlAAIGcTJfa2V5AAMPcTJfbGVlX2Rpc3RhbmNlAAQKiQwF' + + 'WAEDfyAAQYCAAnFBEHQhASAAQQp2QR9xIQIgAEH/B3EhAyACRQRAIAG+DwsgAkEfRgRAIAFBgICA/Acg' + + 'A0ENdHJyvg8LIAEgAkHwAGpBF3QgA0ENdHJyvguKAQEBfwJAAkACQAJAAkAgAg4EAAECAwQLIAAgAUEC' + + 'dGoqAgAPCyAAIAFBAXRqLwEAEAAPCyAAIAFqLAAAsg8LIAAgAUEBdmotAAAhAyABQQFxBEAgA0EPcSED' + + 'BSADQQR2IQMLIANBCGuyDwsgACABQQJ2ai0AACEDIANBAyABQQNxa0EBdHZBA3GzC7cHBQR/BH0Efwt7' + + 'An8gAkECdiEFIANBBEYEQEEAIQYCQANAIAYgBU8NASAEIAZqIAAgBmotAAA6AAAgBkEBaiEGDAALCyAF' + + 'DwsgAUUEQEEADwsgA0UEQCAAIAFBAWsgAmxBAnRqIRxBACEGAkADQCAGIAJPDQEjACAGQQJ0aiAcIAZB' + + 'AnRq/QAEAP0LBAAgBkEEaiEGDAALC0MAAAAA/RMhEkEAIQYCQANAIAYgAk8NASMAIAZBAnRq/QAEACER' + + 'IBIgESAR/eYB/eQBIRIgBkEEaiEGDAALCyASIBL9DQgJCgsMDQ4PAAECAwQFBgchEyASIBP95AEhEiAS' + + 'IBL9DQQFBgcAAQIDCAkKCwwNDg8hEyASIBP95AEhEiAS/R8AIQogCkOVleYkXgRAQwAAgD8gCpGVIQsg' + + 'C/0TIRFBACEGAkADQCAGIAJPDQEjACAGQQJ0aiEIIAggCP0ABAAgEf3mAf0LBAAgBkEEaiEGDAALCwtD' + + 'CKwsPyACs5GVIQwgDP0TIRQgDIz9EyEVQwAAAAD9EyEWQQAhBgJAA0AgBiACTw0BIwAgBkECdGr9AAQA' + + 'IREgESAV/UT9TSEXIBEgFv1E/U0hGCARIBT9RP1NIRlBAv0RQQP9ESAZ/VIhGkEB/REgGiAY/VIhGkEA' + + '/REgGiAX/VIhGiAaIBpBAf2tAf1RIRsgG/0bAEEGdCAb/RsBQQR0ciAb/RsCQQJ0IBv9GwNyciEdIAQg' + + 'BkECdmogHToAACAGQQRqIQYMAAsLIAUPC0EAIQYCQANAIAYgAk8NASABQQFrIAJsIAZqIQcjACAGQQJ0' + + 'aiEIIAggACAHIAMQATgCACAGQQFqIQYMAAsLQwAAAAAhCkEAIQYCQANAIAYgAk8NASMAIAZBAnRqKgIA' + + 'IQkgCiAJIAmUkiEKIAZBAWohBgwACwsgCkOVleYkXgRAQwAAgD8gCpGVIQtBACEGAkADQCAGIAJPDQEj' + + 'ACAGQQJ0aiEIIAggCCoCACALlDgCACAGQQFqIQYMAAsLC0MIrCw/IAKzkZUhDEEAIQYCQANAIAYgBU8N' + + 'ASAEIAZqQQA6AAAgBkEBaiEGDAALC0EAIQYCQANAIAYgAk8NASMAIAZBAnRqKgIAIQlBAyENIAkgDIxf' + + 'BEBBACENBSAJQwAAAABfBEBBASENBSAJIAxfBEBBAiENCwsLIA0gDUEBdnMhDiAGQQJ2IQ9BAyAGQQNx' + + 'a0EBdCEQIAQgD2ogBCAPai0AACAOIBB0cjoAACAGQQFqIQYMAAsLIAULlQEDBn8BfgF/Qf8BIQdCACEI' + + 'QQAhCUEAIQICQANAIAIgAU8NASACQQJ2IQNBAyACQQNxa0EBdCEEIAAgA2otAAAgBHZBA3EhBSAFQQJx' + + 'IAVBAXYgBUEBcXNyIQYgBiAHRwRAIAYhByAJQSBJBEAgCCAGrUE+IAlBAXRrrYaEIQggCUEBaiEJCwsg' + + 'AkEBaiECDAALCyAIC9EBAgN/BXsgAkECdiEDQQAhBUEAIQQCQANAIARBEGogA0sNASAAIARq/QAEACAB' + + 'IARq/QAEAP1RIQYgBv1iIQcgB/19IQggCP1/IQkgCSAJ/Q0ICQoLDA0ODwABAgMEBQYHIQogCSAK/a4B' + + 'IQkgCSAJ/Q0EBQYHAAECAwgJCgsMDQ4PIQogCSAK/a4BIQkgBSAJ/RsAaiEFIARBEGohBAwACwsCQANA' + + 'IAQgA08NASAFIAAgBGotAAAgASAEai0AAHNpaiEFIARBAWohBAwACwsgBQs='; function b64ToBytes(b64: string): Uint8Array { const bin = atob(b64.replace(/\s+/g, '')); @@ -154,6 +166,19 @@ export interface Q2Kernel { * @returns 64-bit key as BigInt (DESIGN.md §2.2) */ key(_packedOffset: number, _n: number): bigint; + + /** + * Compute the Lee distance between two packed Gray-encoded Q² vectors. + * + * Uses SIMD-accelerated XOR + popcnt to exploit the Gray-map isometry + * (DESIGN.md §2.7 Theorem 2.1): d_H(φ(u), φ(v)) = d_L(u, v). + * + * @param aOffset - byte offset in WASM memory for the first packed vector + * @param bOffset - byte offset in WASM memory for the second packed vector + * @param n - original embedding dimension (n/4 packed bytes per vector) + * @returns total Lee distance (sum of per-dimension distances) + */ + leeDistance(_aOffset: number, _bOffset: number, _n: number): number; } // ─── Instantiation ──────────────────────────────────────────────────────────── @@ -180,6 +205,7 @@ async function instantiate(): Promise { mem: WebAssembly.Memory; q2_quantise: (_ip: number, _sl: number, _n: number, _dt: number, _op: number) => number; q2_key: (_sp: number, _n: number) => bigint; + q2_lee_distance: (_ap: number, _bp: number, _n: number) => number; }; const e = instance.exports as WasmExports; @@ -191,6 +217,9 @@ async function instantiate(): Promise { key(packedOffset, n) { return e.q2_key(packedOffset, n); }, + leeDistance(aOffset, bOffset, n) { + return e.q2_lee_distance(aOffset, bOffset, n); + }, }; } @@ -302,3 +331,35 @@ export function l2Normalise(data: Float32Array, n: number): Float32Array { return v; } + +/** + * Compute the Lee distance between two packed Gray-encoded Q² vectors. + * + * Pure-TypeScript reference implementation of the WASM q2_lee_distance + * function. Exploits the Gray-map isometry (DESIGN.md §2.7 Theorem 2.1): + * d_H(φ(u), φ(v)) = d_L(u, v) + * + * Hamming distance on Gray-encoded bits equals Lee distance on Z₄ originals, + * so the total distance is simply popcount(a XOR b) over all packed bytes. + * + * @param a - first packed Gray-encoded vector (n/4 bytes) + * @param b - second packed Gray-encoded vector (n/4 bytes) + * @param n - original embedding dimension + * @returns total Lee distance (sum of per-dimension Lee distances) + */ +export function q2LeeDistanceDirect( + a: Uint8Array, + b: Uint8Array, + n: number, +): number { + const nBytes = n >> 2; + let total = 0; + for (let i = 0; i < nBytes; i++) { + let x = (a[i] ?? 0) ^ (b[i] ?? 0); + // Popcount: count set bits in 8-bit value + x = x - ((x >> 1) & 0x55); + x = (x & 0x33) + ((x >> 2) & 0x33); + total += (x + (x >> 4)) & 0x0F; + } + return total; +} diff --git a/src/q2.wasm b/src/q2.wasm index 50c6188cffd0205183d109c9997c8f108772ce60..33410c92efe4b01675a54b5de546d0cbc2d4a56c 100644 GIT binary patch delta 738 zcmZuv!EO^V5S>}uyXksE*rZL`ZbQeSN01Olq^hw3rwSnu=avFFbX$-}qF$=Vy&`dm zsz`hUaOr`s;J`mX|H2aE4MC}dtkpb^pLgE8@#o~CuHk6(4gmn~s;MrvUkT?NIuR?t zNJu3RYH^GXEa>EcobGnpjSmkGPX~j;x5uOL(fc<8km<)vO%Hx&-xDRBoJgn7qsx8V zy+im!_AE3s2?hj&7EdTTwb8gVTbDnuvqM`Hhph(*UQ6JeEKQt&)_s-*zGa$$=mNpy z)TEqpe~tbZn&uLj51PJSTS;|RztxzFj}z14d~43F2vWsKI9w2K5d1cnl+ylecD-{T zdtCowkJ9PPXws5}a*UKxtG}-pKGG=4wbh XV=1Z;M1& diff --git a/src/q2.wat b/src/q2.wat index 920aad2..3c0f2e7 100644 --- a/src/q2.wat +++ b/src/q2.wat @@ -1,8 +1,8 @@ (module ;; ───────────────────────────────────────────────────────────────────────────── - ;; Q² — Quaternary Quantisation Kernel + ;; Q² — Quaternary Quantisation Kernel (SIMD-accelerated) ;; Source: src/q2.wat - ;; Specification: DESIGN.md §1.5 – §2.2 + ;; Specification: DESIGN.md §1.5 – §2.8 ;; ;; Memory layout (8 pages = 512 KB): ;; [0x00000, 0x10000) page 0 — f32 working buffer (≤ 16 384 dims) @@ -11,9 +11,26 @@ ;; [0x40000, 0x80000) pages 4-7 — host input area ($input_ptr must be ≥ 0x40000) ;; ;; Exports: - ;; mem — shared linear memory (host writes input here, reads output) - ;; q2_quantise(...) — L2-normalise (last token position) + quaternary-quantise → packed Gray bytes - ;; q2_key(...) — run-reduction → 64-bit MSB-aligned transition key + ;; mem — shared linear memory (host writes input here, reads output) + ;; q2_quantise(...) — L2-normalise (last token position) + quaternary-quantise → packed Gray bytes + ;; q2_key(...) — run-reduction → 64-bit MSB-aligned transition key + ;; q2_lee_distance(…) — SIMD Lee distance via XOR + popcnt on packed Gray vectors + ;; + ;; Performance notes (SIMD optimisation): + ;; The fp32 dtype path (dtype=0) uses 128-bit SIMD (v128) throughout: + ;; • v128.load — loads 4× f32 in a single instruction + ;; • f32x4.mul/add — accumulates L2 norm across 4 lanes simultaneously + ;; • f32x4.splat — broadcasts norm_inv for parallel normalisation + ;; • f32x4.gt — vectorised threshold comparison (3 compares per 4 dims) + ;; • v128.bitselect — branchless symbol classification (no if/else branching) + ;; • i32x4.extract_lane — packs 4 Gray codes into one byte + ;; Horizontal f32x4 reduction uses i8x16.shuffle (pairwise swap-and-add). + ;; The Lee distance function (q2_lee_distance) uses v128.xor + i8x16.popcnt + ;; to compute exact cyclic Lee distance in the Z₄ ring without decoding + ;; Gray symbols — the fastest hardware-accelerated distance primitive + ;; (DESIGN.md §2.6, §2.7 Theorem 2.1: d_H(φ(u),φ(v)) = d_L(u,v)). + ;; Non-fp32 dtype paths (fp16, q8, q4, q2) remain scalar since + ;; transformers.js always provides fp32 activations on the hot path. ;; ───────────────────────────────────────────────────────────────────────────── (memory (export "mem") 8) @@ -213,6 +230,21 @@ ;; The n/4 input bytes are copied directly from $input_ptr to $out_ptr, and ;; the function returns without performing load/L2-normalise/threshold/ ;; quantise/Gray-encode/pack steps. + ;; + ;; SIMD fast path (dtype = 0, fp32): + ;; When dtype=0 the kernel uses 128-bit SIMD (v128) to process 4 f32 + ;; dimensions per iteration. Every loop body operates on v128 registers: + ;; Load: v128.load (4 f32s from the input tensor in one op) + ;; Norm²: f32x4.mul + f32x4.add (4-wide FMA accumulation) + ;; Normalise: f32x4.mul with splatted 1/‖v‖ (4-wide broadcast multiply) + ;; Quantise: 3× f32x4.gt + v128.not + v128.bitselect (branchless + ;; symbol classification — no if/else per dimension) + ;; Gray + Pack: v128.xor + i32x4.shr_u → i32x4.extract_lane × 4 + ;; (4 Gray codes combined into one output byte) + ;; Horizontal reduction for the L2 norm uses i8x16.shuffle (pair-wise + ;; swap-and-add) to sum the 4 f32 accumulator lanes. + ;; Result: the fp32 hot path executes ~4× fewer loop iterations with + ;; wider data movement and zero branching in the inner quantisation loop. ;; ───────────────────────────────────────────────────────────────────────────── (func (export "q2_quantise") (param $input_ptr i32) @@ -234,6 +266,20 @@ (local $g i32) (local $byte_idx i32) (local $bit_shift i32) + ;; SIMD locals + (local $v4 v128) ;; current 4× f32 vector + (local $acc4 v128) ;; SIMD accumulator for norm² + (local $hi4 v128) ;; temp for horizontal reduction + (local $tau4 v128) ;; splatted threshold + (local $neg_tau4 v128) ;; splatted −threshold + (local $zero4 v128) ;; splatted 0.0 + (local $mask_a v128) ;; v ≤ −τ (A mask) + (local $mask_b v128) ;; v ≤ 0 (A|B mask) + (local $mask_c v128) ;; v ≤ τ (A|B|C mask) + (local $sym4 v128) ;; 4× i32 symbol values + (local $gray4 v128) ;; 4× i32 Gray codes + (local $src_base i32) ;; byte offset of last-token row start in input + (local $packed_byte i32) ;; assembled output byte (local.set $n_bytes (i32.shr_u (local.get $n) (i32.const 2))) @@ -264,6 +310,180 @@ ) ) + ;; ══════════════════════════════════════════════════════════════════════════ + ;; dtype = 0 (fp32): SIMD fast path — process 4 f32 dimensions per iteration + ;; ══════════════════════════════════════════════════════════════════════════ + (if (i32.eqz (local.get $dtype)) + (then + ;; Byte offset of last-token row start: input_ptr + (seq_len-1)*n*4 + (local.set $src_base + (i32.add (local.get $input_ptr) + (i32.shl + (i32.mul + (i32.sub (local.get $seq_len) (i32.const 1)) + (local.get $n)) + (i32.const 2)))) + + ;; ── SIMD Step 1: Load last token into working buffer (4 f32s/iter) ── + (local.set $d (i32.const 0)) + (block $sload_done + (loop $sload_loop + (br_if $sload_done (i32.ge_u (local.get $d) (local.get $n))) + (v128.store + (i32.add (global.get $ACCUM_BASE) (i32.shl (local.get $d) (i32.const 2))) + (v128.load + (i32.add (local.get $src_base) (i32.shl (local.get $d) (i32.const 2))))) + (local.set $d (i32.add (local.get $d) (i32.const 4))) + (br $sload_loop) + ) + ) + + ;; ── SIMD Step 2: Compute squared L2 norm (4-wide accumulation) ────── + (local.set $acc4 (f32x4.splat (f32.const 0.0))) + (local.set $d (i32.const 0)) + (block $snorm_done + (loop $snorm_loop + (br_if $snorm_done (i32.ge_u (local.get $d) (local.get $n))) + (local.set $v4 + (v128.load + (i32.add (global.get $ACCUM_BASE) (i32.shl (local.get $d) (i32.const 2))))) + (local.set $acc4 + (f32x4.add (local.get $acc4) + (f32x4.mul (local.get $v4) (local.get $v4)))) + (local.set $d (i32.add (local.get $d) (i32.const 4))) + (br $snorm_loop) + ) + ) + + ;; Horizontal sum: acc4 = [a, b, c, d] → a+b+c+d in lane 0 + ;; Step A: swap lanes [2,3] ↔ [0,1] and add + (local.set $hi4 + (i8x16.shuffle 8 9 10 11 12 13 14 15 0 1 2 3 4 5 6 7 + (local.get $acc4) (local.get $acc4))) + (local.set $acc4 (f32x4.add (local.get $acc4) (local.get $hi4))) + ;; Step B: swap lane 1 ↔ lane 0 and add + (local.set $hi4 + (i8x16.shuffle 4 5 6 7 0 1 2 3 8 9 10 11 12 13 14 15 + (local.get $acc4) (local.get $acc4))) + (local.set $acc4 (f32x4.add (local.get $acc4) (local.get $hi4))) + ;; norm_sq is now in lane 0 + (local.set $norm_sq (f32x4.extract_lane 0 (local.get $acc4))) + + ;; ── SIMD Step 3: L2-normalise (skip if ‖v‖ ≈ 0) ──────────────────── + (if (f32.gt (local.get $norm_sq) (f32.const 1e-16)) + (then + (local.set $norm_inv + (f32.div (f32.const 1.0) (f32.sqrt (local.get $norm_sq)))) + (local.set $v4 (f32x4.splat (local.get $norm_inv))) + (local.set $d (i32.const 0)) + (block $snrm_done + (loop $snrm_loop + (br_if $snrm_done (i32.ge_u (local.get $d) (local.get $n))) + (local.set $acc_ptr + (i32.add (global.get $ACCUM_BASE) (i32.shl (local.get $d) (i32.const 2)))) + (v128.store (local.get $acc_ptr) + (f32x4.mul (v128.load (local.get $acc_ptr)) (local.get $v4))) + (local.set $d (i32.add (local.get $d) (i32.const 4))) + (br $snrm_loop) + ) + ) + ) + ) + + ;; ── SIMD Step 4: Compute threshold τ* = 0.6745 / √n ──────────────── + (local.set $tau + (f32.div (f32.const 0.6745) + (f32.sqrt (f32.convert_i32_u (local.get $n))))) + (local.set $tau4 (f32x4.splat (local.get $tau))) + (local.set $neg_tau4 (f32x4.splat (f32.neg (local.get $tau)))) + (local.set $zero4 (f32x4.splat (f32.const 0.0))) + + ;; ── SIMD Step 5+6: Quantise → Gray-encode → pack (4 dims → 1 byte) ─ + ;; For every group of 4 f32 values, produce one packed output byte. + ;; Uses branchless SIMD: 3 vector compares + bitselect (no if/else). + ;; + ;; Classification via cascaded ≤ comparisons (≤ = NOT >): + ;; mask_a = v ≤ −τ → where A (sym = 0) + ;; mask_b = v ≤ 0 → where A or B (sym ≤ 1) + ;; mask_c = v ≤ τ → where A, B, or C (sym ≤ 2) + ;; + ;; Symbol selection (start with 3=D, overlay lower values): + ;; sym = bitselect(2, 3, mask_c) → 2 where v ≤ τ, else 3 + ;; sym = bitselect(1, sym, mask_b) → 1 where v ≤ 0 + ;; sym = bitselect(0, sym, mask_a) → 0 where v ≤ −τ + ;; + ;; Gray-encode: g = sym ⊕ (sym >> 1) (DESIGN.md §2.7, φ(n) = n ⊕ ⌊n/2⌋) + ;; + ;; Pack: extract 4 lanes, shift into MSB-first positions, OR together. + (local.set $d (i32.const 0)) + (block $sq_done + (loop $sq_loop + (br_if $sq_done (i32.ge_u (local.get $d) (local.get $n))) + + ;; Load 4 normalised f32 values + (local.set $v4 + (v128.load + (i32.add (global.get $ACCUM_BASE) (i32.shl (local.get $d) (i32.const 2))))) + + ;; 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)))) + + ;; Branchless symbol selection: D=3 → C=2 → B=1 → A=0 + (local.set $sym4 + (v128.bitselect + (i32x4.splat (i32.const 2)) + (i32x4.splat (i32.const 3)) + (local.get $mask_c))) + (local.set $sym4 + (v128.bitselect + (i32x4.splat (i32.const 1)) + (local.get $sym4) + (local.get $mask_b))) + (local.set $sym4 + (v128.bitselect + (i32x4.splat (i32.const 0)) + (local.get $sym4) + (local.get $mask_a))) + + ;; Gray-encode: g = sym ⊕ (sym >> 1) + (local.set $gray4 + (v128.xor (local.get $sym4) + (i32x4.shr_u (local.get $sym4) (i32.const 1)))) + + ;; Pack 4 Gray codes into one byte (MSB-first): + ;; byte = g[0]<<6 | g[1]<<4 | g[2]<<2 | g[3] + (local.set $packed_byte + (i32.or + (i32.or + (i32.shl (i32x4.extract_lane 0 (local.get $gray4)) (i32.const 6)) + (i32.shl (i32x4.extract_lane 1 (local.get $gray4)) (i32.const 4))) + (i32.or + (i32.shl (i32x4.extract_lane 2 (local.get $gray4)) (i32.const 2)) + (i32x4.extract_lane 3 (local.get $gray4))))) + + ;; Write directly — no pre-zeroing needed + (i32.store8 + (i32.add (local.get $out_ptr) (i32.shr_u (local.get $d) (i32.const 2))) + (local.get $packed_byte)) + + (local.set $d (i32.add (local.get $d) (i32.const 4))) + (br $sq_loop) + ) + ) + + (return (local.get $n_bytes)) + ) + ) + + ;; ══════════════════════════════════════════════════════════════════════════ + ;; Non-fp32 dtypes (1-3): scalar fallback path + ;; ══════════════════════════════════════════════════════════════════════════ + ;; ── Step 1: Load last token position into working buffer ──────────────── ;; element index of last token, dimension d: (seq_len − 1) × n + d (local.set $d (i32.const 0)) @@ -487,4 +707,116 @@ (local.get $key) ) + + ;; ───────────────────────────────────────────────────────────────────────────── + ;; q2_lee_distance — SIMD Lee distance between two packed Gray-encoded Q² vectors. + ;; + ;; Computes the total Lee distance between two Q² vectors, exploiting the + ;; isometry established by Theorem 2.1 (Hammons et al., 1994): + ;; + ;; d_H(φ(u), φ(v)) = d_L(u, v) for all u, v ∈ Z₄ⁿ + ;; + ;; where φ is the Gray map (DESIGN.md §2.7) and d_L is the Lee metric on Z₄. + ;; Because each Q² symbol is Gray-encoded to 2 bits, the Hamming distance + ;; between the encoded bit-vectors equals the Lee distance on the original + ;; Z₄ vectors — and Hamming distance is simply popcnt(XOR). + ;; + ;; This function uses 128-bit SIMD to XOR 16 packed bytes at a time, count + ;; set bits per byte with i8x16.popcnt, and horizontally sum via + ;; i16x8.extadd_pairwise_i8x16_u + i32x4.extadd_pairwise_i16x8_u for a + ;; fully vectorised distance computation (DESIGN.md §2.6). + ;; + ;; Parameters: + ;; $a_ptr pointer to first packed Gray-encoded vector (n/4 bytes) + ;; $b_ptr pointer to second packed Gray-encoded vector (n/4 bytes) + ;; $n original embedding dimension (n/4 = number of packed bytes) + ;; + ;; Returns: + ;; i32 total Lee distance (sum of per-dimension Lee distances) + ;; ───────────────────────────────────────────────────────────────────────────── + (func (export "q2_lee_distance") + (param $a_ptr i32) + (param $b_ptr i32) + (param $n i32) + (result i32) + + (local $n_bytes i32) + (local $d i32) + (local $total i32) + (local $xored v128) + (local $popcnt v128) + (local $pairs v128) + (local $quads v128) + (local $hi v128) + + (local.set $n_bytes (i32.shr_u (local.get $n) (i32.const 2))) + (local.set $total (i32.const 0)) + (local.set $d (i32.const 0)) + + ;; ── SIMD loop: process 16 packed bytes (64 Q² symbols) per iteration ──── + (block $simd_done + (loop $simd_loop + ;; Need at least 16 bytes remaining for a SIMD iteration + (br_if $simd_done + (i32.gt_u + (i32.add (local.get $d) (i32.const 16)) + (local.get $n_bytes))) + + ;; XOR corresponding packed bytes: differing bits = Hamming distance bits + (local.set $xored + (v128.xor + (v128.load (i32.add (local.get $a_ptr) (local.get $d))) + (v128.load (i32.add (local.get $b_ptr) (local.get $d))))) + + ;; Count set bits per byte: each byte's popcount = Lee distance for + ;; the 4 Q² symbols packed in that byte (Theorem 2.1) + (local.set $popcnt (i8x16.popcnt (local.get $xored))) + + ;; Horizontal sum of all 16 byte popcounts → single i32 total + ;; Step 1: pairwise add adjacent u8 → 8× u16 + (local.set $pairs + (i16x8.extadd_pairwise_i8x16_u (local.get $popcnt))) + ;; Step 2: pairwise add adjacent u16 → 4× u32 + (local.set $quads + (i32x4.extadd_pairwise_i16x8_u (local.get $pairs))) + ;; Step 3: horizontal sum of 4 i32 lanes + ;; Swap lanes [2,3] ↔ [0,1] and add + (local.set $hi + (i8x16.shuffle 8 9 10 11 12 13 14 15 0 1 2 3 4 5 6 7 + (local.get $quads) (local.get $quads))) + (local.set $quads (i32x4.add (local.get $quads) (local.get $hi))) + ;; Swap lane 1 ↔ lane 0 and add + (local.set $hi + (i8x16.shuffle 4 5 6 7 0 1 2 3 8 9 10 11 12 13 14 15 + (local.get $quads) (local.get $quads))) + (local.set $quads (i32x4.add (local.get $quads) (local.get $hi))) + + (local.set $total + (i32.add (local.get $total) (i32x4.extract_lane 0 (local.get $quads)))) + + (local.set $d (i32.add (local.get $d) (i32.const 16))) + (br $simd_loop) + ) + ) + + ;; ── Scalar tail: remaining bytes (< 16) ───────────────────────────────── + (block $tail_done + (loop $tail_loop + (br_if $tail_done (i32.ge_u (local.get $d) (local.get $n_bytes))) + + ;; XOR one byte, count bits with i32.popcnt + (local.set $total + (i32.add (local.get $total) + (i32.popcnt + (i32.xor + (i32.load8_u (i32.add (local.get $a_ptr) (local.get $d))) + (i32.load8_u (i32.add (local.get $b_ptr) (local.get $d))))))) + + (local.set $d (i32.add (local.get $d) (i32.const 1))) + (br $tail_loop) + ) + ) + + (local.get $total) + ) ) diff --git a/test/q2.test.ts b/test/q2.test.ts index 642c727..340a384 100644 --- a/test/q2.test.ts +++ b/test/q2.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'; import { q2EncodeDirect, q2KeyDirect, + q2LeeDistanceDirect, l2Normalise, Q2_DTYPE_FP32, Q2_DTYPE_FP16, @@ -172,6 +173,49 @@ describe('l2Normalise', () => { }); }); +describe('q2LeeDistanceDirect', () => { + it('returns 0 for identical vectors', () => { + // All D: Gray 10₂ → 0xAA per byte + const a = new Uint8Array([0xAA, 0xAA]); + expect(q2LeeDistanceDirect(a, a, 8)).toBe(0); + }); + + it('returns maximum distance (2) for complement pairs A↔C', () => { + // A=00₂, C=11₂ → complement (DESIGN.md §2.8, distance 2 per dim) + // All A → 0x00; All C → 0xFF + const a = new Uint8Array([0x00, 0x00]); // 8× A + const c = new Uint8Array([0xFF, 0xFF]); // 8× C (11₂ packed) + // Each of 8 dims differs by 2 bits → total Hamming = 16 = total Lee + expect(q2LeeDistanceDirect(a, c, 8)).toBe(16); + }); + + it('returns 1 per dimension for adjacent symbols A↔B', () => { + // A=00₂, B=01₂ → adjacent, Lee distance 1 per dim + // All A → 0x00; All B → 01_01_01_01₂ = 0x55 + const a = new Uint8Array([0x00, 0x00]); // 8× A + const b = new Uint8Array([0x55, 0x55]); // 8× B + expect(q2LeeDistanceDirect(a, b, 8)).toBe(8); + }); + + it('returns 1 per dimension for cyclic-adjacent D↔A (strong extremes)', () => { + // D=10₂, A=00₂ → cyclic adjacent, Lee distance 1 per dim + // XOR = 10₂ per dim → 1 bit per dim → Hamming 1 = Lee 1 + const d = new Uint8Array([0xAA, 0xAA]); // 8× D (10₂) + const a = new Uint8Array([0x00, 0x00]); // 8× A (00₂) + expect(q2LeeDistanceDirect(d, a, 8)).toBe(8); + }); + + it('is symmetric: distance(a,b) = distance(b,a)', () => { + const a = new Uint8Array([0xAA, 0x55]); // mixed D and B + const b = new Uint8Array([0xFF, 0x00]); // mixed C and A + expect(q2LeeDistanceDirect(a, b, 8)).toBe(q2LeeDistanceDirect(b, a, 8)); + }); + + it('works with n=0 (empty vectors)', () => { + expect(q2LeeDistanceDirect(new Uint8Array(0), new Uint8Array(0), 0)).toBe(0); + }); +}); + describe('Q2Dtype constants and DTYPE_TO_Q2 map', () => { it('maps dtype strings to correct ids', () => { expect(DTYPE_TO_Q2['fp32']).toBe(Q2_DTYPE_FP32); From 49b308740e168a8b10b7910d7f9b2f116e480274 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 21 Mar 2026 01:53:30 +0000 Subject: [PATCH 3/4] Address code review: improve test comments for clarity Co-authored-by: devlux76 <86517969+devlux76@users.noreply.github.com> Agent-Logs-Url: https://github.com/devlux76/q2/sessions/7c4e359f-3a13-46a7-8a66-a9e387936697 --- test/q2.test.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/q2.test.ts b/test/q2.test.ts index 340a384..557cd5f 100644 --- a/test/q2.test.ts +++ b/test/q2.test.ts @@ -182,9 +182,9 @@ describe('q2LeeDistanceDirect', () => { it('returns maximum distance (2) for complement pairs A↔C', () => { // A=00₂, C=11₂ → complement (DESIGN.md §2.8, distance 2 per dim) - // All A → 0x00; All C → 0xFF + // All A → 0x00; All C → 0xFF = 11_11_11_11₂ (four C symbols per byte) const a = new Uint8Array([0x00, 0x00]); // 8× A - const c = new Uint8Array([0xFF, 0xFF]); // 8× C (11₂ packed) + const c = new Uint8Array([0xFF, 0xFF]); // 8× C // Each of 8 dims differs by 2 bits → total Hamming = 16 = total Lee expect(q2LeeDistanceDirect(a, c, 8)).toBe(16); }); @@ -197,7 +197,7 @@ describe('q2LeeDistanceDirect', () => { expect(q2LeeDistanceDirect(a, b, 8)).toBe(8); }); - it('returns 1 per dimension for cyclic-adjacent D↔A (strong extremes)', () => { + it('returns 1 per dimension for cyclic-adjacent D↔A', () => { // D=10₂, A=00₂ → cyclic adjacent, Lee distance 1 per dim // XOR = 10₂ per dim → 1 bit per dim → Hamming 1 = Lee 1 const d = new Uint8Array([0xAA, 0xAA]); // 8× D (10₂) From 894dd62424a019043b054c370c41ce44c9b6bc82 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 21 Mar 2026 02:20:04 +0000 Subject: [PATCH 4/4] Address PR review: f32x4.le for NaN consistency, deduplicate Lee distance, fix test name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace v128.not(f32x4.gt(...)) with f32x4.le(...) in SIMD path to match scalar f32.le NaN→D semantics (NaN comparisons return false with both le and gt, so le gives correct mask directly) - Remove duplicate q2LeeDistanceDirect implementation; re-export leeDistancePacked from q2stats.ts to avoid divergence - Rename complement test to clarify it returns sum over dimensions - Update test calls to match new (a, b) signature (no n parameter) Co-authored-by: devlux76 <86517969+devlux76@users.noreply.github.com> Agent-Logs-Url: https://github.com/devlux76/q2/sessions/6a9ae6ac-eb19-4555-a97e-4760c53b5299 --- src/q2.ts | 57 ++++++++++++++++++------------------------------ src/q2.wasm | Bin 1664 -> 1658 bytes src/q2.wat | 9 ++++---- test/q2.test.ts | 16 +++++++------- 4 files changed, 34 insertions(+), 48 deletions(-) diff --git a/src/q2.ts b/src/q2.ts index 186b47e..22e111b 100644 --- a/src/q2.ts +++ b/src/q2.ts @@ -84,11 +84,11 @@ export const DTYPE_TO_Q2: Record = { */ const WASM_B64 = 'AGFzbQEAAAABIwVgAX8BfWADf39/AX1gBX9/f39/AX9gAn9/AX5gA39/fwF/AwYFAAECAwQFAwEACAYG' + - 'AX8AQQALBzAEA21lbQIAC3EyX3F1YW50aXNlAAIGcTJfa2V5AAMPcTJfbGVlX2Rpc3RhbmNlAAQKiQwF' + + 'AX8AQQALBzAEA21lbQIAC3EyX3F1YW50aXNlAAIGcTJfa2V5AAMPcTJfbGVlX2Rpc3RhbmNlAAQKgwwF' + 'WAEDfyAAQYCAAnFBEHQhASAAQQp2QR9xIQIgAEH/B3EhAyACRQRAIAG+DwsgAkEfRgRAIAFBgICA/Acg' + 'A0ENdHJyvg8LIAEgAkHwAGpBF3QgA0ENdHJyvguKAQEBfwJAAkACQAJAAkAgAg4EAAECAwQLIAAgAUEC' + 'dGoqAgAPCyAAIAFBAXRqLwEAEAAPCyAAIAFqLAAAsg8LIAAgAUEBdmotAAAhAyABQQFxBEAgA0EPcSED' + - 'BSADQQR2IQMLIANBCGuyDwsgACABQQJ2ai0AACEDIANBAyABQQNxa0EBdHZBA3GzC7cHBQR/BH0Efwt7' + + 'BSADQQR2IQMLIANBCGuyDwsgACABQQJ2ai0AACEDIANBAyABQQNxa0EBdHZBA3GzC7EHBQR/BH0Efwt7' + 'An8gAkECdiEFIANBBEYEQEEAIQYCQANAIAYgBU8NASAEIAZqIAAgBmotAAA6AAAgBkEBaiEGDAALCyAF' + 'DwsgAUUEQEEADwsgA0UEQCAAIAFBAWsgAmxBAnRqIRxBACEGAkADQCAGIAJPDQEjACAGQQJ0aiAcIAZB' + 'AnRq/QAEAP0LBAAgBkEEaiEGDAALC0MAAAAA/RMhEkEAIQYCQANAIAYgAk8NASMAIAZBAnRq/QAEACER' + @@ -96,21 +96,21 @@ const WASM_B64 = 'IBL9DQQFBgcAAQIDCAkKCwwNDg8hEyASIBP95AEhEiAS/R8AIQogCkOVleYkXgRAQwAAgD8gCpGVIQsg' + 'C/0TIRFBACEGAkADQCAGIAJPDQEjACAGQQJ0aiEIIAggCP0ABAAgEf3mAf0LBAAgBkEEaiEGDAALCwtD' + 'CKwsPyACs5GVIQwgDP0TIRQgDIz9EyEVQwAAAAD9EyEWQQAhBgJAA0AgBiACTw0BIwAgBkECdGr9AAQA' + - 'IREgESAV/UT9TSEXIBEgFv1E/U0hGCARIBT9RP1NIRlBAv0RQQP9ESAZ/VIhGkEB/REgGiAY/VIhGkEA' + - '/REgGiAX/VIhGiAaIBpBAf2tAf1RIRsgG/0bAEEGdCAb/RsBQQR0ciAb/RsCQQJ0IBv9GwNyciEdIAQg' + - 'BkECdmogHToAACAGQQRqIQYMAAsLIAUPC0EAIQYCQANAIAYgAk8NASABQQFrIAJsIAZqIQcjACAGQQJ0' + - 'aiEIIAggACAHIAMQATgCACAGQQFqIQYMAAsLQwAAAAAhCkEAIQYCQANAIAYgAk8NASMAIAZBAnRqKgIA' + - 'IQkgCiAJIAmUkiEKIAZBAWohBgwACwsgCkOVleYkXgRAQwAAgD8gCpGVIQtBACEGAkADQCAGIAJPDQEj' + - 'ACAGQQJ0aiEIIAggCCoCACALlDgCACAGQQFqIQYMAAsLC0MIrCw/IAKzkZUhDEEAIQYCQANAIAYgBU8N' + - 'ASAEIAZqQQA6AAAgBkEBaiEGDAALC0EAIQYCQANAIAYgAk8NASMAIAZBAnRqKgIAIQlBAyENIAkgDIxf' + - 'BEBBACENBSAJQwAAAABfBEBBASENBSAJIAxfBEBBAiENCwsLIA0gDUEBdnMhDiAGQQJ2IQ9BAyAGQQNx' + - 'a0EBdCEQIAQgD2ogBCAPai0AACAOIBB0cjoAACAGQQFqIQYMAAsLIAULlQEDBn8BfgF/Qf8BIQdCACEI' + - 'QQAhCUEAIQICQANAIAIgAU8NASACQQJ2IQNBAyACQQNxa0EBdCEEIAAgA2otAAAgBHZBA3EhBSAFQQJx' + - 'IAVBAXYgBUEBcXNyIQYgBiAHRwRAIAYhByAJQSBJBEAgCCAGrUE+IAlBAXRrrYaEIQggCUEBaiEJCwsg' + - 'AkEBaiECDAALCyAIC9EBAgN/BXsgAkECdiEDQQAhBUEAIQQCQANAIARBEGogA0sNASAAIARq/QAEACAB' + - 'IARq/QAEAP1RIQYgBv1iIQcgB/19IQggCP1/IQkgCSAJ/Q0ICQoLDA0ODwABAgMEBQYHIQogCSAK/a4B' + - 'IQkgCSAJ/Q0EBQYHAAECAwgJCgsMDQ4PIQogCSAK/a4BIQkgBSAJ/RsAaiEFIARBEGohBAwACwsCQANA' + - 'IAQgA08NASAFIAAgBGotAAAgASAEai0AAHNpaiEFIARBAWohBAwACwsgBQs='; + 'IREgESAV/UUhFyARIBb9RSEYIBEgFP1FIRlBAv0RQQP9ESAZ/VIhGkEB/REgGiAY/VIhGkEA/REgGiAX' + + '/VIhGiAaIBpBAf2tAf1RIRsgG/0bAEEGdCAb/RsBQQR0ciAb/RsCQQJ0IBv9GwNyciEdIAQgBkECdmog' + + 'HToAACAGQQRqIQYMAAsLIAUPC0EAIQYCQANAIAYgAk8NASABQQFrIAJsIAZqIQcjACAGQQJ0aiEIIAgg' + + 'ACAHIAMQATgCACAGQQFqIQYMAAsLQwAAAAAhCkEAIQYCQANAIAYgAk8NASMAIAZBAnRqKgIAIQkgCiAJ' + + 'IAmUkiEKIAZBAWohBgwACwsgCkOVleYkXgRAQwAAgD8gCpGVIQtBACEGAkADQCAGIAJPDQEjACAGQQJ0' + + 'aiEIIAggCCoCACALlDgCACAGQQFqIQYMAAsLC0MIrCw/IAKzkZUhDEEAIQYCQANAIAYgBU8NASAEIAZq' + + 'QQA6AAAgBkEBaiEGDAALC0EAIQYCQANAIAYgAk8NASMAIAZBAnRqKgIAIQlBAyENIAkgDIxfBEBBACEN' + + 'BSAJQwAAAABfBEBBASENBSAJIAxfBEBBAiENCwsLIA0gDUEBdnMhDiAGQQJ2IQ9BAyAGQQNxa0EBdCEQ' + + 'IAQgD2ogBCAPai0AACAOIBB0cjoAACAGQQFqIQYMAAsLIAULlQEDBn8BfgF/Qf8BIQdCACEIQQAhCUEA' + + 'IQICQANAIAIgAU8NASACQQJ2IQNBAyACQQNxa0EBdCEEIAAgA2otAAAgBHZBA3EhBSAFQQJxIAVBAXYg' + + 'BUEBcXNyIQYgBiAHRwRAIAYhByAJQSBJBEAgCCAGrUE+IAlBAXRrrYaEIQggCUEBaiEJCwsgAkEBaiEC' + + 'DAALCyAIC9EBAgN/BXsgAkECdiEDQQAhBUEAIQQCQANAIARBEGogA0sNASAAIARq/QAEACABIARq/QAE' + + 'AP1RIQYgBv1iIQcgB/19IQggCP1/IQkgCSAJ/Q0ICQoLDA0ODwABAgMEBQYHIQogCSAK/a4BIQkgCSAJ' + + '/Q0EBQYHAAECAwgJCgsMDQ4PIQogCSAK/a4BIQkgBSAJ/RsAaiEFIARBEGohBAwACwsCQANAIAQgA08N' + + 'ASAFIAAgBGotAAAgASAEai0AAHNpaiEFIARBAWohBAwACwsgBQs='; function b64ToBytes(b64: string): Uint8Array { const bin = atob(b64.replace(/\s+/g, '')); @@ -335,31 +335,16 @@ export function l2Normalise(data: Float32Array, n: number): Float32Array { /** * Compute the Lee distance between two packed Gray-encoded Q² vectors. * - * Pure-TypeScript reference implementation of the WASM q2_lee_distance - * function. Exploits the Gray-map isometry (DESIGN.md §2.7 Theorem 2.1): + * This is a convenience re-export of {@link leeDistancePacked} from q2stats.ts. + * Both exploit the Gray-map isometry (DESIGN.md §2.7 Theorem 2.1): * d_H(φ(u), φ(v)) = d_L(u, v) * * Hamming distance on Gray-encoded bits equals Lee distance on Z₄ originals, * so the total distance is simply popcount(a XOR b) over all packed bytes. + * Distance is computed over min(a.length, b.length) bytes. * * @param a - first packed Gray-encoded vector (n/4 bytes) * @param b - second packed Gray-encoded vector (n/4 bytes) - * @param n - original embedding dimension * @returns total Lee distance (sum of per-dimension Lee distances) */ -export function q2LeeDistanceDirect( - a: Uint8Array, - b: Uint8Array, - n: number, -): number { - const nBytes = n >> 2; - let total = 0; - for (let i = 0; i < nBytes; i++) { - let x = (a[i] ?? 0) ^ (b[i] ?? 0); - // Popcount: count set bits in 8-bit value - x = x - ((x >> 1) & 0x55); - x = (x & 0x33) + ((x >> 2) & 0x33); - total += (x + (x >> 4)) & 0x0F; - } - return total; -} +export { leeDistancePacked as q2LeeDistanceDirect } from './q2stats.js'; diff --git a/src/q2.wasm b/src/q2.wasm index 33410c92efe4b01675a54b5de546d0cbc2d4a56c..de9433f76ffa5d74f43eb9fc3859686028446b0d 100644 GIT binary patch delta 41 wcmZqR{lzn(l(BhY*>lE?o8uU_G77pXiYo{zi2Zd{lmOBqe_c0gF^4e&05AOv)c^nh delta 47 xcmeyx)4)5Sl(BPS*>lG2o8uU_GK#zW^;HyC5L6KR3t~wCSt20TW)) + ;; Cascaded ≤ masks (native f32x4.le preserves NaN→D semantics + ;; matching the scalar path's f32.le, which returns false for NaN) (local.set $mask_a - (v128.not (f32x4.gt (local.get $v4) (local.get $neg_tau4)))) + (f32x4.le (local.get $v4) (local.get $neg_tau4))) (local.set $mask_b - (v128.not (f32x4.gt (local.get $v4) (local.get $zero4)))) + (f32x4.le (local.get $v4) (local.get $zero4))) (local.set $mask_c - (v128.not (f32x4.gt (local.get $v4) (local.get $tau4)))) + (f32x4.le (local.get $v4) (local.get $tau4))) ;; Branchless symbol selection: D=3 → C=2 → B=1 → A=0 (local.set $sym4 diff --git a/test/q2.test.ts b/test/q2.test.ts index 557cd5f..3f12d11 100644 --- a/test/q2.test.ts +++ b/test/q2.test.ts @@ -177,16 +177,16 @@ describe('q2LeeDistanceDirect', () => { it('returns 0 for identical vectors', () => { // All D: Gray 10₂ → 0xAA per byte const a = new Uint8Array([0xAA, 0xAA]); - expect(q2LeeDistanceDirect(a, a, 8)).toBe(0); + expect(q2LeeDistanceDirect(a, a)).toBe(0); }); - it('returns maximum distance (2) for complement pairs A↔C', () => { + it('returns 2 per dimension (16 total for n=8) for complement pairs A↔C', () => { // A=00₂, C=11₂ → complement (DESIGN.md §2.8, distance 2 per dim) // All A → 0x00; All C → 0xFF = 11_11_11_11₂ (four C symbols per byte) const a = new Uint8Array([0x00, 0x00]); // 8× A const c = new Uint8Array([0xFF, 0xFF]); // 8× C // Each of 8 dims differs by 2 bits → total Hamming = 16 = total Lee - expect(q2LeeDistanceDirect(a, c, 8)).toBe(16); + expect(q2LeeDistanceDirect(a, c)).toBe(16); }); it('returns 1 per dimension for adjacent symbols A↔B', () => { @@ -194,7 +194,7 @@ describe('q2LeeDistanceDirect', () => { // All A → 0x00; All B → 01_01_01_01₂ = 0x55 const a = new Uint8Array([0x00, 0x00]); // 8× A const b = new Uint8Array([0x55, 0x55]); // 8× B - expect(q2LeeDistanceDirect(a, b, 8)).toBe(8); + expect(q2LeeDistanceDirect(a, b)).toBe(8); }); it('returns 1 per dimension for cyclic-adjacent D↔A', () => { @@ -202,17 +202,17 @@ describe('q2LeeDistanceDirect', () => { // XOR = 10₂ per dim → 1 bit per dim → Hamming 1 = Lee 1 const d = new Uint8Array([0xAA, 0xAA]); // 8× D (10₂) const a = new Uint8Array([0x00, 0x00]); // 8× A (00₂) - expect(q2LeeDistanceDirect(d, a, 8)).toBe(8); + expect(q2LeeDistanceDirect(d, a)).toBe(8); }); it('is symmetric: distance(a,b) = distance(b,a)', () => { const a = new Uint8Array([0xAA, 0x55]); // mixed D and B const b = new Uint8Array([0xFF, 0x00]); // mixed C and A - expect(q2LeeDistanceDirect(a, b, 8)).toBe(q2LeeDistanceDirect(b, a, 8)); + expect(q2LeeDistanceDirect(a, b)).toBe(q2LeeDistanceDirect(b, a)); }); - it('works with n=0 (empty vectors)', () => { - expect(q2LeeDistanceDirect(new Uint8Array(0), new Uint8Array(0), 0)).toBe(0); + it('works with empty vectors', () => { + expect(q2LeeDistanceDirect(new Uint8Array(0), new Uint8Array(0))).toBe(0); }); });