fix(jit): ARM64 bundle and tail encodings; v_random honours its seed (Refs #989)

.github/workflows/reproduce.yml

@@ -6,7 +6,7 @@ name: Reproduce the headless profile
 # empty directory and setup-zig keeps no cache, so nothing built before is reused.
 # The record of the run -- each command with its exit code, the compile failures by
 # root, the run steps, every compiled test root and every install with its sha256 --
-# is uploaded as the artifact reproduce-headless.
+# is uploaded as the artifact reproduce-headless-<arch>.
 
 on:
   pull_request:
@@ -28,8 +28,19 @@ on:
 
 jobs:
   headless:
-    name: Headless profile from a clean clone
-    runs-on: ubuntu-latest
+    # Both architectures the spec has floors for. The ARM64 JIT kernels run only on an ARM64
+    # machine: with x86_64 alone, two encoding errors in src/jit_arm64.zig (a CMLT for a CMGT,
+    # a 32-bit LDRSB under 64-bit arithmetic) had no runner that could execute them.
+    name: Headless profile from a clean clone (${{ matrix.arch }})
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - arch: x86_64
+            runner: ubuntu-latest
+          - arch: aarch64
+            runner: ubuntu-24.04-arm
+    runs-on: ${{ matrix.runner }}
     timeout-minutes: 120
     steps:
       - uses: actions/checkout@v4
@@ -54,7 +65,8 @@ jobs:
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: reproduce-headless
+          # One artifact per architecture: upload-artifact@v4 refuses a second one by the same name.
+          name: reproduce-headless-${{ matrix.arch }}
           # The record and, beside it, the whole output of each command.
           path: |
             reproduce-headless.json

examples/vm.zig

@@ -30,6 +30,7 @@ pub fn main() !void {
     std.debug.print("  v2 = bind(v0, v1)\n", .{});
     std.debug.print("  v3 = unbind(v2, v1)\n", .{});
     std.debug.print("  f0 = cosine(v0, v3)\n\n", .{});
-    std.debug.print("Result: f0 = {d:.4} (expected ~1.0)\n", .{vm.registers.f0});
+    // unbind recovers v0 wherever v1 is non-zero, two thirds of the trits: sqrt(2/3).
+    std.debug.print("Result: f0 = {d:.4} (expected ~0.82)\n", .{vm.registers.f0});
     std.debug.print("Cycles: {}\n\n", .{vm.cycle_count});
 }

specs/reproduce/headless.t27

@@ -56,14 +56,17 @@ pub const INSTALLS : [46]str = [
 ];
 ; The test graph: the floors it may not drop below. A graph that shrank is a partial graph, and a run step
 ; that ran no test is not a passing one: one does today, the step of external/zig-hslm/src/root.zig, the only
-; compiled test root without a test declaration. A skipped test is not a passing one either: the eight that
-; skip are the four x86-64 JIT tests of src/jit.zig, compiled into src/trinity.zig and src/e2e_test.zig, and
-; only off x86_64. The floor of passed tests is the aarch64 count, the lower of the two. Not every pass tested
-; something: 23 of the 24 tests of src/jit_arm64.zig, and 4 of the 5 of src/jit_x86_64.zig, return early off
-; their architecture and are counted as passed there.
+; compiled test root without a test declaration. A skipped test is not a passing one either, and what skips
+; depends on the machine: the four x86-64 JIT tests of src/jit.zig and the differential tests of the x86-64
+; kernels skip off x86_64, the differential test of the ARM64 kernels off aarch64 -- each counted once per
+; root that compiles it. So the floors are per architecture, ARCHES[i] with TESTS_PASSED_MIN[i] and
+; TESTS_SKIPPED_MAX[i], and a record from another architecture is not judged against either. Not every pass
+; tested something: 23 of the 24 older tests of src/jit_arm64.zig, and 4 of the 5 of src/jit_x86_64.zig,
+; return early off their architecture and are counted as passed there.
 pub const TEST_STEPS_MIN : u32 = 135;
-pub const TESTS_PASSED_MIN : u32 = 2760;
-pub const TESTS_SKIPPED_MAX : u32 = 8;
+pub const ARCHES : [2]str = ["x86_64", "aarch64"];
+pub const TESTS_PASSED_MIN : [2]u32 = [2773, 2763];
+pub const TESTS_SKIPPED_MAX : [2]u32 = [3, 13];
 pub const RUN_STEPS_WITH_TESTS_MIN : u32 = 66;
 pub const EMPTY_RUN_STEPS_MAX : u32 = 1;
 ; Every root the test step compiles, read from `--verbose` with empty caches; a root that leaves this list is
@@ -121,13 +124,15 @@ test the_profile_is_headless_and_pinned {
 }
 
 test the_floors_hold_together {
-    assert TESTS_PASSED_MIN == 2760;
-    assert TESTS_SKIPPED_MAX == 8;
+    assert TESTS_PASSED_MIN[0] == 2773;
+    assert TESTS_PASSED_MIN[1] == 2763;
+    assert TESTS_SKIPPED_MAX[0] + TESTS_PASSED_MIN[0] == TESTS_SKIPPED_MAX[1] + TESTS_PASSED_MIN[1];
     assert RUN_STEPS_WITH_TESTS_MIN + EMPTY_RUN_STEPS_MAX == 67;
     assert TEST_STEPS_MIN > RUN_STEPS_WITH_TESTS_MIN;
 }
 
 invariant every_root_runs {
     assert EMPTY_RUN_STEPS_MAX <= 1;
-    assert TESTS_SKIPPED_MAX <= 8;
+    assert TESTS_SKIPPED_MAX[0] <= 3;
+    assert TESTS_SKIPPED_MAX[1] <= 13;
 }

src/e2e_test.zig

@@ -223,10 +223,10 @@ test "E2E: VM program — random → bind → cosine pipeline" {
     });
     try machine.run();
 
