Use AVX2 for bitmap AND fallback stores

Author: perfloop-agent

bitsetops_avx512_amd64.go

@@ -30,8 +30,10 @@ func andCardSliceAVX512(dst, a, b []uint64) uint64
 
 // useAVX512BitsetOps requires AVX512_VPOPCNTDQ because the fused kernels use
 // VPOPCNTQ; the plain writes only need AVX512F, but they are gated together so
-// a single flag governs the whole file. x/sys/cpu verifies operating-system
-// support for the ZMM state and honors GODEBUG=cpu.avx512vpopcntdq=off.
+// a single flag governs the whole file. When AVX-512 is unavailable, the AND
+// paths use the AVX2 store helpers before falling back to Go. x/sys/cpu verifies
+// operating-system support for the ZMM state and honors
+// GODEBUG=cpu.avx512vpopcntdq=off.
 var useAVX512BitsetOps = cpu.X86.HasAVX512VPOPCNTDQ
 
 func orSlice(dst, a, b []uint64) {
@@ -47,6 +49,10 @@ func andSlice(dst, a, b []uint64) {
 		andSliceAVX512(dst, a, b)
 		return
 	}
+	if useAVX2 {
+		_andStoreSliceAVX2(dst, a, b)
+		return
+	}
 	andSliceGo(dst, a, b)
 }
 
@@ -80,6 +86,9 @@ func andCardSlice(dst, a, b []uint64) uint64 {
 	if useAVX512BitsetOps {
 		return andCardSliceAVX512(dst, a, b)
 	}
+	if useAVX2 {
+		return _andCardStoreSliceAVX2(dst, a, b)
+	}
 	andSliceGo(dst, a, b)
 	return popcntSlice(dst)
 }

popcnt_avx2_amd64.go

@@ -7,7 +7,8 @@ package roaring
 // They are only used when the CPU supports AVX2 (see useAVX2); otherwise the
 // pure-Go fallbacks in popcnt_slices.go are used. This keeps behavior identical
 // on every target: appengine and non-amd64 builds compile popcnt_generic.go
-// instead, and amd64 CPUs without AVX2 take the scalar path at runtime.
+// instead, and amd64 CPUs without AVX2 take the scalar path at runtime. The
+// AND store helpers are also selected by the bitmap word-operation dispatch.
 
 //go:noescape
 func _hasAVX2() bool
@@ -21,6 +22,12 @@ func _popcntMaskSliceAVX2(s, m []uint64) uint64
 //go:noescape
 func _popcntAndSliceAVX2(s, m []uint64) uint64
 
+//go:noescape
+func _andStoreSliceAVX2(dst, a, b []uint64)
+
+//go:noescape
+func _andCardStoreSliceAVX2(dst, a, b []uint64) uint64
+
 //go:noescape
 func _popcntOrSliceAVX2(s, m []uint64) uint64
 

popcnt_avx2_amd64.s

@@ -358,3 +358,86 @@ TEXT ·_hasAVX2(SB), NOSPLIT, $0-1
 noavx2:
 	SETEQ ret+0(FP)          // ZF is still set by whichever TESTL ran last
 	RET
+
+// func _andStoreSliceAVX2(dst, a, b []uint64)
+// Writes a[i] & b[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. dst may alias either input.
+TEXT ·_andStoreSliceAVX2(SB), NOSPLIT, $0-72
+	MOVQ dst_base+0(FP), DI
+	MOVQ a_base+24(FP), SI
+	MOVQ b_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
+
+// func _andCardStoreSliceAVX2(dst, a, b []uint64) uint64
+// Writes a[i] & b[i] to dst[i] and returns the population count of the
+// result. Four uint64 words are processed per AVX2 iteration; a scalar tail
+// handles any remaining words. dst may alias either input.
+TEXT ·_andCardStoreSliceAVX2(SB), NOSPLIT, $0-80
+	MOVQ dst_base+0(FP), R10
+	MOVQ a_base+24(FP), SI
+	MOVQ b_base+48(FP), DI
+	MOVQ dst_len+8(FP), CX
+	XORL AX, AX
+
+	MOVQ CX, R8
+	SHRQ $2, R8
+	JZ andcardstoretail
+	SETUP
+andcardstoreloop:
+	VMOVDQU (SI), Ydata
+	VPAND (DI), Ydata, Ydata
+	VMOVDQU Ydata, (R10)
+	COUNTBLOCK
+	ADDQ $32, R10
+	ADDQ $32, SI
+	ADDQ $32, DI
+	DECQ R8
+	JNZ andcardstoreloop
+	HSUM
+andcardstoretail:
+	ANDL $3, CX
+	JZ andcardstoredone
+andcardstoretailloop:
+	MOVQ (SI), DX
+	ANDQ (DI), DX
+	MOVQ DX, (R10)
+	POPCNTQ DX, DX
+	ADDQ DX, AX
+	ADDQ $8, R10
+	ADDQ $8, SI
+	ADDQ $8, DI
+	DECL CX
+	JNZ andcardstoretailloop
+andcardstoredone:
+	VZEROUPPER
+	MOVQ AX, ret+72(FP)
+	RET

