Skip to content

Commit 454f892

Browse files
gitRasheedgitRasheed
committed
arm64: NEON-vectorized intersection2by2 and intersection2by2Cardinality
All-pairs 8x8 block compare with a two-way range gate that fast-forwards disjoint runs on scalar boundary loads. The materializing kernel compacts matched lanes through a mask-indexed shuffle table with dual store bounds for the exact-capacity and in-place caller contracts, and spills a retained block when output may alias unread input. Dispatch floor 16 from a size sweep on three core families; the galloping paths and the boolean variant keep the existing scalar code.
1 parent 438e356 commit 454f892

7 files changed

Lines changed: 951 additions & 28 deletions

‎roaring_test.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4166,3 +4166,36 @@ func TestBitmapOrBulkMergeCopyOnWriteTailOwnership(t *testing.T) {
41664166
t.Fatalf("source became invalid after tail mutations: %v", err)
41674167
}
41684168
}
4169+
4170+
// iandArray passes the container's array as both inputs and the output.
4171+
func TestAndSelfInPlace(t *testing.T) {
4172+
cases := []struct {
4173+
name string
4174+
n int
4175+
}{
4176+
{"below-threshold", 10},
4177+
{"neon-sized", 500},
4178+
}
4179+
for _, tc := range cases {
4180+
t.Run(tc.name, func(t *testing.T) {
4181+
rb := New()
4182+
for v := 0; v < tc.n; v++ {
4183+
rb.Add(uint32(v * 3))
4184+
}
4185+
want := rb.ToArray()
4186+
rb.And(rb)
4187+
if err := rb.Validate(); err != nil {
4188+
t.Fatalf("invalid after self-And: %v", err)
4189+
}
4190+
got := rb.ToArray()
4191+
if len(got) != len(want) {
4192+
t.Fatalf("self-And changed cardinality: got %d want %d", len(got), len(want))
4193+
}
4194+
for i := range want {
4195+
if got[i] != want[i] {
4196+
t.Fatalf("self-And corrupted index %d: got %d want %d", i, got[i], want[i])
4197+
}
4198+
}
4199+
})
4200+
}
4201+
}

‎setutil.go‎

Lines changed: 0 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -169,34 +169,6 @@ func union2by2Cardinality(set1 []uint16, set2 []uint16) int {
169169
return pos
170170
}
171171

