Add AVX2 bitmap intersection stores

Author: perfloop-agent

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]
+	}
+}

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)
+}

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)
+}

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)
+			}
+		}
+	}
+}

bitmapcontainer.go

@@ -867,19 +867,21 @@ 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)
 	ac.content = ac.content[:newcardinality]
 	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 {

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

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)
+	}
+}