diff --git a/bitmapcontainer.go b/bitmapcontainer.go index 416da732..354af17e 100644 --- a/bitmapcontainer.go +++ b/bitmapcontainer.go @@ -294,13 +294,24 @@ func bitmapEquals(a, b []uint64) bool { return true } +// bitmapContainerVectorFillMinCardinality is where the vector decoder starts +// paying back its fixed per-word cost. Measured break-even is near 2000; 4096 +// is the array-to-bitmap conversion point and leaves a comfortable margin. +const bitmapContainerVectorFillMinCardinality = 4096 + func (bc *bitmapContainer) fillLeastSignificant16bits(x []uint32, i int, mask uint32) int { - // On amd64 this loop compiles to TZCNT/BLSR; the remaining headroom is - // vectorized decode (cf. CRoaring bitset_extract_setbits_avx2/avx512). + if useVectorFill && bc.cardinality >= bitmapContainerVectorFillMinCardinality { + return fillLeastSignificant16bitsVector(bc.bitmap, x, i, mask) + } + return fillLeastSignificant16bitsScalar(bc.bitmap, x, i, mask) +} + +// fillLeastSignificant16bitsScalar compiles to a TZCNT/BLSR loop on amd64. +func fillLeastSignificant16bitsScalar(bitmap []uint64, x []uint32, i int, mask uint32) int { pos := i base := mask - for k := 0; k < len(bc.bitmap); k++ { - bitset := bc.bitmap[k] + for k := 0; k < len(bitmap); k++ { + bitset := bitmap[k] for bitset != 0 { x[pos] = base + uint32(bits.TrailingZeros64(bitset)) pos++ diff --git a/bitmapcontainer_test.go b/bitmapcontainer_test.go index b74fb5c6..48e423dc 100644 --- a/bitmapcontainer_test.go +++ b/bitmapcontainer_test.go @@ -639,3 +639,45 @@ func TestBitmapContainerFillLeastSignificant16bitsProperties(t *testing.T) { runTest(t, vals, 0x55550000) }) } + +// TestBitmapContainerFillLeastSignificant16bitsVector checks the vector decoder +// against the scalar one, including at the cardinality where they switch over. +func TestBitmapContainerFillLeastSignificant16bitsVector(t *testing.T) { + if !useVectorFill { + t.Skip("no vector decoder on this platform") + } + + r := rand.New(rand.NewSource(42)) + for _, cardinality := range []int{ + bitmapContainerVectorFillMinCardinality - 1, + bitmapContainerVectorFillMinCardinality, + bitmapContainerVectorFillMinCardinality + 1, + 20000, 40000, 60000, 65535, 65536, + } { + bc := newBitmapContainer() + for bc.cardinality < cardinality { + bc.iadd(uint16(r.Intn(65536))) + } + + const startIdx = 3 + const sentinel = 0xDEADC0DE + want := make([]uint32, startIdx+cardinality+8) + got := make([]uint32, len(want)) + for j := range want { + want[j] = sentinel + got[j] = sentinel + } + + wantPos := fillLeastSignificant16bitsScalar(bc.bitmap, want, startIdx, 0xFFFF0000) + gotPos := bc.fillLeastSignificant16bits(got, startIdx, 0xFFFF0000) + + if gotPos != wantPos { + t.Fatalf("cardinality %d: returned %d, want %d", cardinality, gotPos, wantPos) + } + for j := range want { + if got[j] != want[j] { + t.Fatalf("cardinality %d: at %d got %#x, want %#x", cardinality, j, got[j], want[j]) + } + } + } +} diff --git a/fillbits_generic.go b/fillbits_generic.go new file mode 100644 index 00000000..d31c913e --- /dev/null +++ b/fillbits_generic.go @@ -0,0 +1,10 @@ +//go:build !amd64 || appengine +// +build !amd64 appengine + +package roaring + +const useVectorFill = false + +func fillLeastSignificant16bitsVector(bitmap []uint64, x []uint32, pos int, mask uint32) int { + return fillLeastSignificant16bitsScalar(bitmap, x, pos, mask) +} diff --git a/fillbits_vbmi2_amd64.go b/fillbits_vbmi2_amd64.go new file mode 100644 index 00000000..6d7d2538 --- /dev/null +++ b/fillbits_vbmi2_amd64.go @@ -0,0 +1,19 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +package roaring + +import "golang.org/x/sys/cpu" + +// fillLeastSignificant16bitsVector is implemented in fillbits_vbmi2_amd64.s. +// It is only called when useVectorFill reports that the CPU can run it. +// +//go:noescape +func fillLeastSignificant16bitsVector(bitmap []uint64, x []uint32, pos int, mask uint32) int + +// useVectorFill reports whether the assembly decoder may be used. It needs +// VPCOMPRESSB, that is AVX512_VBMI2 (Ice Lake and later, Zen 4 and later). +// x/sys/cpu also verifies that the operating system saves the opmask and ZMM +// state, and it honors GODEBUG=cpu.avx512vbmi2=off, so the vector path can be +// turned off at run time without rebuilding. +var useVectorFill = cpu.X86.HasAVX512VBMI2 diff --git a/fillbits_vbmi2_amd64.s b/fillbits_vbmi2_amd64.s new file mode 100644 index 00000000..397a6ec1 --- /dev/null +++ b/fillbits_vbmi2_amd64.s @@ -0,0 +1,106 @@ +//go:build amd64 && !appengine + +#include "textflag.h" + +// Byte values 0x00..0x3f: the bit positions inside one 64-bit word. +DATA bitPositions<>+0x00(SB)/8, $0x0706050403020100 +DATA bitPositions<>+0x08(SB)/8, $0x0f0e0d0c0b0a0908 +DATA bitPositions<>+0x10(SB)/8, $0x1716151413121110 +DATA bitPositions<>+0x18(SB)/8, $0x1f1e1d1c1b1a1918 +DATA bitPositions<>+0x20(SB)/8, $0x2726252423222120 +DATA bitPositions<>+0x28(SB)/8, $0x2f2e2d2c2b2a2928 +DATA bitPositions<>+0x30(SB)/8, $0x3736353433323130 +DATA bitPositions<>+0x38(SB)/8, $0x3f3e3d3c3b3a3938 +GLOBL bitPositions<>(SB), RODATA|NOPTR, $64 + +// func fillLeastSignificant16bitsVector(bitmap []uint64, x []uint32, pos int, mask uint32) int +// +// One VPCOMPRESSB per 64-bit word turns the whole word into 64 byte-sized bit +// positions in a single shot; the positions are then widened 16 at a time with +// VPMOVZXBD and added to the running base. +// +// The compress-then-widen idea is taken from simdjson's bit_indexer::write +// (icelake kernel, Apache-2.0): https://github.com/simdjson/simdjson +// +// It differs here in the stores. simdjson writes whole 64-byte blocks and +// relies on its output buffer having up to 64 uint32 of slack. ToArray +// allocates exactly GetCardinality() values, with no padding, so each block is +// written with a mask derived from BZHI(-1, popcount) instead. That keeps the +// writes exactly popcount(word) wide, and it is also faster on dense +// containers because no store bandwidth is spent on values nobody asked for. +TEXT ·fillLeastSignificant16bitsVector(SB), NOSPLIT, $0-72 + MOVQ bitmap_base+0(FP), SI + MOVQ bitmap_len+8(FP), CX + MOVQ x_base+24(FP), DI + MOVQ pos+48(FP), R8 + LEAQ (DI)(R8*4), DI + MOVL mask+56(FP), R9 + + VMOVDQU64 bitPositions<>(SB), Z0 + VPBROADCASTD R9, Z1 + MOVL $64, AX + VPBROADCASTD AX, Z3 + + TESTQ CX, CX + JZ done + +loop: + MOVQ (SI), AX + TESTQ AX, AX + JZ next + + KMOVQ AX, K1 + VPCOMPRESSB.Z Z0, K1, Z2 + POPCNTQ AX, R10 + MOVQ $-1, R11 + BZHIQ R10, R11, R11 + + KMOVW R11, K2 + VPMOVZXBD X2, Z4 + VPADDD Z1, Z4, Z4 + VMOVDQU32 Z4, K2, (DI) + + CMPQ R10, $16 + JLE advance + + SHRQ $16, R11 + KMOVW R11, K2 + VEXTRACTI32X4 $1, Z2, X5 + VPMOVZXBD X5, Z4 + VPADDD Z1, Z4, Z4 + VMOVDQU32 Z4, K2, 64(DI) + + CMPQ R10, $32 + JLE advance + + SHRQ $16, R11 + KMOVW R11, K2 + VEXTRACTI32X4 $2, Z2, X5 + VPMOVZXBD X5, Z4 + VPADDD Z1, Z4, Z4 + VMOVDQU32 Z4, K2, 128(DI) + + CMPQ R10, $48 + JLE advance + + SHRQ $16, R11 + KMOVW R11, K2 + VEXTRACTI32X4 $3, Z2, X5 + VPMOVZXBD X5, Z4 + VPADDD Z1, Z4, Z4 + VMOVDQU32 Z4, K2, 192(DI) + +advance: + LEAQ (DI)(R10*4), DI + ADDQ R10, R8 + +next: + VPADDD Z3, Z1, Z1 + ADDQ $8, SI + DECQ CX + JNZ loop + +done: + VZEROUPPER + MOVQ R8, ret+64(FP) + RET 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=