Vectorize dense bitmap decoding in ToArray with AVX-512 VBMI2

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++

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])
+			}
+		}
+	}
+}

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)
+}

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

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

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 (

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=