-    // unbind(bind(v0,v1), v1) should be similar to v0
-    // Note: JIT engine may return negative cosine (known sign issue),
-    // and at MAX_TRITS dim with ~1/3 zeros, recovery is approximate
-    try std.testing.expect(@abs(machine.registers.f0) > 0.3);
+    // unbind(bind(v0,v1), v1) is v0 wherever v1 is non-zero -- two thirds of the trits -- so
+    // cos(v3, v0) is sqrt(2/3), about 0.82. This test used to accept |f0| > 0.3, because the
+    // ARM64 JIT returned -1.0 here ("known sign issue"); the kernel is fixed and the sign counts.
+    try std.testing.expect(machine.registers.f0 > 0.7);
     try std.testing.expectEqual(@as(u64, 6), machine.cycle_count);
 }
 
@@ -504,8 +504,8 @@ fn runVerdict(allocator: std.mem.Allocator) !VerdictResult {
         if (vsa.countNonZero(&machine.registers.v2) > 0) vm_passed += 1;
     }
 
-    // Check 4: VM cosine self-similarity
-    // Note: JIT engine has known sign issue on cosine, so test absolute value
+    // Check 4: VM cosine self-similarity (it used to take the absolute value, to pass the
+    // ARM64 JIT's -1.0; that kernel is fixed)
     {
         var machine = vm.VSAVM.init(allocator);
         defer machine.deinit();
@@ -515,7 +515,7 @@ fn runVerdict(allocator: std.mem.Allocator) !VerdictResult {
             .{ .opcode = .halt },
         });
         try machine.run();
-        if (@abs(machine.registers.f0) > 0.99) vm_passed += 1;
+        if (machine.registers.f0 > 0.99) vm_passed += 1;
     }
 
     // Check 5: VM cycle count tracking

src/jit.zig

@@ -701,3 +701,59 @@ test "JitCompiler dot product correctness" {
     const result = func(&a, &b);
     try std.testing.expectEqual(expected, result);
 }