172-
func intersection2by2(
173-
set1 []uint16,
174-
set2 []uint16,
175-
buffer []uint16,
176-
) int {
177-
if len(set1)*64 < len(set2) {
178-
return onesidedgallopingintersect2by2(set1, set2, buffer)
179-
} else if len(set2)*64 < len(set1) {
180-
return onesidedgallopingintersect2by2(set2, set1, buffer)
181-
} else {
182-
return localintersect2by2(set1, set2, buffer)
183-
}
184-
}
185-
186-
// intersection2by2Cardinality computes the cardinality of the intersection
187-
func intersection2by2Cardinality(
188-
set1 []uint16,
189-
set2 []uint16,
190-
) int {
191-
if len(set1)*64 < len(set2) {
192-
return onesidedgallopingintersect2by2Cardinality(set1, set2)
193-
} else if len(set2)*64 < len(set1) {
194-
return onesidedgallopingintersect2by2Cardinality(set2, set1)
195-
} else {
196-
return localintersect2by2Cardinality(set1, set2)
197-
}
198-
}
199-
200172
// intersects2by2 computes whether the two sets intersect
201173
func intersects2by2(
202174
set1 []uint16,

‎setutil_generic.go‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,34 @@
33

44
package roaring
55

6+
func intersection2by2(
7+
set1 []uint16,
8+
set2 []uint16,
9+
buffer []uint16,
10+
) int {
11+
if len(set1)*64 < len(set2) {
12+
return onesidedgallopingintersect2by2(set1, set2, buffer)
13+
} else if len(set2)*64 < len(set1) {
14+
return onesidedgallopingintersect2by2(set2, set1, buffer)
15+
} else {
16+
return localintersect2by2(set1, set2, buffer)
17+
}
18+
}
19+
20+
// intersection2by2Cardinality computes the cardinality of the intersection
21+
func intersection2by2Cardinality(
22+
set1 []uint16,
23+
set2 []uint16,
24+
) int {
25+
if len(set1)*64 < len(set2) {
26+
return onesidedgallopingintersect2by2Cardinality(set1, set2)
27+
} else if len(set2)*64 < len(set1) {
28+
return onesidedgallopingintersect2by2Cardinality(set2, set1)
29+
} else {
30+
return localintersect2by2Cardinality(set1, set2)
31+
}
32+
}
33+
634
func union2by2(set1 []uint16, set2 []uint16, buffer []uint16) int {
735
pos := 0
836
k1 := 0

‎setutil_intersect_arm64.go‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
//go:build arm64 && !gccgo && !appengine
2+
// +build arm64,!gccgo,!appengine
3+
4+
package roaring
5+
6+
//go:noescape
7+
func intersectCardKernelNEON(set1, set2 []uint16) (card, pos1, pos2 int)
8+
9+
//go:noescape
10+
func intersectKernelNEON(set1, set2, buffer []uint16, shuf *byte, spill *[8]uint16) (outLen, pos1, pos2, spilled int)
11+
12+
// Below this the kernel setup usually loses to the scalar path.
13+
// Duplicate values get unspecified results; stores stay within capacity.
14+
const (
15+
neonIntersectCardThreshold = 16
16+
neonIntersectThreshold = 16
17+
)
18+
19+
func intersection2by2(
20+
set1 []uint16,
21+
set2 []uint16,
22+
buffer []uint16,
23+
) int {
24+
if len(set1)*64 < len(set2) {
25+
return onesidedgallopingintersect2by2(set1, set2, buffer)
26+
} else if len(set2)*64 < len(set1) {
27+
return onesidedgallopingintersect2by2(set2, set1, buffer)
28+
}
29+
if len(set1) < neonIntersectThreshold || len(set2) < neonIntersectThreshold {
30+
return localintersect2by2(set1, set2, buffer)
31+
}
32+
if set1[len(set1)-1] < set2[0] || set2[len(set2)-1] < set1[0] {
33+
return 0
34+
}
35+
// andArray passes len 0; COW cloning keeps iandArray's spare cap private.
36+
buffer = buffer[:cap(buffer)]
37+
var spill [8]uint16
38+
outLen, pos1, pos2, spilled := intersectKernelNEON(set1, set2, buffer, &uniqshuf[0], &spill)
39+
if spilled != 0 {
40+
// set1 may alias buffer: drain from this copy, never reread.
41+
i := 0
42+
for i < 8 && pos2 < len(set2) {
43+
switch {
44+
case spill[i] < set2[pos2]:
45+
i++
46+
case set2[pos2] < spill[i]:
47+
pos2++
48+
default:
49+
buffer[outLen] = spill[i]
50+
outLen++
51+
i++
52+
pos2++
53+
}
54+
}
55+
}
56+
return outLen + localintersect2by2(set1[pos1:], set2[pos2:], buffer[outLen:])
57+
}
58+
59+
func intersection2by2Cardinality(
60+
set1 []uint16,
61+
set2 []uint16,
62+
) int {
63+
if len(set1)*64 < len(set2) {
64+
return onesidedgallopingintersect2by2Cardinality(set1, set2)
65+
} else if len(set2)*64 < len(set1) {
66+
return onesidedgallopingintersect2by2Cardinality(set2, set1)
67+
}
68+
if len(set1) < neonIntersectCardThreshold || len(set2) < neonIntersectCardThreshold {
69+
return localintersect2by2Cardinality(set1, set2)
70+
}
71+
if set1[len(set1)-1] < set2[0] || set2[len(set2)-1] < set1[0] {
72+
return 0
73+
}
74+
card, pos1, pos2 := intersectCardKernelNEON(set1, set2)
75+
return card + localintersect2by2Cardinality(set1[pos1:], set2[pos2:])
76+
}
77+
78+
var uniqshuf = buildUniqshuf()
79+
80+
func buildUniqshuf() (t [256 * 16]byte) {
81+
for m := 0; m < 256; m++ {
82+
pos := 0
83+
for lane := 0; lane < 8; lane++ {
84+
if m&(1<<lane) == 0 {
85+
t[m*16+pos*2] = byte(2 * lane)
86+
t[m*16+pos*2+1] = byte(2*lane + 1)
87+
pos++
88+
}
89+
}
90+
for ; pos < 8; pos++ {
91+
t[m*16+pos*2] = 0xFF
92+
t[m*16+pos*2+1] = 0xFF
93+
}
94+
}
95+
return
96+
}

‎setutil_intersect_bench_arm64_test.go‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
//go:build arm64 && !gccgo && !appengine
2+
// +build arm64,!gccgo,!appengine
3+
4+
package roaring
5+
6+
import (
7+
"fmt"
8+
"math/rand"
9+
"sort"
10+
"testing"
11+
)
12+
13+
const intersectBenchVariants = 8
14+
15+
func benchIntersectPair(shape string, n int, seed int64) (a, b []uint16) {
16+
r := rand.New(rand.NewSource(int64(42+n) + seed*7919))
17+
switch shape {
18+
case "dense50":
19+
return genSortedUnique(r, n, 2*n), genSortedUnique(r, n, 2*n)
20+
case "coinflip": // disjoint 8-blocks, shuffled ownership: zero matches
21+
nb := (n + 7) / 8
22+
own := make([]bool, 2*nb)
23+
for i := 0; i < nb; i++ {
24+
own[i] = true
25+
}
26+
r.Shuffle(len(own), func(i, j int) { own[i], own[j] = own[j], own[i] })
27+
a = make([]uint16, 0, nb*8)
28+
b = make([]uint16, 0, nb*8)
29+
base := 0
30+
for _, toA := range own {
31+
for x := 0; x < 8; x++ {
32+
if toA {
33+
a = append(a, uint16(base+x))
34+
} else {
35+
b = append(b, uint16(base+x))
36+
}
37+
}
38+
base += 8
39+
}
40+
return a[:n], b[:n]
41+
case "skew8": // one side 8x longer, below the 64:1 galloping cutoff
42+
big := 8 * n
43+
if big > 32768 {
44+
big = 32768
45+
}
46+
return genSortedUnique(r, n, 65536), genSortedUnique(r, big, 65536)
47+
case "overlap95": // ~95% shared elements: near-total match density
48+
a = genSortedUnique(r, n, 4*n)
49+
b = append([]uint16(nil), a...)
50+
for i := 10; i < n; i += 20 {
51+
b[i] ^= 1
52+
}
53+
sort.Slice(b, func(i, j int) bool { return b[i] < b[j] })
54+
out := b[:0]
55+
for i, v := range b {
56+
if i == 0 || v != b[i-1] {
57+
out = append(out, v)
58+
}
59+
}
60+
return a, out
61+
}
62+
panic("unknown shape")
63+
}
64+
65+
var benchIntersectShapes = []string{"dense50", "coinflip", "skew8", "overlap95"}
66+
var benchIntersectSizes = []int{8, 16, 24, 64, 256, 4096}
67+
68+
func BenchmarkIntersect2By2(b *testing.B) {
69+
for _, shape := range benchIntersectShapes {
70+
for _, n := range benchIntersectSizes {
71+
as := make([][]uint16, intersectBenchVariants)
72+
bs := make([][]uint16, intersectBenchVariants)
73+
for v := 0; v < intersectBenchVariants; v++ {
74+
as[v], bs[v] = benchIntersectPair(shape, n, int64(v))
75+
}
76+
// andArray allocates exactly min(len1,len2); mirror that cap.
77+
mins := make([]int, intersectBenchVariants)
78+
for v := 0; v < intersectBenchVariants; v++ {
79+
mins[v] = len(as[v])
80+
if len(bs[v]) < mins[v] {
81+
mins[v] = len(bs[v])
82+
}
83+
}
84+
buffer := make([]uint16, n+8)
85+
scratch := make([]uint16, n+8)
86+
for _, impl := range []struct {
87+
name string
88+
fn func([]uint16, []uint16, []uint16) int
89+
}{
90+
{"dispatch", intersection2by2},
91+
{"scalar", localintersect2by2},
92+
} {
93+
b.Run(fmt.Sprintf("%s/%d/%s", shape, n, impl.name), func(b *testing.B) {
94+
sink := 0
95+
for i := 0; i < b.N; i++ {
96+
v := i % intersectBenchVariants
97+
sink += impl.fn(as[v], bs[v], buffer[:0:mins[v]])
98+
}
99+
_ = sink
100+
})
101+
}
102+
// iandArray geometry; both rows pay the same restore copy.
103+
if shape != "dense50" || (n != 16 && n != 4096) {
104+
continue
105+
}
106+
for _, impl := range []struct {
107+
name string
108+
fn func([]uint16, []uint16, []uint16) int
109+
}{
110+
{"inplace", intersection2by2},
111+
{"inplaceScalar", localintersect2by2},
112+
} {
113+
b.Run(fmt.Sprintf("%s/%d/%s", shape, n, impl.name), func(b *testing.B) {
114+
sink := 0
115+
for i := 0; i < b.N; i++ {
116+
v := i % intersectBenchVariants
117+
m := copy(scratch, as[v])
118+
sink += impl.fn(scratch[:m], bs[v], scratch[:0:m])
119+
}
120+
_ = sink
121+
})
122+
}
123+
}
124+
}
125+
}
126+
127+
func BenchmarkIntersectCard2By2(b *testing.B) {
128+
for _, shape := range benchIntersectShapes {
129+
for _, n := range benchIntersectSizes {
130+
as := make([][]uint16, intersectBenchVariants)
131+
bs := make([][]uint16, intersectBenchVariants)
132+
for v := 0; v < intersectBenchVariants; v++ {
133+
as[v], bs[v] = benchIntersectPair(shape, n, int64(v))
134+
}
135+
for _, impl := range []struct {
136+
name string
137+
fn func([]uint16, []uint16) int
138+
}{
139+
{"dispatch", intersection2by2Cardinality},
140+
{"scalar", localintersect2by2Cardinality},
141+
} {
142+
b.Run(fmt.Sprintf("%s/%d/%s", shape, n, impl.name), func(b *testing.B) {
143+
sink := 0
144+
for i := 0; i < b.N; i++ {
145+
v := i % intersectBenchVariants
146+
sink += impl.fn(as[v], bs[v])
147+
}
148+
_ = sink
149+
})
150+
}
151+
}
152+
}
153+
}

0 commit comments

Comments
 (0)