diff --git a/bitsetops_avx512_amd64.go b/bitsetops_avx512_amd64.go index b25909cc..4351a9ab 100644 --- a/bitsetops_avx512_amd64.go +++ b/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) } diff --git a/bitsetops_test.go b/bitsetops_test.go index 8d277e22..0e0b5fe5 100644 --- a/bitsetops_test.go +++ b/bitsetops_test.go @@ -173,6 +173,13 @@ func BenchmarkDenseBitmapOps(b *testing.B) { sinkU += And(x, y).GetCardinality() } }) + b.Run("IAnd", func(b *testing.B) { + for b.Loop() { + z := x.Clone() + z.And(y) + sinkU += z.GetCardinality() + } + }) b.Run("Xor", func(b *testing.B) { for b.Loop() { sinkU += Xor(x, y).GetCardinality() diff --git a/popcnt_avx2_amd64.go b/popcnt_avx2_amd64.go index 7594c3f4..90cac24b 100644 --- a/popcnt_avx2_amd64.go +++ b/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 diff --git a/popcnt_avx2_amd64.s b/popcnt_avx2_amd64.s index 52e2c989..572d61bc 100644 --- a/popcnt_avx2_amd64.s +++ b/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 diff --git a/popcnt_avx2_amd64_test.go b/popcnt_avx2_amd64_test.go index c0e91b0f..8a75009d 100644 --- a/popcnt_avx2_amd64_test.go +++ b/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) + } +}