popcnt_avx2_amd64_test.go

@@ -127,3 +127,80 @@ func TestAVX2PopcntDifferential(t *testing.T) {
 		}
 	}
 }
+
+func TestAVX2AndStoreDifferential(t *testing.T) {
+	if !useAVX2 {
+		t.Skip("AVX2 not available on this CPU")
+	}
+	r := rand.New(rand.NewSource(43))
+	for _, n := range avx2TestLengths {
+		for iter := 0; iter < 64; iter++ {
+			a := randomUint64Slice(r, n)
+			b := randomUint64Slice(r, n)
+			want := make([]uint64, n)
+			andSliceGo(want, a, b)
+
+			got := make([]uint64, n)
+			_andStoreSliceAVX2(got, a, b)
+			assert.Equalf(t, want, got, "separate destination len=%d", n)
+
+			aliasA := append(make([]uint64, 0, n), a...)
+			_andStoreSliceAVX2(aliasA, aliasA, b)
+			assert.Equalf(t, want, aliasA, "destination aliases a len=%d", n)
+
+			aliasB := append(make([]uint64, 0, n), b...)
+			_andStoreSliceAVX2(aliasB, a, aliasB)
+			assert.Equalf(t, want, aliasB, "destination aliases b len=%d", n)
+
+			wantCard := andCardSliceGo(want, a, b)
+			gotCard := make([]uint64, n)
+			card := _andCardStoreSliceAVX2(gotCard, a, b)
+			assert.Equalf(t, want, gotCard, "fused destination len=%d", n)
+			assert.Equalf(t, wantCard, card, "fused cardinality len=%d", n)
+
+			aliasA = append(make([]uint64, 0, n), a...)
+			card = _andCardStoreSliceAVX2(aliasA, aliasA, b)
+			assert.Equalf(t, want, aliasA, "fused destination aliases a len=%d", n)
+			assert.Equalf(t, wantCard, card, "fused aliases a cardinality len=%d", n)
+
+			aliasB = append(make([]uint64, 0, n), b...)
+			card = _andCardStoreSliceAVX2(aliasB, a, aliasB)
+			assert.Equalf(t, want, aliasB, "fused destination aliases b len=%d", n)
+			assert.Equalf(t, wantCard, card, "fused aliases b cardinality len=%d", n)
+		}
+	}
+}
+
+func TestAVX2AndStoreDispatch(t *testing.T) {
+	if !useAVX2 {
+		t.Skip("AVX2 not available on this CPU")
+	}
+	savedAVX512 := useAVX512BitsetOps
+	savedAVX512Popcnt := useAVX512Popcnt
+	savedAVX2 := useAVX2
+	defer func() {
+		useAVX512BitsetOps = savedAVX512
+		useAVX512Popcnt = savedAVX512Popcnt
+		useAVX2 = savedAVX2
+	}()
+	useAVX512BitsetOps = false
+	useAVX512Popcnt = false
+	useAVX2 = true
+
+	r := rand.New(rand.NewSource(44))
+	for _, n := range avx2TestLengths {
+		a := randomUint64Slice(r, n)
+		b := randomUint64Slice(r, n)
+		want := make([]uint64, n)
+		andSliceGo(want, a, b)
+		got := make([]uint64, n)
+		andSlice(got, a, b)
+		assert.Equalf(t, want, got, "andSlice AVX2 dispatch len=%d", n)
+
+		gotCard := make([]uint64, n)
+		wantCard := andCardSliceGo(want, a, b)
+		card := andCardSlice(gotCard, a, b)
+		assert.Equalf(t, want, gotCard, "andCardSlice AVX2 dispatch len=%d", n)
+		assert.Equalf(t, wantCard, card, "andCardSlice AVX2 cardinality len=%d", n)
+	}
+}