diff --git a/and_slices.go b/and_slices.go new file mode 100644 index 00000000..fe75c8ab --- /dev/null +++ b/and_slices.go @@ -0,0 +1,7 @@ +package roaring + +func andStoreSliceGo(dst, s, m []uint64) { + for i := range dst { + dst[i] = s[i] & m[i] + } +} diff --git a/and_slices_amd64.go b/and_slices_amd64.go new file mode 100644 index 00000000..586a5d84 --- /dev/null +++ b/and_slices_amd64.go @@ -0,0 +1,15 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +package roaring + +//go:noescape +func _andStoreSliceAVX2(dst, s, m []uint64) + +func andStoreSlice(dst, s, m []uint64) { + if useAVX2 { + _andStoreSliceAVX2(dst, s, m) + return + } + andStoreSliceGo(dst, s, m) +} diff --git a/and_slices_generic.go b/and_slices_generic.go new file mode 100644 index 00000000..994f88ed --- /dev/null +++ b/and_slices_generic.go @@ -0,0 +1,8 @@ +//go:build !amd64 || appengine +// +build !amd64 appengine + +package roaring + +func andStoreSlice(dst, s, m []uint64) { + andStoreSliceGo(dst, s, m) +} diff --git a/and_slices_test.go b/and_slices_test.go new file mode 100644 index 00000000..5899d468 --- /dev/null +++ b/and_slices_test.go @@ -0,0 +1,32 @@ +package roaring + +import "testing" + +func TestAndStoreSlice(t *testing.T) { + for _, n := range []int{0, 1, 3, 4, 5, bitmapContainerSize} { + left := make([]uint64, n) + right := make([]uint64, n) + for i := range left { + left[i] = uint64(i+1) * 0x5555555555555555 + right[i] = ^uint64(i * 3) + } + + got := make([]uint64, n) + andStoreSlice(got, left, right) + for i := range got { + want := left[i] & right[i] + if got[i] != want { + t.Fatalf("separate destination len=%d index=%d: got %#x, want %#x", n, i, got[i], want) + } + } + + inPlace := append([]uint64(nil), left...) + andStoreSlice(inPlace, inPlace, right) + for i := range inPlace { + want := left[i] & right[i] + if inPlace[i] != want { + t.Fatalf("in-place destination len=%d index=%d: got %#x, want %#x", n, i, inPlace[i], want) + } + } + } +} diff --git a/bitmapcontainer.go b/bitmapcontainer.go index 416da732..65904bda 100644 --- a/bitmapcontainer.go +++ b/bitmapcontainer.go @@ -867,12 +867,7 @@ func (bc *bitmapContainer) getCardinalityInRange(start, end uint) int { func (bc *bitmapContainer) andBitmap(value2 *bitmapContainer) container { newcardinality := int(popcntAndSlice(bc.bitmap, value2.bitmap)) if newcardinality > arrayDefaultMaxSize { - answer := newBitmapContainer() - for k := 0; k < len(answer.bitmap); k++ { - answer.bitmap[k] = bc.bitmap[k] & value2.bitmap[k] - } - answer.cardinality = newcardinality - return answer + return andBitmapStore(bc, value2, newcardinality) } ac := newArrayContainerSize(newcardinality) fillArrayAND(ac.content, bc.bitmap, value2.bitmap) @@ -880,6 +875,13 @@ func (bc *bitmapContainer) andBitmap(value2 *bitmapContainer) container { return ac } +func andBitmapStore(left, right *bitmapContainer, cardinality int) *bitmapContainer { + answer := newBitmapContainer() + andStoreSlice(answer.bitmap, left.bitmap, right.bitmap) + answer.cardinality = cardinality + return answer +} + func (bc *bitmapContainer) intersectsArray(value2 *arrayContainer) bool { c := value2.getCardinality() for k := 0; k < c; k++ { @@ -902,9 +904,7 @@ func (bc *bitmapContainer) intersectsBitmap(value2 *bitmapContainer) bool { func (bc *bitmapContainer) iandBitmap(value2 *bitmapContainer) container { newcardinality := int(popcntAndSlice(bc.bitmap, value2.bitmap)) - for k := 0; k < len(bc.bitmap); k++ { - bc.bitmap[k] = bc.bitmap[k] & value2.bitmap[k] - } + andStoreSlice(bc.bitmap, bc.bitmap, value2.bitmap) bc.cardinality = newcardinality if newcardinality <= arrayDefaultMaxSize { diff --git a/bitmapcontainer_bench_test.go b/bitmapcontainer_bench_test.go index 530772ac..a8a7d23e 100644 --- a/bitmapcontainer_bench_test.go +++ b/bitmapcontainer_bench_test.go @@ -68,3 +68,53 @@ func BenchmarkParOrBitmapContainers(b *testing.B) { } } } + +func roaringAndBitmapFixture(ranges ...[2]int) *Bitmap { + words := make([]uint64, bitmapContainerSize) + for _, r := range ranges { + for value := r[0]; value < r[1]; value++ { + words[value/64] |= uint64(1) << uint(value&63) + } + } + return FromDense(words, false) +} + +func benchmarkRoaringAndBitmap(b *testing.B, left, right *Bitmap, expected uint64) { + b.Helper() + for _, bitmap := range []*Bitmap{left, right} { + if bitmap.highlowcontainer.size() != 1 { + b.Fatal("workload did not produce one container") + } + if _, ok := bitmap.highlowcontainer.getContainerAtIndex(0).(*bitmapContainer); !ok { + b.Fatal("workload did not produce bitmap containers") + } + } + + if result := And(left, right); result.GetCardinality() != expected { + b.Fatalf("unexpected cardinality: got %d, want %d", result.GetCardinality(), expected) + } + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + And(left, right) + } +} + +func BenchmarkRoaringAndBitmapDense(b *testing.B) { + left := roaringAndBitmapFixture([2]int{0, maxCapacity}) + right := roaringAndBitmapFixture([2]int{0, 59000}) + benchmarkRoaringAndBitmap(b, left, right, 59000) +} + +func BenchmarkRoaringAndBitmapSparse(b *testing.B) { + left := roaringAndBitmapFixture([2]int{0, 5000}, [2]int{10000, 15000}) + right := roaringAndBitmapFixture([2]int{0, 5000}, [2]int{20000, 25000}) + benchmarkRoaringAndBitmap(b, left, right, 5000) +} + +func BenchmarkRoaringAndBitmapDisjoint(b *testing.B) { + left := roaringAndBitmapFixture([2]int{0, 5000}) + right := roaringAndBitmapFixture([2]int{10000, 15000}) + benchmarkRoaringAndBitmap(b, left, right, 0) +} diff --git a/popcnt_avx2_amd64.s b/popcnt_avx2_amd64.s index 52e2c989..8317d730 100644 --- a/popcnt_avx2_amd64.s +++ b/popcnt_avx2_amd64.s @@ -358,3 +358,42 @@ TEXT ·_hasAVX2(SB), NOSPLIT, $0-1 noavx2: SETEQ ret+0(FP) // ZF is still set by whichever TESTL ran last RET + +// func _andStoreSliceAVX2(dst, s, m []uint64) +// +// Writes s[i] & m[i] to dst[i]. The three slices must have equal lengths. +// Four uint64 words are processed per AVX2 iteration; a scalar tail handles +// any remaining words. +TEXT ·_andStoreSliceAVX2(SB), NOSPLIT, $0-72 + MOVQ dst_base+0(FP), DI + MOVQ s_base+24(FP), SI + MOVQ m_base+48(FP), DX + MOVQ dst_len+8(FP), CX + + MOVQ CX, R8 + SHRQ $2, R8 + JZ andstoretail +andstoreloop: + VMOVDQU (SI), Y0 + VPAND (DX), Y0, Y0 + VMOVDQU Y0, (DI) + ADDQ $32, DI + ADDQ $32, SI + ADDQ $32, DX + DECQ R8 + JNZ andstoreloop +andstoretail: + ANDL $3, CX + JZ andstoredone +andstoretailloop: + MOVQ (SI), AX + ANDQ (DX), AX + MOVQ AX, (DI) + ADDQ $8, DI + ADDQ $8, SI + ADDQ $8, DX + DECL CX + JNZ andstoretailloop +andstoredone: + VZEROUPPER + RET diff --git a/popcnt_avx2_amd64_test.go b/popcnt_avx2_amd64_test.go index c0e91b0f..9fda79be 100644 --- a/popcnt_avx2_amd64_test.go +++ b/popcnt_avx2_amd64_test.go @@ -127,3 +127,120 @@ func TestAVX2PopcntDifferential(t *testing.T) { } } } + +func TestAVX2AndStoreSliceDispatch(t *testing.T) { + saved := useAVX2 + defer func() { useAVX2 = saved }() + + r := rand.New(rand.NewSource(11)) + for _, on := range []bool{false, true} { + if on && !saved { + continue // CPU has no AVX2; only the fallback exists + } + useAVX2 = on + for _, n := range avx2TestLengths { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + expected := append([]uint64(nil), s...) + andStoreSliceGo(expected, s, m) + + dst := append([]uint64(nil), s...) + andStoreSlice(dst, s, m) + assert.Equalf(t, expected, dst, "separate destination avx2=%v len=%d", on, n) + + inPlaceLeft := append([]uint64(nil), s...) + andStoreSlice(inPlaceLeft, inPlaceLeft, m) + assert.Equalf(t, expected, inPlaceLeft, "left alias avx2=%v len=%d", on, n) + + inPlaceRight := append([]uint64(nil), m...) + andStoreSlice(inPlaceRight, s, inPlaceRight) + assert.Equalf(t, expected, inPlaceRight, "right alias avx2=%v len=%d", on, n) + } + } +} + +func TestAVX2AndStoreSliceDifferential(t *testing.T) { + if !useAVX2 { + t.Skip("AVX2 not available on this CPU") + } + r := rand.New(rand.NewSource(13)) + for _, n := range avx2TestLengths { + for iter := 0; iter < 64; iter++ { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + expected := append([]uint64(nil), s...) + andStoreSliceGo(expected, s, m) + + dst := append([]uint64(nil), s...) + _andStoreSliceAVX2(dst, s, m) + assert.Equalf(t, expected, dst, "separate destination len=%d", n) + + inPlaceLeft := append([]uint64(nil), s...) + _andStoreSliceAVX2(inPlaceLeft, inPlaceLeft, m) + assert.Equalf(t, expected, inPlaceLeft, "left alias len=%d", n) + + inPlaceRight := append([]uint64(nil), m...) + _andStoreSliceAVX2(inPlaceRight, s, inPlaceRight) + assert.Equalf(t, expected, inPlaceRight, "right alias len=%d", n) + } + } +} + +func TestAVX2BitmapAndStoreDispatch(t *testing.T) { + saved := useAVX2 + defer func() { useAVX2 = saved }() + + newFull := func() *bitmapContainer { + bc := newBitmapContainer() + for i := range bc.bitmap { + bc.bitmap[i] = ^uint64(0) + } + bc.cardinality = maxCapacity + return bc + } + + right := newBitmapContainer() + for i := range right.bitmap { + right.bitmap[i] = 0xaaaaaaaaaaaaaaaa + } + right.cardinality = maxCapacity / 2 + + small := newBitmapContainer() + small.bitmap[0] = 1 + small.cardinality = 1 + + for _, on := range []bool{false, true} { + if on && !saved { + continue // CPU has no AVX2; only the fallback exists + } + useAVX2 = on + + answer, ok := newFull().andBitmap(right).(*bitmapContainer) + if !ok { + t.Fatalf("expected bitmap result with avx2=%v", on) + } + assert.Equalf(t, right.bitmap, answer.bitmap, "bitmap result avx2=%v", on) + assert.Equalf(t, right.cardinality, answer.cardinality, "bitmap cardinality avx2=%v", on) + + inPlace := newFull() + answer, ok = inPlace.iandBitmap(right).(*bitmapContainer) + if !ok { + t.Fatalf("expected in-place bitmap result with avx2=%v", on) + } + assert.Samef(t, inPlace, answer, "in-place bitmap result avx2=%v", on) + assert.Equalf(t, right.bitmap, inPlace.bitmap, "in-place bitmap contents avx2=%v", on) + + arrayResult, ok := newFull().andBitmap(small).(*arrayContainer) + if !ok { + t.Fatalf("expected array result with avx2=%v", on) + } + assert.Equalf(t, []uint16{0}, arrayResult.content, "array result avx2=%v", on) + + inPlace = newFull() + arrayResult, ok = inPlace.iandBitmap(small).(*arrayContainer) + if !ok { + t.Fatalf("expected in-place array result with avx2=%v", on) + } + assert.Equalf(t, []uint16{0}, arrayResult.content, "in-place array result avx2=%v", on) + } +}