+
+// Every kernel against a scalar reference, on random trits and on dimensions with and without a
+// tail: the tests above use one fixed sixteen- or eight-trit vector each.
+test "JitCompiler kernels agree with a scalar reference, with tails and negative trits" {
+    const builtin = @import("builtin");
+    if (builtin.cpu.arch != .x86_64) return error.SkipZigTest;
+    const dims = [_]usize{ 1, 7, 15, 16, 17, 33, 1024, 1031, 32767, 59049 };
+    const n = 59049;
+    const alloc = std.testing.allocator;
+    const a = try alloc.alloc(i8, n);
+    defer alloc.free(a);
+    const b = try alloc.alloc(i8, n);
+    defer alloc.free(b);
+    const out = try alloc.alloc(i8, n);
+    defer alloc.free(out);
+    var wrong: usize = 0;
+    for (dims) |dim| {
+        var prng = std.Random.DefaultPrng.init(dim);
+        for (a[0..dim], b[0..dim]) |*x, *y| {
+            x.* = @as(i8, @intCast(prng.random().intRangeAtMost(u8, 0, 2))) - 1;
+            y.* = @as(i8, @intCast(prng.random().intRangeAtMost(u8, 0, 2))) - 1;
+        }
+        var dot: i64 = 0;
+        for (a[0..dim], b[0..dim]) |x, y| dot += @as(i64, x) * y;
+        {
+            var c = JitCompiler.init(alloc);
+            defer c.deinit();
+            try c.compileDotProduct(dim);
+            // finalize() types its pointer for the in-place kernels; the dot product returns rax.
+            const f: *const fn ([*]const i8, [*]const i8) callconv(.c) i64 = @ptrCast(try c.finalize());
+            const r = f(a.ptr, b.ptr);
+            if (r != dot) {
+                std.debug.print("dot dim={d}: got {d}, want {d}\n", .{ dim, r, dot });
+                wrong += 1;
+            }
+        }
+        for ([_]bool{ false, true }) |bundle| {
+            @memcpy(out[0..dim], a[0..dim]);
+            var c = JitCompiler.init(alloc);
+            defer c.deinit();
+            if (bundle) try c.compileBundleDirect(dim) else try c.compileBindDirect(dim);
+            (try c.finalize())(@ptrCast(out.ptr), @ptrCast(b.ptr));
+            var bad: usize = 0;
+            for (a[0..dim], b[0..dim], out[0..dim]) |x, y, z| {
+                const s: i16 = @as(i16, x) + y;
+                const want: i8 = if (bundle) (if (s > 0) 1 else if (s < 0) -1 else 0) else x * y;
+                if (z != want) bad += 1;
+            }
+            if (bad > 0) {
+                std.debug.print("{s} dim={d}: {d} of {d} trits wrong\n", .{ if (bundle) "bundle" else "bind", dim, bad, dim });
+                wrong += 1;
+            }
+        }
+    }
+    try std.testing.expectEqual(@as(usize, 0), wrong);
+}

src/jit_arm64.zig

