Fuse dense bitmap intersection output and popcount

bitmapcontainer.go

@@ -839,10 +839,19 @@ func (bc *bitmapContainer) getCardinalityInRange(start, end uint) int {
 }
 
 func (bc *bitmapContainer) andBitmap(value2 *bitmapContainer) container {
+	// The intersection is necessarily a bitmap when the two cardinalities
+	// exceed the universe by more than the array threshold. Build and count it
+	// together so the operand bitmaps are traversed only once.
+	if bc.cardinality+value2.cardinality > maxCapacity+arrayDefaultMaxSize {
+		answer := newBitmapContainer()
+		answer.cardinality = int(andPopcntSlice(answer.bitmap, bc.bitmap, value2.bitmap))
+		return answer
+	}
+
 	newcardinality := int(popcntAndSlice(bc.bitmap, value2.bitmap))
 	if newcardinality > arrayDefaultMaxSize {
 		answer := newBitmapContainer()
-		for k := 0; k < len(answer.bitmap); k++ {
+		for k := range answer.bitmap {
 			answer.bitmap[k] = bc.bitmap[k] & value2.bitmap[k]
 		}
 		answer.cardinality = newcardinality

bitmapcontainer_test.go

@@ -58,6 +58,42 @@ func TestBitmapcontainerAndCardinality(t *testing.T) {
 	}
 }
 
+func TestBitmapContainerAndBitmapRepresentation(t *testing.T) {
+	full := newBitmapContainer()
+	fill(full.bitmap, ^uint64(0))
+	full.cardinality = maxCapacity
+
+	t.Run("bitmap result", func(t *testing.T) {
+		other := newBitmapContainer()
+		for i := range other.bitmap {
+			other.bitmap[i] = 0x5555555555555555
+		}
+		other.cardinality = int(popcntSlice(other.bitmap))
+
+		result, ok := full.andBitmap(other).(*bitmapContainer)
+		require.True(t, ok)
+		assert.Equal(t, other.cardinality, result.cardinality)
+		assert.Equal(t, other.bitmap, result.bitmap)
+		assert.NoError(t, result.validate())
+	})
+
+	t.Run("array threshold", func(t *testing.T) {
+		other := newBitmapContainer()
+		for i := 0; i < arrayDefaultMaxSize/64; i++ {
+			other.bitmap[i] = ^uint64(0)
+		}
+		other.cardinality = arrayDefaultMaxSize
+
+		result, ok := full.andBitmap(other).(*arrayContainer)
+		require.True(t, ok)
+		require.Len(t, result.content, arrayDefaultMaxSize)
+		for i, value := range result.content {
+			assert.Equal(t, uint16(i), value)
+		}
+		assert.NoError(t, result.validate())
+	})
+}
+
 func TestIssue181(t *testing.T) {
 	t.Run("Initial issue 181", func(t *testing.T) {
 		a := New()

parallel_benchmark_test.go

@@ -5,6 +5,30 @@ import (
 	"testing"
 )
 
+func bitmapIntersectionBenchmarkFixture(containerCount int, word uint64) *Bitmap {
+	bitmap := make([]uint64, containerCount*(maxCapacity/64))
+	for i := range bitmap {
+		bitmap[i] = word
+	}
+	return FromDense(bitmap, false)
+}
+
+// BenchmarkIntersectionBitmapParallel exercises ParAnd with bitmap containers
+// whose intersection stays above the array-container threshold.
+func BenchmarkIntersectionBitmapParallel(b *testing.B) {
+	const containerCount = 128
+	left := bitmapIntersectionBenchmarkFixture(containerCount, ^uint64(0))
+	right := bitmapIntersectionBenchmarkFixture(containerCount, 0x5555555555555555)
+	want := left.AndCardinality(right)
+
+	for b.Loop() {
+		result := ParAnd(0, left, right)
+		if result.GetCardinality() != want {
+			b.Fatalf("got %d values, want %d", result.GetCardinality(), want)
+		}
+	}
+}
+
 func BenchmarkIntersectionLargeParallel(b *testing.B) {
 	b.StopTimer()
 

popcnt_avx2_amd64.go

@@ -21,6 +21,9 @@ func _popcntMaskSliceAVX2(s, m []uint64) uint64
 //go:noescape
 func _popcntAndSliceAVX2(s, m []uint64) uint64
 
+//go:noescape
+func _andPopcntSliceAVX2(dst, s, m []uint64) uint64
+
 //go:noescape
 func _popcntOrSliceAVX2(s, m []uint64) uint64
 
@@ -52,6 +55,13 @@ func popcntAndSlice(s, m []uint64) uint64 {
 	return popcntAndSliceGo(s, m)
 }
 
+func andPopcntSlice(dst, s, m []uint64) uint64 {
+	if useAVX2 {
+		return _andPopcntSliceAVX2(dst, s, m)
+	}
+	return andPopcntSliceGo(dst, s, m)
+}
+
 func popcntOrSlice(s, m []uint64) uint64 {
 	if useAVX2 {
 		return _popcntOrSliceAVX2(s, m)

popcnt_avx2_amd64.s

@@ -26,9 +26,10 @@
 //     "VPAND Ymask, Ydata, Ylo" means Ylo = Ydata AND Ymask.
 //   - Yn are the 256-bit AVX registers; Xn aliases the low 128 bits of Yn.
 //   - Arguments/results are read from the frame pointer (FP). A Go slice is a
-//     3-word header {ptr,len,cap}: s_base+0(FP), s_len+8(FP); a second slice
-//     argument starts at +24(FP). The uint64 result slot follows the args
-//     (e.g. ret+24(FP) for one slice arg, ret+48(FP) for two).
+//     3-word header {ptr,len,cap}: s_base+0(FP), s_len+8(FP); subsequent slice
+//     arguments start at +24-byte intervals. The uint64 result slot follows the
+//     args (e.g. ret+24(FP) for one slice arg, ret+48(FP) for two, and
+//     ret+72(FP) for three).
 //   - Every routine is a leaf (makes no calls): NOSPLIT with a $0 local frame.
 //   - Loads/stores use VMOVDQU (unaligned): container slices are only 8-byte
 //     aligned, not 32. VZEROUPPER precedes every RET to avoid the AVX<->SSE
@@ -190,6 +191,53 @@ anddone:
 	MOVQ AX, ret+48(FP)      // +48: result follows two 24-byte slice headers
 	RET
 
+// func _andPopcntSliceAVX2(dst, s, m []uint64) uint64
+// Writes dst[i] = s[i] & m[i] and returns the resulting cardinality. The three
+// slices are assumed to have equal length. Combining the store and popcount
+// avoids rescanning dst after the non-inplace intersection is materialized.
+TEXT ·_andPopcntSliceAVX2(SB), NOSPLIT, $0-80
+	MOVQ dst_base+0(FP), BX  // BX = &dst[0]
+	MOVQ s_base+24(FP), SI   // SI = &s[0]
+	MOVQ m_base+48(FP), DI   // DI = &m[0]
+	MOVQ dst_len+8(FP), CX   // CX = len(dst)
+	XORQ AX, AX
+	SETUP
+	MOVQ CX, R8
+	SHRQ $2, R8
+	TESTQ R8, R8
+	JZ andstoretail
+andstoreloop:
+	VMOVDQU (SI), Ydata
+	VMOVDQU (DI), Yb
+	VPAND Yb, Ydata, Ydata
+	VMOVDQU Ydata, (BX)
+	COUNTBLOCK
+	ADDQ $32, BX
+	ADDQ $32, SI
+	ADDQ $32, DI
+	DECQ R8
+	JNZ andstoreloop
+	HSUM
+andstoretail:
+	ANDQ $3, CX
+	TESTQ CX, CX
+	JZ andstoredone
+andstoretailloop:
+	MOVQ (SI), DX
+	ANDQ (DI), DX
+	MOVQ DX, (BX)
+	POPCNTQ DX, DX
+	ADDQ DX, AX
+	ADDQ $8, BX
+	ADDQ $8, SI
+	ADDQ $8, DI
+	DECQ CX
+	JNZ andstoretailloop
+andstoredone:
+	VZEROUPPER
+	MOVQ AX, ret+72(FP)      // +72: result follows three 24-byte slice headers
+	RET
+
 // func _popcntOrSliceAVX2(s, m []uint64) uint64
 // Returns the sum of popcount(s[i] | m[i]); see _popcntAndSliceAVX2 for the
 // shared structure.

popcnt_avx2_amd64_test.go

@@ -97,6 +97,11 @@ func TestAVX2PopcntDispatch(t *testing.T) {
 			m := randomUint64Slice(r, n)
 			assert.Equalf(t, popcntSliceGo(s), popcntSlice(s), "popcntSlice avx2=%v len=%d", on, n)
 			assert.Equalf(t, popcntAndSliceGo(s, m), popcntAndSlice(s, m), "popcntAndSlice avx2=%v len=%d", on, n)
+			dst := make([]uint64, n)
+			assert.Equalf(t, popcntAndSliceGo(s, m), andPopcntSlice(dst, s, m), "andPopcntSlice avx2=%v len=%d", on, n)
+			for i := range dst {
+				assert.Equalf(t, s[i]&m[i], dst[i], "andPopcntSlice avx2=%v len=%d index=%d", on, n, i)
+			}
 			assert.Equalf(t, popcntOrSliceGo(s, m), popcntOrSlice(s, m), "popcntOrSlice avx2=%v len=%d", on, n)
 			assert.Equalf(t, popcntXorSliceGo(s, m), popcntXorSlice(s, m), "popcntXorSlice avx2=%v len=%d", on, n)
 			assert.Equalf(t, popcntMaskSliceGo(s, m), popcntMaskSlice(s, m), "popcntMaskSlice avx2=%v len=%d", on, n)
@@ -118,6 +123,12 @@ func TestAVX2PopcntDifferential(t *testing.T) {
 				"popcntSlice len=%d", n)
 			assert.Equalf(t, popcntAndSliceGo(s, m), _popcntAndSliceAVX2(s, m),
 				"popcntAndSlice len=%d", n)
+			dst := make([]uint64, n)
+			assert.Equalf(t, popcntAndSliceGo(s, m), _andPopcntSliceAVX2(dst, s, m),
+				"andPopcntSlice len=%d", n)
+			for i := range dst {
+				assert.Equalf(t, s[i]&m[i], dst[i], "andPopcntSlice len=%d index=%d", n, i)
+			}
 			assert.Equalf(t, popcntOrSliceGo(s, m), _popcntOrSliceAVX2(s, m),
 				"popcntOrSlice len=%d", n)
 			assert.Equalf(t, popcntXorSliceGo(s, m), _popcntXorSliceAVX2(s, m),

popcnt_generic.go

@@ -15,6 +15,10 @@ func popcntAndSlice(s, m []uint64) uint64 {
 	return popcntAndSliceGo(s, m)
 }
 
+func andPopcntSlice(dst, s, m []uint64) uint64 {
+	return andPopcntSliceGo(dst, s, m)
+}
+
 func popcntOrSlice(s, m []uint64) uint64 {
 	return popcntOrSliceGo(s, m)
 }

popcnt_neon_arm64.go

@@ -50,6 +50,10 @@ func popcntAndSlice(s, m []uint64) uint64 {
 	return popcntAndSliceGo(s, m)
 }
 
+func andPopcntSlice(dst, s, m []uint64) uint64 {
+	return andPopcntSliceGo(dst, s, m)
+}
+
 func popcntOrSlice(s, m []uint64) uint64 {
 	if useNEON {
 		return _popcntOrSliceNEON(s, m)

popcnt_slices.go

@@ -26,6 +26,15 @@ func popcntAndSliceGo(s, m []uint64) uint64 {
 	return cnt
 }
 
+// andPopcntSliceGo writes s & m to dst and returns its cardinality. All slices
+// must have the same length.
+func andPopcntSliceGo(dst, s, m []uint64) uint64 {
+	for i := range dst {
+		dst[i] = s[i] & m[i]
+	}
+	return popcntSlice(dst)
+}
+
 func popcntOrSliceGo(s, m []uint64) uint64 {
 	cnt := uint64(0)
 	for i := range s {