diff --git a/go.mod b/go.mod index 34ab80ef..89089272 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/google/uuid v1.6.0 github.com/mschoch/smat v0.2.0 github.com/stretchr/testify v1.11.1 + golang.org/x/sys v0.30.0 ) require ( diff --git a/go.sum b/go.sum index 4666bde7..a5482567 100644 --- a/go.sum +++ b/go.sum @@ -18,6 +18,8 @@ golang.org/x/mod v0.23.0 h1:Zb7khfcRGKk+kqfxFaP5tZqCnDZMjC5VtUBs87Hr6QM= golang.org/x/mod v0.23.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY= golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w= golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= +golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= golang.org/x/tools v0.30.0 h1:BgcpHewrV5AUp2G9MebG4XPFI1E2W41zU1SaqVA9vJY= diff --git a/popcnt_avx2_amd64.go b/popcnt_avx2_amd64.go index 25b6065c..7594c3f4 100644 --- a/popcnt_avx2_amd64.go +++ b/popcnt_avx2_amd64.go @@ -32,6 +32,9 @@ func _popcntXorSliceAVX2(s, m []uint64) uint64 var useAVX2 = _hasAVX2() func popcntSlice(s []uint64) uint64 { + if useAVX512Popcnt { + return popcntSliceAVX512(s) + } if useAVX2 { return _popcntSliceAVX2(s) } @@ -39,6 +42,9 @@ func popcntSlice(s []uint64) uint64 { } func popcntMaskSlice(s, m []uint64) uint64 { + if useAVX512Popcnt { + return popcntMaskSliceAVX512(s, m) + } if useAVX2 { return _popcntMaskSliceAVX2(s, m) } @@ -46,6 +52,9 @@ func popcntMaskSlice(s, m []uint64) uint64 { } func popcntAndSlice(s, m []uint64) uint64 { + if useAVX512Popcnt { + return popcntAndSliceAVX512(s, m) + } if useAVX2 { return _popcntAndSliceAVX2(s, m) } @@ -53,6 +62,9 @@ func popcntAndSlice(s, m []uint64) uint64 { } func popcntOrSlice(s, m []uint64) uint64 { + if useAVX512Popcnt { + return popcntOrSliceAVX512(s, m) + } if useAVX2 { return _popcntOrSliceAVX2(s, m) } @@ -60,6 +72,9 @@ func popcntOrSlice(s, m []uint64) uint64 { } func popcntXorSlice(s, m []uint64) uint64 { + if useAVX512Popcnt { + return popcntXorSliceAVX512(s, m) + } if useAVX2 { return _popcntXorSliceAVX2(s, m) } diff --git a/popcnt_avx512_amd64.go b/popcnt_avx512_amd64.go new file mode 100644 index 00000000..9b9882a1 --- /dev/null +++ b/popcnt_avx512_amd64.go @@ -0,0 +1,31 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +package roaring + +import "golang.org/x/sys/cpu" + +// The functions below are implemented in popcnt_avx512_amd64.s using +// AVX512_VPOPCNTDQ. They are used in preference to the AVX2 kernels when the +// CPU provides VPOPCNTQ (see useAVX512Popcnt). + +//go:noescape +func popcntSliceAVX512(s []uint64) uint64 + +//go:noescape +func popcntMaskSliceAVX512(s, m []uint64) uint64 + +//go:noescape +func popcntAndSliceAVX512(s, m []uint64) uint64 + +//go:noescape +func popcntOrSliceAVX512(s, m []uint64) uint64 + +//go:noescape +func popcntXorSliceAVX512(s, m []uint64) uint64 + +// useAVX512Popcnt selects the AVX-512 implementations when the running CPU has +// AVX512_VPOPCNTDQ. x/sys/cpu verifies that the operating system saves the +// opmask and ZMM state, and honors GODEBUG=cpu.avx512vpopcntdq=off so the path +// can be disabled at run time. Evaluated once at package initialization. +var useAVX512Popcnt = cpu.X86.HasAVX512VPOPCNTDQ diff --git a/popcnt_avx512_amd64.s b/popcnt_avx512_amd64.s new file mode 100644 index 00000000..6e287598 --- /dev/null +++ b/popcnt_avx512_amd64.s @@ -0,0 +1,311 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +#include "textflag.h" + +// AVX-512 population-count routines for amd64, counterparts to the AVX2 ones +// in popcnt_avx2_amd64.s. They run when the CPU has AVX512_VPOPCNTDQ (see +// useAVX512Popcnt); otherwise the AVX2 or scalar paths are used instead. +// +// Algorithm +// --------- +// AVX2 has no vector population count, so the AVX2 routines build one out of a +// VPSHUFB nibble lookup plus VPSADBW and handle 256 bits per iteration. +// AVX512_VPOPCNTDQ provides VPOPCNTQ, which counts all eight 64-bit lanes of a +// ZMM register in one instruction, so each iteration here handles 2048 bits: +// four ZMM loads, four VPOPCNTQ, four adds. +// +// Four accumulators (Z4..Z7) keep the adds off a single dependency chain; they +// are folded together after the loop and reduced to a scalar by HSUM512. A +// scalar POPCNTQ tail handles the trailing len%32 words, so any slice length is +// counted correctly. +// +// Go assembler conventions are as in popcnt_avx2_amd64.s: operands are written +// source(s) first and destination last, Yn/Xn alias the low halves of Zn, and a +// []uint64 argument is a {ptr,len,cap} header, so a second slice argument +// starts at +24(FP) and the uint64 result follows the arguments. + +// HSUM512 horizontally sums the eight 64-bit lanes of Z4 into out. +#define HSUM512(out) \ + VEXTRACTI64X4 $1, Z4, Y1 \ + VPADDQ Y1, Y4, Y1 \ + VEXTRACTI128 $1, Y1, X2 \ + VPADDQ X2, X1, X1 \ + VPSHUFD $0x4e, X1, X2 \ + VPADDQ X2, X1, X1 \ + VMOVQ X1, out + +// ZEROACC clears the four accumulators. +#define ZEROACC \ + VPXORQ Z4, Z4, Z4 \ + VPXORQ Z5, Z5, Z5 \ + VPXORQ Z6, Z6, Z6 \ + VPXORQ Z7, Z7, Z7 + +// FOLDACC sums the four accumulators into Z4. +#define FOLDACC \ + VPADDQ Z5, Z4, Z4 \ + VPADDQ Z7, Z6, Z6 \ + VPADDQ Z6, Z4, Z4 + +// func popcntSliceAVX512(s []uint64) uint64 +TEXT ·popcntSliceAVX512(SB), NOSPLIT, $0-32 + MOVQ s_base+0(FP), SI + MOVQ s_len+8(FP), DX + ZEROACC + MOVQ DX, CX + SHRQ $5, CX + TESTQ CX, CX + JZ ps_tail + +ps_loop: + VPOPCNTQ (SI), Z0 + VPOPCNTQ 64(SI), Z1 + VPOPCNTQ 128(SI), Z2 + VPOPCNTQ 192(SI), Z3 + VPADDQ Z0, Z4, Z4 + VPADDQ Z1, Z5, Z5 + VPADDQ Z2, Z6, Z6 + VPADDQ Z3, Z7, Z7 + ADDQ $256, SI + DECQ CX + JNZ ps_loop + +ps_tail: + FOLDACC + HSUM512(AX) + ANDQ $31, DX + JZ ps_done + +ps_scalar: + POPCNTQ (SI), BX + ADDQ BX, AX + ADDQ $8, SI + DECQ DX + JNZ ps_scalar + +ps_done: + VZEROUPPER + MOVQ AX, ret+24(FP) + RET + +// func popcntAndSliceAVX512(s, m []uint64) uint64 +TEXT ·popcntAndSliceAVX512(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ s_len+8(FP), DX + MOVQ m_base+24(FP), DI + ZEROACC + MOVQ DX, CX + SHRQ $5, CX + TESTQ CX, CX + JZ pa_tail + +pa_loop: + VMOVDQU64 (SI), Z0 + VMOVDQU64 64(SI), Z1 + VMOVDQU64 128(SI), Z2 + VMOVDQU64 192(SI), Z3 + VPANDQ (DI), Z0, Z0 + VPANDQ 64(DI), Z1, Z1 + VPANDQ 128(DI), Z2, Z2 + VPANDQ 192(DI), Z3, Z3 + VPOPCNTQ Z0, Z0 + VPOPCNTQ Z1, Z1 + VPOPCNTQ Z2, Z2 + VPOPCNTQ Z3, Z3 + VPADDQ Z0, Z4, Z4 + VPADDQ Z1, Z5, Z5 + VPADDQ Z2, Z6, Z6 + VPADDQ Z3, Z7, Z7 + ADDQ $256, SI + ADDQ $256, DI + DECQ CX + JNZ pa_loop + +pa_tail: + FOLDACC + HSUM512(AX) + ANDQ $31, DX + JZ pa_done + +pa_scalar: + MOVQ (SI), BX + ANDQ (DI), BX + POPCNTQ BX, BX + ADDQ BX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ DX + JNZ pa_scalar + +pa_done: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET + +// func popcntOrSliceAVX512(s, m []uint64) uint64 +TEXT ·popcntOrSliceAVX512(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ s_len+8(FP), DX + MOVQ m_base+24(FP), DI + ZEROACC + MOVQ DX, CX + SHRQ $5, CX + TESTQ CX, CX + JZ po_tail + +po_loop: + VMOVDQU64 (SI), Z0 + VMOVDQU64 64(SI), Z1 + VMOVDQU64 128(SI), Z2 + VMOVDQU64 192(SI), Z3 + VPORQ (DI), Z0, Z0 + VPORQ 64(DI), Z1, Z1 + VPORQ 128(DI), Z2, Z2 + VPORQ 192(DI), Z3, Z3 + VPOPCNTQ Z0, Z0 + VPOPCNTQ Z1, Z1 + VPOPCNTQ Z2, Z2 + VPOPCNTQ Z3, Z3 + VPADDQ Z0, Z4, Z4 + VPADDQ Z1, Z5, Z5 + VPADDQ Z2, Z6, Z6 + VPADDQ Z3, Z7, Z7 + ADDQ $256, SI + ADDQ $256, DI + DECQ CX + JNZ po_loop + +po_tail: + FOLDACC + HSUM512(AX) + ANDQ $31, DX + JZ po_done + +po_scalar: + MOVQ (SI), BX + ORQ (DI), BX + POPCNTQ BX, BX + ADDQ BX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ DX + JNZ po_scalar + +po_done: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET + +// func popcntXorSliceAVX512(s, m []uint64) uint64 +TEXT ·popcntXorSliceAVX512(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ s_len+8(FP), DX + MOVQ m_base+24(FP), DI + ZEROACC + MOVQ DX, CX + SHRQ $5, CX + TESTQ CX, CX + JZ px_tail + +px_loop: + VMOVDQU64 (SI), Z0 + VMOVDQU64 64(SI), Z1 + VMOVDQU64 128(SI), Z2 + VMOVDQU64 192(SI), Z3 + VPXORQ (DI), Z0, Z0 + VPXORQ 64(DI), Z1, Z1 + VPXORQ 128(DI), Z2, Z2 + VPXORQ 192(DI), Z3, Z3 + VPOPCNTQ Z0, Z0 + VPOPCNTQ Z1, Z1 + VPOPCNTQ Z2, Z2 + VPOPCNTQ Z3, Z3 + VPADDQ Z0, Z4, Z4 + VPADDQ Z1, Z5, Z5 + VPADDQ Z2, Z6, Z6 + VPADDQ Z3, Z7, Z7 + ADDQ $256, SI + ADDQ $256, DI + DECQ CX + JNZ px_loop + +px_tail: + FOLDACC + HSUM512(AX) + ANDQ $31, DX + JZ px_done + +px_scalar: + MOVQ (SI), BX + XORQ (DI), BX + POPCNTQ BX, BX + ADDQ BX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ DX + JNZ px_scalar + +px_done: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET + +// func popcntMaskSliceAVX512(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] &^ m[i]). As in the AVX2 routine, VPANDN +// negates its first source and only the second may come from memory, so m goes +// into the register and s is read from memory: "VPANDNQ (SI), Zm, Zm" gives +// (NOT m) AND s. +TEXT ·popcntMaskSliceAVX512(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ s_len+8(FP), DX + MOVQ m_base+24(FP), DI + ZEROACC + MOVQ DX, CX + SHRQ $5, CX + TESTQ CX, CX + JZ pm_tail + +pm_loop: + VMOVDQU64 (DI), Z0 + VMOVDQU64 64(DI), Z1 + VMOVDQU64 128(DI), Z2 + VMOVDQU64 192(DI), Z3 + VPANDNQ (SI), Z0, Z0 + VPANDNQ 64(SI), Z1, Z1 + VPANDNQ 128(SI), Z2, Z2 + VPANDNQ 192(SI), Z3, Z3 + VPOPCNTQ Z0, Z0 + VPOPCNTQ Z1, Z1 + VPOPCNTQ Z2, Z2 + VPOPCNTQ Z3, Z3 + VPADDQ Z0, Z4, Z4 + VPADDQ Z1, Z5, Z5 + VPADDQ Z2, Z6, Z6 + VPADDQ Z3, Z7, Z7 + ADDQ $256, SI + ADDQ $256, DI + DECQ CX + JNZ pm_loop + +pm_tail: + FOLDACC + HSUM512(AX) + ANDQ $31, DX + JZ pm_done + +pm_scalar: + MOVQ (DI), BX + NOTQ BX + ANDQ (SI), BX + POPCNTQ BX, BX + ADDQ BX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ DX + JNZ pm_scalar + +pm_done: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET diff --git a/popcnt_avx512_amd64_test.go b/popcnt_avx512_amd64_test.go new file mode 100644 index 00000000..a04e9976 --- /dev/null +++ b/popcnt_avx512_amd64_test.go @@ -0,0 +1,77 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +package roaring + +import ( + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +// edge lengths exercise the AVX-512 main loop (multiples of 32) and the scalar +// POPCNTQ tail (len % 32 != 0), including the empty and sub-block cases. +var avx512TestLengths = []int{0, 1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 32, 33, 63, 64, 65, 1023, 1024, 1025} + +func TestPopcntAVX512MatchesGo(t *testing.T) { + if !useAVX512Popcnt { + t.Skip("AVX512_VPOPCNTDQ not available") + } + r := rand.New(rand.NewSource(42)) + for _, n := range avx512TestLengths { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + + assert.Equalf(t, popcntSliceGo(s), popcntSliceAVX512(s), "popcntSlice len=%d", n) + assert.Equalf(t, popcntMaskSliceGo(s, m), popcntMaskSliceAVX512(s, m), "popcntMaskSlice len=%d", n) + assert.Equalf(t, popcntAndSliceGo(s, m), popcntAndSliceAVX512(s, m), "popcntAndSlice len=%d", n) + assert.Equalf(t, popcntOrSliceGo(s, m), popcntOrSliceAVX512(s, m), "popcntOrSlice len=%d", n) + assert.Equalf(t, popcntXorSliceGo(s, m), popcntXorSliceAVX512(s, m), "popcntXorSlice len=%d", n) + } +} + +// The dispatchers must agree with the portable implementations whatever the +// running CPU supports. +func TestPopcntDispatchMatchesGo(t *testing.T) { + r := rand.New(rand.NewSource(7)) + for _, n := range avx512TestLengths { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + + assert.Equalf(t, popcntSliceGo(s), popcntSlice(s), "popcntSlice len=%d", n) + assert.Equalf(t, popcntMaskSliceGo(s, m), popcntMaskSlice(s, m), "popcntMaskSlice len=%d", n) + assert.Equalf(t, popcntAndSliceGo(s, m), popcntAndSlice(s, m), "popcntAndSlice len=%d", n) + assert.Equalf(t, popcntOrSliceGo(s, m), popcntOrSlice(s, m), "popcntOrSlice len=%d", n) + assert.Equalf(t, popcntXorSliceGo(s, m), popcntXorSlice(s, m), "popcntXorSlice len=%d", n) + } +} + +func BenchmarkPopcntSliceAVX512(b *testing.B) { + if !useAVX512Popcnt { + b.Skip("AVX512_VPOPCNTDQ not available") + } + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + b.SetBytes(int64(len(s) * 8)) + var sink uint64 + for b.Loop() { + sink = popcntSliceAVX512(s) + } + _ = sink +} + +func BenchmarkPopcntAndSliceAVX512(b *testing.B) { + if !useAVX512Popcnt { + b.Skip("AVX512_VPOPCNTDQ not available") + } + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + m := randomUint64Slice(r, 1024) + b.SetBytes(int64(len(s) * 8)) + var sink uint64 + for b.Loop() { + sink = popcntAndSliceAVX512(s, m) + } + _ = sink +}