@@ -202,8 +202,12 @@ pub const Arm64JitCompiler = struct {
 
     /// LDRSB (register) - ldrsb Wt, [Xn, Xm]
     fn ldrsbReg(self: *Self, rt: u5, rn: u5, rm: u5) !void {
-        // LDRSB (reg, 32-bit): 00 111 0 00 11 1 Rm 011 0 10 Rn Rt
-        const instr: u32 = 0x38E06800 |
+        // LDRSB (reg, 64-bit): 00 111 0 00 10 1 Rm 011 0 10 Rn Rt -- sign-extends into all
+        // of Xt. The 32-bit form (opc 11, 0x38E06800) zeroes bits 63:32, so a -1 trit read
+        // as 0x00000000FFFFFFFF by the 64-bit MUL, ADD, CMP and SCVTF of the scalar tails:
+        // the fused cosine and the bundle came out wrong for every vector with a negative
+        // trit past the last full 16. SMULL reads only the low 32 bits, which are the same.
+        const instr: u32 = 0x38A06800 |
             (@as(u32, rm) << 16) |
             (@as(u32, rn) << 5) |
             @as(u32, rt);
@@ -587,7 +591,9 @@ pub const Arm64JitCompiler = struct {
 
     /// CMGT Vd.16B, Vn.16B, #0 - Compare greater than zero
     fn cmgt_16b_zero(self: *Self, vd: u5, vn: u5) !void {
-        const instr: u32 = 0x4E20A800 |
+        // 0 Q 0 01110 size 10000 01000 10 Rn Rd. 0x4E20A800 (opcode 01010) is CMLT #0, which
+        // made the SIMD bundle return 0 for every positive sum.
+        const instr: u32 = 0x4E208800 |
             (@as(u32, vn) << 5) |
             @as(u32, vd);
         try self.emit32(instr);
@@ -2173,3 +2179,96 @@ test "ARM64 bundle SIMD non-aligned" {
     _ = func(@ptrCast(&a), @ptrCast(&b));
     // Bundle SIMD correctness to be verified via integration tests
 }
+
+// Every kernel against a scalar reference, on random trits and on dimensions with and without a
+// tail past the last full 16. The tests above use hand-picked inputs -- all ones, or one fixed
+// eight-trit vector -- and the bundle test checks no result at all; with them all passing, the
+// SIMD bundle returned 0 for every positive sum (CMLT encoded where CMGT was meant), and the
+// fused cosine was wrong for any vector with a negative trit in its tail (LDRSB into W, then
+// 64-bit arithmetic), which is how the VM's v_cosine came to say -1.0 for two identical vectors.
+test "ARM64 kernels agree with a scalar reference, with tails and negative trits" {
+    if (!is_arm64) return error.SkipZigTest;
+    const dims = [_]usize{ 1, 7, 15, 16, 17, 33, 1024, 1031, 32767, 59049 };
+    const n = 59049;
+    const alloc = std.testing.allocator;
+    const a = try alloc.alloc(i8, n);
+    defer alloc.free(a);
+    const b = try alloc.alloc(i8, n);
+    defer alloc.free(b);
+    const out = try alloc.alloc(i8, n);
+    defer alloc.free(out);
+    var wrong: usize = 0;
+    for (dims) |dim| {
+        var prng = std.Random.DefaultPrng.init(dim);
+        for (a[0..dim], b[0..dim]) |*x, *y| {
+            x.* = @as(i8, @intCast(prng.random().intRangeAtMost(u8, 0, 2))) - 1;
+            y.* = @as(i8, @intCast(prng.random().intRangeAtMost(u8, 0, 2))) - 1;
+        }
+        var dot: i64 = 0;
+        var aa: i64 = 0;
+        var bb: i64 = 0;
+        var ham: i64 = 0;
+        for (a[0..dim], b[0..dim]) |x, y| {
+            dot += @as(i64, x) * y;
+            aa += @as(i64, x) * x;
+            bb += @as(i64, y) * y;
+            if (x != y) ham += 1;
+        }
+        const cos: f64 = @as(f64, @floatFromInt(dot)) / @sqrt(@as(f64, @floatFromInt(aa)) * @as(f64, @floatFromInt(bb)));
+        const Reduce = enum { dot, dot_hybrid, hamming, cosine };
+        for ([_]Reduce{ .dot, .dot_hybrid, .hamming, .cosine }) |k| {
+            var c = Arm64JitCompiler.init(alloc);
+            defer c.deinit();
+            switch (k) {
+                .dot => try c.compileDotProduct(dim),
+                .dot_hybrid => try c.compileDotProductHybrid(dim),
+                .hamming => try c.compileHammingSIMD(dim),
+                .cosine => try c.compileFusedCosine(dim),
+            }
+            const r = (try c.finalize())(@ptrCast(a.ptr), @ptrCast(b.ptr));
+            const ok = switch (k) {
+                .dot, .dot_hybrid => r == dot,
+                .hamming => r == ham,
+                // 0/0 for a zero vector is the caller's to map (vsa_jit maps it to 0).
+                .cosine => if (aa == 0 or bb == 0) std.math.isNan(@as(f64, @bitCast(r))) else @abs(@as(f64, @bitCast(r)) - cos) < 1e-12,
+            };
+            if (!ok) {
+                std.debug.print("{s} dim={d}: got {d}, want dot={d} hamming={d} cosine={d}\n", .{ @tagName(k), dim, r, dot, ham, cos });
+                wrong += 1;
+            }
+        }
+        if (dim % 16 == 0) {
+            var c = Arm64JitCompiler.init(alloc);
+            defer c.deinit();
+            try c.compileDotProductSIMD(dim);
+            const r = (try c.finalize())(@ptrCast(a.ptr), @ptrCast(b.ptr));
+            if (r != dot) {
+                std.debug.print("dot_simd dim={d}: got {d}, want {d}\n", .{ dim, r, dot });
+                wrong += 1;
+            }
+        }
+        const Map = enum { bind_direct, bind_simd, bundle_simd };
+        for ([_]Map{ .bind_direct, .bind_simd, .bundle_simd }) |k| {
+            @memcpy(out[0..dim], a[0..dim]);
+            var c = Arm64JitCompiler.init(alloc);
+            defer c.deinit();
+            switch (k) {
+                .bind_direct => try c.compileBindDirect(dim),
+                .bind_simd => try c.compileBindSIMD(dim),
+                .bundle_simd => try c.compileBundleSIMD(dim),
+            }
+            _ = (try c.finalize())(@ptrCast(out.ptr), @ptrCast(b.ptr));
+            var bad: usize = 0;
+            for (a[0..dim], b[0..dim], out[0..dim]) |x, y, z| {
+                const s: i16 = @as(i16, x) + y;
+                const want: i8 = if (k == .bundle_simd) (if (s > 0) 1 else if (s < 0) -1 else 0) else x * y;
+                if (z != want) bad += 1;
+            }
+            if (bad > 0) {
+                std.debug.print("{s} dim={d}: {d} of {d} trits wrong\n", .{ @tagName(k), dim, bad, dim });
+                wrong += 1;
+            }
+        }
+    }
+    try std.testing.expectEqual(@as(usize, 0), wrong);
+}

src/jit_x86_64.zig

@@ -469,3 +469,57 @@ test "x86-64 JIT large dimension" {
     const result = func(@ptrCast(&a), @ptrCast(&b));
     try std.testing.expectEqual(expected, result);
 }
+
+// Every kernel against a scalar reference, on random trits and on dimensions with and without a
+// tail. The tests above use one fixed vector each; the ARM64 compiler, tested the same way,
+// carried two encoding errors that only random trits and a tail exposed.
+test "x86-64 kernels agree with a scalar reference, with tails and negative trits" {
+    if (!is_x86_64) return error.SkipZigTest;
+    const dims = [_]usize{ 1, 7, 15, 16, 17, 33, 1024, 1031, 32767, 59049 };
+    const n = 59049;
+    const alloc = std.testing.allocator;
+    const a = try alloc.alloc(i8, n);
+    defer alloc.free(a);
+    const b = try alloc.alloc(i8, n);
+    defer alloc.free(b);
+    const out = try alloc.alloc(i8, n);
+    defer alloc.free(out);
+    var wrong: usize = 0;
+    for (dims) |dim| {
+        var prng = std.Random.DefaultPrng.init(dim);
+        for (a[0..dim], b[0..dim]) |*x, *y| {
+            x.* = @as(i8, @intCast(prng.random().intRangeAtMost(u8, 0, 2))) - 1;
+            y.* = @as(i8, @intCast(prng.random().intRangeAtMost(u8, 0, 2))) - 1;
+        }
+        var dot: i64 = 0;
+        for (a[0..dim], b[0..dim]) |x, y| dot += @as(i64, x) * y;
+        {
+            var c = X86_64JitCompiler.init(alloc);
+            defer c.deinit();
+            try c.compileDotProduct(dim);
+            const r = (try c.finalize())(@ptrCast(a.ptr), @ptrCast(b.ptr));
+            if (r != dot) {
+                std.debug.print("dot dim={d}: got {d}, want {d}\n", .{ dim, r, dot });
+                wrong += 1;
+            }
+        }
+        for ([_]bool{ false, true }) |bundle| {
+            @memcpy(out[0..dim], a[0..dim]);
+            var c = X86_64JitCompiler.init(alloc);
+            defer c.deinit();
+            if (bundle) try c.compileBundleDirect(dim) else try c.compileBindDirect(dim);
+            _ = (try c.finalize())(@ptrCast(out.ptr), @ptrCast(b.ptr));
+            var bad: usize = 0;
+            for (a[0..dim], b[0..dim], out[0..dim]) |x, y, z| {
+                const s: i16 = @as(i16, x) + y;
+                const want: i8 = if (bundle) (if (s > 0) 1 else if (s < 0) -1 else 0) else x * y;
+                if (z != want) bad += 1;
+            }
+            if (bad > 0) {
+                std.debug.print("{s} dim={d}: {d} of {d} trits wrong\n", .{ if (bundle) "bundle" else "bind", dim, bad, dim });
+                wrong += 1;
+            }
+        }
+    }
+    try std.testing.expectEqual(@as(usize, 0), wrong);
+}

src/vm.zig

@@ -259,9 +259,11 @@ pub const VSAVM = struct {
 
     fn execVRandom(self: *VSAVM, inst: VSAInstruction) void {
         const dst = self.getVReg(inst.dst);
+        // The immediate is the seed. It used to be ignored -- every v_random drew seed 0 -- so
+        // a program that bound two "random" vectors bound a vector with itself, and the bind and
+        // unbind examples could not tell unbind from the identity.
         const seed: u64 = @bitCast(inst.imm);
-        _ = seed; // TODO: actually use seed
-        dst.* = tvc_vsa.randomVector(MAX_TRITS, 0);
+        dst.* = tvc_vsa.randomVector(MAX_TRITS, seed);
     }
 
     fn execVBind(self: *VSAVM, inst: VSAInstruction) void {
@@ -971,7 +973,6 @@ test "VSA VM bind/unbind" {
 
 test "VSA VM bundle similarity" {
     var vm = VSAVM.init(std.testing.allocator);
-    vm.jit_enabled = false; // Disable JIT (has bug in cosineSimilarity)
     defer vm.deinit();
 
     const program = [_]VSAInstruction{
@@ -992,7 +993,6 @@ test "VSA VM bundle similarity" {
 
 test "VSA VM permute" {
     var vm = VSAVM.init(std.testing.allocator);
-    vm.jit_enabled = false; // Disable JIT (has bug in cosineSimilarity)
     defer vm.deinit();
 
     const program = [_]VSAInstruction{