From 4d673e0914f89533b5e4b292e4ce220cd9f2f68b Mon Sep 17 00:00:00 2001 From: gitRasheed Date: Fri, 31 Jul 2026 22:03:20 +0100 Subject: [PATCH 1/4] arm64: NEON-vectorized union2by2 128-bit SIMD kernel for array-container unions, dispatched when both inputs have at least 384 elements; smaller inputs keep the scalar assembly. Based on CRoaring's union_vector16, with a transposed bitonic merge network (shorter dependency chain than the SSE rotate network, which measured slower than the scalar code on Neoverse cores), a two-multiply movemask replacement, and a skip over the merge network when the incoming block cannot interleave with the carried values. Differential tests cover the caller contracts: iorArray's aliased buffer, lazyorArray's zero-length buffer, block-boundary duplicates, and store overrun guards. --- setutil_arm64.go | 112 ++++++++++++++- setutil_arm64.s | 2 +- setutil_neon_arm64.s | 146 +++++++++++++++++++ setutil_neon_arm64_test.go | 231 +++++++++++++++++++++++++++++++ setutil_neon_bench_arm64_test.go | 73 ++++++++++ 5 files changed, 562 insertions(+), 2 deletions(-) create mode 100644 setutil_neon_arm64.s create mode 100644 setutil_neon_arm64_test.go create mode 100644 setutil_neon_bench_arm64_test.go diff --git a/setutil_arm64.go b/setutil_arm64.go index 3e089650..bf52c34c 100644 --- a/setutil_arm64.go +++ b/setutil_arm64.go @@ -4,4 +4,114 @@ package roaring //go:noescape -func union2by2(set1 []uint16, set2 []uint16, buffer []uint16) (size int) +func union2by2scalar(set1 []uint16, set2 []uint16, buffer []uint16) (size int) + +//go:noescape +func unionKernelNEON(set1, set2, buffer []uint16, shuf *byte, leftover *[16]uint16) (outLen, pos1, pos2, leftoverLen int) + +// uniqshuf[m] is the TBL index vector that compacts the lanes not set in +// mask m to the front. A variable initializer, not init(): package-level +// initializers in other files run first and would read a zero table. +var uniqshuf = buildUniqshuf() + +func buildUniqshuf() (t [256 * 16]byte) { + for m := 0; m < 256; m++ { + pos := 0 + for lane := 0; lane < 8; lane++ { + if m&(1<= carry's max the merged result is + // just carry then fresh, so skip the network. Boundary equality falls + // to the dedup chain, hence >= rather than >. + VMOV V0.H[7], R21 + CMP R21, R20 + BHS disjoint + MERGE + STOREUNIQ(R2, R17) + VORR V2.B16, V2.B16, V1.B16 + B loop + +disjoint: + VORR V2.B16, V2.B16, V22.B16 // stash fresh + VORR V0.B16, V0.B16, V2.B16 // emit old carry as this round's minimum + STOREUNIQ(R2, R17) + VORR V2.B16, V2.B16, V1.B16 // laststore = old carry + VORR V22.B16, V22.B16, V0.B16 // carry = fresh + B loop + +done: + // flush carry through the dedup into the leftover buffer + VORR V0.B16, V0.B16, V2.B16 + STOREUNIQ(R7, R17) + MOVD R17, leftoverLen+112(FP) + + // outLen = (out - outBase) / 2 + SUB R5, R2, R2 + LSR $1, R2, R2 + MOVD R2, outLen+88(FP) + + // pos1/pos2 = blocks consumed + MOVD set1_base+0(FP), R8 + SUB R8, R0, R0 + LSR $4, R0, R0 + MOVD R0, pos1+96(FP) + MOVD set2_base+24(FP), R8 + SUB R8, R1, R1 + LSR $4, R1, R1 + MOVD R1, pos2+104(FP) + RET diff --git a/setutil_neon_arm64_test.go b/setutil_neon_arm64_test.go new file mode 100644 index 00000000..eaa7e893 --- /dev/null +++ b/setutil_neon_arm64_test.go @@ -0,0 +1,231 @@ +//go:build arm64 && !gccgo && !appengine +// +build arm64,!gccgo,!appengine + +package roaring + +import ( + "math/rand" + "reflect" + "testing" +) + +func refUnion(a, b []uint16) []uint16 { + out := make([]uint16, len(a)+len(b)) + n := scalarMergeUnion(a, b, out) + return out[:n] +} + +func genSortedUnique(r *rand.Rand, n, valRange int) []uint16 { + if n > valRange { + n = valRange + } + seen := make(map[uint16]bool, n) + for len(seen) < n { + seen[uint16(r.Intn(valRange))] = true + } + out := make([]uint16, 0, n) + for v := 0; v < valRange; v++ { + if seen[uint16(v)] { + out = append(out, uint16(v)) + } + } + return out +} + +func checkUnion(t *testing.T, a, b []uint16, label string) { + t.Helper() + want := refUnion(a, b) + buffer := make([]uint16, len(a)+len(b)) + got := buffer[:union2by2(a, b, buffer)] + if !reflect.DeepEqual(want, got) { + t.Fatalf("%s: len(a)=%d len(b)=%d: want %d elems, got %d\nwant %v\ngot %v", + label, len(a), len(b), len(want), len(got), want, got) + } +} + +func TestUnion2By2NEONAdversarial(t *testing.T) { + for n := 8; n <= 2048; n *= 2 { + identical := make([]uint16, n) + evens := make([]uint16, n) + odds := make([]uint16, n) + low := make([]uint16, n) + high := make([]uint16, n) + for i := 0; i < n; i++ { + identical[i] = uint16(3 * i) + evens[i] = uint16(2 * i) + odds[i] = uint16(2*i + 1) + low[i] = uint16(i) + high[i] = uint16(65535 - n + 1 + i) + } + checkUnion(t, identical, identical, "identical") + checkUnion(t, evens, odds, "interleaved") + checkUnion(t, odds, evens, "interleaved-swap") + checkUnion(t, low, high, "disjoint-extremes") + checkUnion(t, high, low, "disjoint-extremes-swap") + } +} + +func TestUnion2By2NEONRandom(t *testing.T) { + r := rand.New(rand.NewSource(12345)) + for iter := 0; iter < 2000; iter++ { + valRange := 256 + r.Intn(65280) + a := genSortedUnique(r, r.Intn(5000), valRange) + b := genSortedUnique(r, r.Intn(5000), valRange) + checkUnion(t, a, b, "random") + } +} + +func checkKernel(t *testing.T, a, b []uint16, label string) { + t.Helper() + if len(a) < 8 || len(b) < 8 { + checkUnion(t, a, b, label) + return + } + want := refUnion(a, b) + buffer := make([]uint16, 0, len(a)+len(b)) + n := unionNEON(a, b, buffer) + got := buffer[:n] + if !reflect.DeepEqual(want, got) { + t.Fatalf("%s: la=%d lb=%d want %v got %v", label, len(a), len(b), want, got) + } +} + +func TestUnion2By2NEONBoundaryMatrix(t *testing.T) { + for _, sz := range [][2]int{{1, 1}, {2, 8}, {7, 8}, {8, 8}, {8, 9}, {8, 16}, {15, 16}, {16, 17}} { + a := make([]uint16, sz[0]) + b := make([]uint16, sz[1]) + for i := range a { + a[i] = uint16(2 * i) + } + for i := range b { + b[i] = uint16(3*i + 1) + } + checkKernel(t, a, b, "boundary") + checkKernel(t, b, a, "boundary-swap") + } +} + +// A duplicate straddling lanes 7 and 0, and blocks touching 0xFFFF. +func TestUnion2By2NEONLaneStraddleAndHighEnd(t *testing.T) { + a := []uint16{0, 2, 4, 6, 7, 20, 22, 24} + b := []uint16{1, 3, 5, 7, 21, 23, 25, 27} + checkKernel(t, a, b, "straddle") + checkKernel(t, b, a, "straddle-swap") + + high := make([]uint16, 8) + for i := range high { + high[i] = uint16(65528 + i) + } + checkKernel(t, high, high, "identical-high") + low := make([]uint16, 8) + for i := range low { + low[i] = uint16(65520 + i) + } + checkKernel(t, low, high, "adjacent-high") +} + +// Every block seam shares its boundary value, pinning the fast path's +// choice of >= over >: the duplicate must fall to the dedup chain. +func TestUnion2By2NEONEqualityStreak(t *testing.T) { + for blocks := 2; blocks <= 64; blocks *= 2 { + var a, b []uint16 + next := uint16(0) + for i := 0; i < blocks; i++ { + for l := 0; l < 8; l++ { + a = append(a, next+uint16(l)) + } + next += 7 // b's block starts at a's block's last value + for l := 0; l < 8; l++ { + b = append(b, next+uint16(l)) + } + next += 7 // and a's next block starts at b's last value + } + checkKernel(t, a, b, "equality-streak") + checkKernel(t, b, a, "equality-streak-swap") + + want := refUnion(a, b) + shared := make([]uint16, len(a)+len(b)) + copy(shared[len(b):], a) + n := unionNEON(shared[len(b):], b, shared) + if !reflect.DeepEqual(want, shared[:n]) { + t.Fatalf("equality-streak aliased, blocks=%d: want %d got %d elems", + blocks, len(want), n) + } + } +} + +// Tightest alias geometry: set1 at output offset len(set2) with a +// one-block set2, forcing the lookahead tail. +func TestUnion2By2NEONMinimumGapAlias(t *testing.T) { + set2 := []uint16{0, 1, 2, 3, 4, 5, 6, 7} + set1 := make([]uint16, 64) + for i := range set1 { + set1[i] = uint16(100 + i) + } + want := refUnion(set1, set2) + shared := make([]uint16, len(set1)+len(set2)) + copy(shared[len(set2):], set1) + n := unionNEON(shared[len(set2):], set2, shared) + if !reflect.DeepEqual(want, shared[:n]) { + t.Fatalf("minimum-gap alias: want %d elems got %d", len(want), n) + } +} + +// The kernel's 16-byte stores must never touch beyond len(set1)+len(set2). +func TestUnion2By2NEONBufferCanaries(t *testing.T) { + r := rand.New(rand.NewSource(31337)) + for iter := 0; iter < 300; iter++ { + a := genSortedUnique(r, 8+r.Intn(1000), 8192) + b := genSortedUnique(r, 8+r.Intn(1000), 8192) + need := len(a) + len(b) + full := make([]uint16, need+16) + for i := need; i < len(full); i++ { + full[i] = 0xDEAD + } + n := unionNEON(a, b, full[0:0:need]) + if !reflect.DeepEqual(refUnion(a, b), full[:n]) { + t.Fatalf("canary iter %d: wrong result", iter) + } + for i := need; i < len(full); i++ { + if full[i] != 0xDEAD { + t.Fatalf("canary iter %d: guard word %d clobbered (0x%04X)", iter, i, full[i]) + } + } + } +} + +// lazyorArray passes a zero-length buffer with spare capacity. +func TestUnion2By2NEONZeroLenBuffer(t *testing.T) { + r := rand.New(rand.NewSource(99)) + for iter := 0; iter < 200; iter++ { + a := genSortedUnique(r, 8+r.Intn(500), 4096) + b := genSortedUnique(r, 8+r.Intn(500), 4096) + want := refUnion(a, b) + buffer := make([]uint16, 0, len(a)+len(b)) + n := union2by2(a, b, buffer) + got := buffer[:n] + if !reflect.DeepEqual(want, got) { + t.Fatalf("zero-len buffer: mismatch (want %d got %d elems)", len(want), n) + } + } +} + +// iorArray's geometry: set1 lives in the upper region of the output array. +func TestUnion2By2NEONAliasedBuffer(t *testing.T) { + r := rand.New(rand.NewSource(7)) + for iter := 0; iter < 500; iter++ { + valRange := 256 + r.Intn(65280) + set1 := genSortedUnique(r, 8+r.Intn(2000), valRange) + set2 := genSortedUnique(r, 8+r.Intn(2000), valRange) + want := refUnion(set1, set2) + + max := len(set1) + len(set2) + shared := make([]uint16, max) + copy(shared[len(set2):max], set1) + got := shared[:union2by2(shared[len(set2):max], set2, shared)] + if !reflect.DeepEqual(want, got) { + t.Fatalf("aliased: len1=%d len2=%d: mismatch (want %d got %d elems)", + len(set1), len(set2), len(want), len(got)) + } + } +} diff --git a/setutil_neon_bench_arm64_test.go b/setutil_neon_bench_arm64_test.go new file mode 100644 index 00000000..17a0d8cf --- /dev/null +++ b/setutil_neon_bench_arm64_test.go @@ -0,0 +1,73 @@ +//go:build arm64 && !gccgo && !appengine +// +build arm64,!gccgo,!appengine + +package roaring + +import ( + "fmt" + "math/rand" + "testing" +) + +func benchPairSeed(shape string, n int, seed int64) (a, b []uint16) { + r := rand.New(rand.NewSource(int64(42+n) + seed*7919)) + switch shape { + case "dense50": + return genSortedUnique(r, n, 2*n), genSortedUnique(r, n, 2*n) + case "sparse6": + vr := n * 16 + if vr > 65536 { + vr = 65536 + } + return genSortedUnique(r, n, vr), genSortedUnique(r, n, vr) + case "runs16": + a = make([]uint16, n) + b = make([]uint16, n) + for i := 0; i < n; i++ { + blk, off := i/16, i%16 + a[i] = uint16(blk*32 + off) + b[i] = uint16(blk*32 + 16 + off) + } + return a, b + case "spread": // density mismatch: a confined to 1/8 of b's range + vr := 2 * n + if vr > 8192 { + vr = 8192 + } + return genSortedUnique(r, n, vr), genSortedUnique(r, n, 65536) + } + panic("unknown shape") +} + +// Rotating fixed-seed variants keeps the branch predictor from memorizing +// one pair's decision sequence, which flatters the scalar path. +const benchVariants = 16 + +func BenchmarkUnion2By2(b *testing.B) { + for _, shape := range []string{"dense50", "sparse6", "runs16", "spread"} { + for _, n := range []int{256, 384, 512, 1024, 4096} { + as := make([][]uint16, benchVariants) + bs := make([][]uint16, benchVariants) + for v := 0; v < benchVariants; v++ { + as[v], bs[v] = benchPairSeed(shape, n, int64(v)) + } + buffer := make([]uint16, 2*n) + for _, impl := range []struct { + name string + fn func([]uint16, []uint16, []uint16) int + }{ + {"dispatch", union2by2}, + {"scalar", union2by2scalar}, + } { + b.Run(fmt.Sprintf("%s/%d/%s", shape, n, impl.name), func(b *testing.B) { + sink := 0 + for i := 0; i < b.N; i++ { + v := i % benchVariants + sink += impl.fn(as[v], bs[v], buffer) + } + _ = sink + }) + } + } + } +} From 605f0e1cb54484369b04a110fac75913ff196695 Mon Sep 17 00:00:00 2001 From: gitRasheed Date: Fri, 31 Jul 2026 22:03:21 +0100 Subject: [PATCH 2/4] arm64: fast-forward loop for runs of non-interleaving blocks When the disjoint check passes, stay in a small loop that selects blocks with a predicted branch instead of returning to the main loop's CSEL select. On run-structured data the branch predicts and the untaken cursor stays off the loop-carried dependency chain; the first interleaving block merges and rejoins the main loop. Run-structured unions improve about 2.5x over the base kernel; random data never enters the loop. --- setutil_neon_arm64.s | 58 +++++++++++++++++++++++++++++++++----------- 1 file changed, 44 insertions(+), 14 deletions(-) diff --git a/setutil_neon_arm64.s b/setutil_neon_arm64.s index f3ec77c0..8e642c06 100644 --- a/setutil_neon_arm64.s +++ b/setutil_neon_arm64.s @@ -37,13 +37,13 @@ VZIP1 V5.H8, V4.H8, V2.H8 \ VZIP2 V5.H8, V4.H8, V0.H8 -// Store the lanes of V2 that differ from their predecessor (previous lane, +// Store the lanes of Vin that differ from their predecessor (previous lane, // or last lane of V1 for lane 0) at (Rout), advancing Rout by 2 bytes per // lane kept. Writes a full 16 bytes; the caller guarantees slack. // Rcnt receives the number of lanes kept. Clobbers V3-V5, R14, R15, R16. -#define STOREUNIQ(Rout, Rcnt) \ - VEXT $14, V2.B16, V1.B16, V3.B16 \ - VCMEQ V3.H8, V2.H8, V4.H8 \ +#define STOREUNIQ(Vin, Rout, Rcnt) \ + VEXT $14, Vin.B16, V1.B16, V3.B16 \ + VCMEQ V3.H8, Vin.H8, V4.H8 \ VUZP1 V4.B16, V4.B16, V4.B16 \ VMOV V4.D[0], R14 \ AND R12, R14, R14 \ @@ -54,7 +54,7 @@ SUB R14, R19, Rcnt \ ADD R15<<4, R6, R16 \ VLD1 (R16), [V5.B16] \ - VTBL V5.B16, [V2.B16], V3.B16 \ + VTBL V5.B16, [Vin.B16], V3.B16 \ VST1 [V3.B16], (Rout) \ ADD Rcnt<<1, Rout, Rout @@ -86,7 +86,7 @@ TEXT ·unionKernelNEON(SB), NOSPLIT, $0-120 MERGE // laststore = all ones (never equal to a first stored lane) VCMEQ V0.H8, V0.H8, V1.H8 - STOREUNIQ(R2, R17) + STOREUNIQ(V2, R2, R17) VORR V2.B16, V2.B16, V1.B16 loop: @@ -111,22 +111,52 @@ loop: CMP R21, R20 BHS disjoint MERGE - STOREUNIQ(R2, R17) + STOREUNIQ(V2, R2, R17) VORR V2.B16, V2.B16, V1.B16 B loop +// Fast-forward loop for consecutive disjoint blocks. The head select is a +// plain branch: on run-structured data it predicts, keeping the untaken +// cursor off the loop-carried chain that the main loop's CSEL select +// serializes. Interleaving input rejoins the main loop after one merge. disjoint: - VORR V2.B16, V2.B16, V22.B16 // stash fresh - VORR V0.B16, V0.B16, V2.B16 // emit old carry as this round's minimum - STOREUNIQ(R2, R17) - VORR V2.B16, V2.B16, V1.B16 // laststore = old carry - VORR V22.B16, V22.B16, V0.B16 // carry = fresh + STOREUNIQ(V0, R2, R17) + VORR V0.B16, V0.B16, V1.B16 // laststore = old carry + VORR V2.B16, V2.B16, V0.B16 // carry = fresh + +ffloop: + CMP R3, R0 + BHS done + CMP R4, R1 + BHS done + MOVHU (R0), R8 + MOVHU (R1), R9 + CMP R9, R8 + BHI fftake2 + VLD1.P 16(R0), [V2.H8] // take set1's block on tie/less + MOVD R8, R20 + B ffcheck +fftake2: + VLD1.P 16(R1), [V2.H8] + MOVD R9, R20 +ffcheck: + VMOV V0.H[7], R21 + CMP R21, R20 + BLO ffmerge + STOREUNIQ(V0, R2, R17) + VORR V0.B16, V0.B16, V1.B16 // laststore = old carry + VORR V2.B16, V2.B16, V0.B16 // carry = fresh + B ffloop + +ffmerge: + MERGE + STOREUNIQ(V2, R2, R17) + VORR V2.B16, V2.B16, V1.B16 B loop done: // flush carry through the dedup into the leftover buffer - VORR V0.B16, V0.B16, V2.B16 - STOREUNIQ(R7, R17) + STOREUNIQ(V0, R7, R17) MOVD R17, leftoverLen+112(FP) // outLen = (out - outBase) / 2 From ce010c16aebfc95a8df5fc1bb772988185a9d155 Mon Sep 17 00:00:00 2001 From: gitRasheed Date: Mon, 3 Aug 2026 17:54:52 +0100 Subject: [PATCH 3/4] arm64: fall back to scalar union when set2 aliases the output --- roaring_test.go | 57 ++++++++++++++++++++++++++++++++++++++ setutil_arm64.go | 7 +++++ setutil_neon_arm64_test.go | 18 ++++++++++++ 3 files changed, 82 insertions(+) diff --git a/roaring_test.go b/roaring_test.go index 6b31c61f..4cb90bd0 100644 --- a/roaring_test.go +++ b/roaring_test.go @@ -4166,3 +4166,60 @@ func TestBitmapOrBulkMergeCopyOnWriteTailOwnership(t *testing.T) { t.Fatalf("source became invalid after tail mutations: %v", err) } } + +// Self-union must leave the bitmap unchanged. Spare capacity sends iorArray +// through its reuse branch, where set2 and the output share a backing array. +func TestOrSelfInPlace(t *testing.T) { + buildRange := func(n int) *Bitmap { + rb := New() + for v := 0; v < n; v++ { + rb.Add(uint32(v)) + } + return rb + } + cases := []struct { + name string + build func() *Bitmap + }{ + {"small", func() *Bitmap { + rb := buildRange(1024) + rb.RemoveRange(383, 1024) + return rb + }}, + {"neon-sized", func() *Bitmap { + rb := buildRange(1024) + rb.RemoveRange(384, 1024) + return rb + }}, + {"strided-exact-cap", func() *Bitmap { + rb := buildRange(1024) + for v := 1; v < 1024; v += 2 { + rb.Remove(uint32(v)) + } + return rb + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + rb := tc.build() + ac, ok := rb.highlowcontainer.getContainerAtIndex(0).(*arrayContainer) + if !ok || cap(ac.content) < 2*len(ac.content) { + t.Fatal("container no longer hits iorArray's capacity-reuse branch") + } + want := rb.ToArray() + rb.Or(rb) + if err := rb.Validate(); err != nil { + t.Fatalf("invalid after self-Or: %v", err) + } + got := rb.ToArray() + if len(got) != len(want) { + t.Fatalf("self-Or changed cardinality: got %d want %d", len(got), len(want)) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("self-Or corrupted index %d: got %d want %d", i, got[i], want[i]) + } + } + }) + } +} diff --git a/setutil_arm64.go b/setutil_arm64.go index bf52c34c..24245fa2 100644 --- a/setutil_arm64.go +++ b/setutil_arm64.go @@ -46,6 +46,13 @@ func union2by2(set1 []uint16, set2 []uint16, buffer []uint16) int { func unionNEON(set1 []uint16, set2 []uint16, buffer []uint16) int { // Callers such as lazyorArray pass a zero-length buffer with capacity. buffer = buffer[:cap(buffer)] + // iorArray's in-place self-union passes set2 and buffer sharing a + // backing array from offset 0, with set1 a copy of set2. The kernel's + // block stores would clobber unread set2; on identical inputs the + // scalar merge never writes a position it has not already read. + if &buffer[0] == &set2[0] { + return union2by2scalar(set1, set2, buffer) + } var leftover [16]uint16 outLen, pos1, pos2, ll := unionKernelNEON(set1, set2, buffer, &uniqshuf[0], &leftover) // The leftovers and the exhausted input's tail are two sorted runs. diff --git a/setutil_neon_arm64_test.go b/setutil_neon_arm64_test.go index eaa7e893..20e289b9 100644 --- a/setutil_neon_arm64_test.go +++ b/setutil_neon_arm64_test.go @@ -229,3 +229,21 @@ func TestUnion2By2NEONAliasedBuffer(t *testing.T) { } } } + +// iorArray's self-union geometry: set2 and the output share a backing array. +func TestUnion2By2NEONSelfAliasedBuffer(t *testing.T) { + r := rand.New(rand.NewSource(11)) + for iter := 0; iter < 200; iter++ { + valRange := 256 + r.Intn(65280) + set := genSortedUnique(r, 8+r.Intn(2000), valRange) + + n := len(set) + shared := make([]uint16, 2*n) + copy(shared, set) + copy(shared[n:], set) + got := shared[:union2by2(shared[n:], shared[:n], shared)] + if !reflect.DeepEqual(set, got) { + t.Fatalf("self-aliased: n=%d: mismatch (got %d elems)", n, len(got)) + } + } +} From b20d1128092fd8a5c2efda1d265c92a244dcdb9a Mon Sep 17 00:00:00 2001 From: gitRasheed Date: Mon, 3 Aug 2026 19:58:02 +0100 Subject: [PATCH 4/4] arm64: independent union oracle and multi-block 0xFFFF coverage --- setutil_neon_arm64_test.go | 39 +++++++++++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/setutil_neon_arm64_test.go b/setutil_neon_arm64_test.go index 20e289b9..45eff11a 100644 --- a/setutil_neon_arm64_test.go +++ b/setutil_neon_arm64_test.go @@ -9,9 +9,10 @@ import ( "testing" ) +// The oracle is the pre-existing scalar asm, sharing no code with unionNEON. func refUnion(a, b []uint16) []uint16 { out := make([]uint16, len(a)+len(b)) - n := scalarMergeUnion(a, b, out) + n := union2by2scalar(a, b, out) return out[:n] } @@ -124,6 +125,42 @@ func TestUnion2By2NEONLaneStraddleAndHighEnd(t *testing.T) { checkKernel(t, low, high, "adjacent-high") } +// 20 lands in both the kernel leftovers and set1's tail; the wrapper dedups. +func TestUnion2By2NEONLeftoverTailDedup(t *testing.T) { + a := []uint16{0, 1, 2, 3, 4, 5, 6, 7, 20} + b := []uint16{15, 16, 17, 18, 19, 20, 21, 22} + checkKernel(t, a, b, "leftover-tail") + checkKernel(t, b, a, "leftover-tail-swap") +} + +// laststore initializes to all ones and the random generators cannot +// emit 0xFFFF. These shapes run it through the merge loop, the disjoint +// carry, and the wrapper dedup, not just prime-and-flush. +func TestUnion2By2NEONHighEndMultiBlock(t *testing.T) { + span := func(lo, hi, step int) []uint16 { + var s []uint16 + for v := lo; v <= hi; v += step { + s = append(s, uint16(v)) + } + return s + } + + a := span(65520, 65535, 1) + b := append(span(0, 7, 1), span(65528, 65535, 1)...) + checkKernel(t, a, b, "high-two-blocks") + checkKernel(t, b, a, "high-two-blocks-swap") + + // 0xFFFF shared between the kernel leftovers and set1's tail. + c := append(span(0, 7, 1), span(65529, 65535, 1)...) + d := span(65528, 65535, 1) + checkKernel(t, c, d, "high-leftover-tail") + checkKernel(t, d, c, "high-leftover-tail-swap") + + // Identical blocks to 0xFFFF: equality, disjoint, and ffmerge transitions. + g := span(65472, 65535, 1) + checkKernel(t, g, g, "high-identical-multiblock") +} + // Every block seam shares its boundary value, pinning the fast path's // choice of >= over >: the duplicate must fall to the dedup chain. func TestUnion2By2NEONEqualityStreak(t *testing.T) {