Accelerate dense bitmap decoding in ToArray

bitmapcontainer.go

@@ -294,13 +294,22 @@ func bitmapEquals(a, b []uint64) bool {
 	return true
 }
 
+// The vector path's fixed per-word setup is worthwhile above this density;
+// sparser bitmap containers keep the scalar decoder.
+const bitmapContainerAVX512MinCardinality = 1 << 14
+
 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 useAVX512 && bc.cardinality >= bitmapContainerAVX512MinCardinality {
+		return fillLeastSignificant16bitsAVX512(bc.bitmap, x, i, mask)
+	}
+	return fillLeastSignificant16bitsScalar(bc.bitmap, x, i, mask)
+}
+
+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_bench_test.go

@@ -1,6 +1,7 @@
 package roaring
 
 import (
+	"math/bits"
 	"math/rand"
 	"testing"
 )
@@ -25,6 +26,56 @@ func BenchmarkBitmapContainerFillLeastSignificant16bits(b *testing.B) {
 	}
 }
 
+// BenchmarkBitmapToArrayDenseContainers measures ToArray across four dense
+// bitmap containers and validates its complete ordered output before timing.
+func BenchmarkBitmapToArrayDenseContainers(b *testing.B) {
+	const containerCount = 4
+	words := make([]uint64, bitmapContainerSize*containerCount)
+	state := uint64(1)
+	for i := range words {
+		state += 0x9e3779b97f4a7c15
+		word := state
+		word = (word ^ (word >> 30)) * 0xbf58476d1ce4e5b9
+		word = (word ^ (word >> 27)) * 0x94d049bb133111eb
+		words[i] = word ^ (word >> 31)
+	}
+
+	bitmap := FromDense(words, false)
+	if len(bitmap.highlowcontainer.containers) != containerCount {
+		b.Fatalf("workload has %d containers, want %d", len(bitmap.highlowcontainer.containers), containerCount)
+	}
+	for _, c := range bitmap.highlowcontainer.containers {
+		if _, ok := c.(*bitmapContainer); !ok {
+			b.Fatal("workload did not produce bitmap containers")
+		}
+	}
+
+	expected := make([]uint32, 0, len(words)*32)
+	for i, word := range words {
+		base := uint32(i/bitmapContainerSize)<<16 + uint32(i%bitmapContainerSize*64)
+		for word != 0 {
+			expected = append(expected, base+uint32(bits.TrailingZeros64(word)))
+			word &= word - 1
+		}
+	}
+	actual := bitmap.ToArray()
+	if len(actual) != len(expected) {
+		b.Fatalf("validation returned %d values, want %d", len(actual), len(expected))
+	}
+	for i := range expected {
+		if actual[i] != expected[i] {
+			b.Fatalf("validation mismatch at %d: got %d, want %d", i, actual[i], expected[i])
+		}
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for b.Loop() {
+		result := bitmap.ToArray()
+		sink += result[len(result)-1]
+	}
+}
+
 // BenchmarkParOrBitmapContainers measures ParOr across four inputs with dense
 // bitmap containers at 64 shared keys using four workers.
 func BenchmarkParOrBitmapContainers(b *testing.B) {

fillbits_avx512_amd64.go

@@ -0,0 +1,15 @@
+//go:build amd64 && !appengine
+// +build amd64,!appengine
+
+package roaring
+
+//go:noescape
+func fillLeastSignificant16bitsAVX512(bitmap []uint64, x []uint32, pos int, mask uint32) int
+
+//go:noescape
+func _hasAVX512() bool
+
+// AVX-512 is used only when the operating system saves the ZMM state and the
+// CPU exposes AVX-512 Foundation. The scalar implementation remains available
+// for CPUs without that feature and for appengine builds.
+var useAVX512 = _hasAVX512()

fillbits_avx512_amd64.s

@@ -0,0 +1,127 @@
+//go:build amd64 && !appengine
+// +build amd64,!appengine
+
+#include "textflag.h"
+
+// The vector contains the 16 possible bit positions in one 16-bit input
+// block. VPCOMPRESSD writes only the lanes selected by the input mask, so its
+// store is exactly the variable-sized output for that block.
+DATA fillLeastSignificant16bitsOffsets<>+0(SB)/4, $0
+DATA fillLeastSignificant16bitsOffsets<>+4(SB)/4, $1
+DATA fillLeastSignificant16bitsOffsets<>+8(SB)/4, $2
+DATA fillLeastSignificant16bitsOffsets<>+12(SB)/4, $3
+DATA fillLeastSignificant16bitsOffsets<>+16(SB)/4, $4
+DATA fillLeastSignificant16bitsOffsets<>+20(SB)/4, $5
+DATA fillLeastSignificant16bitsOffsets<>+24(SB)/4, $6
+DATA fillLeastSignificant16bitsOffsets<>+28(SB)/4, $7
+DATA fillLeastSignificant16bitsOffsets<>+32(SB)/4, $8
+DATA fillLeastSignificant16bitsOffsets<>+36(SB)/4, $9
+DATA fillLeastSignificant16bitsOffsets<>+40(SB)/4, $10
+DATA fillLeastSignificant16bitsOffsets<>+44(SB)/4, $11
+DATA fillLeastSignificant16bitsOffsets<>+48(SB)/4, $12
+DATA fillLeastSignificant16bitsOffsets<>+52(SB)/4, $13
+DATA fillLeastSignificant16bitsOffsets<>+56(SB)/4, $14
+DATA fillLeastSignificant16bitsOffsets<>+60(SB)/4, $15
+GLOBL fillLeastSignificant16bitsOffsets<>(SB), RODATA|NOPTR, $64
+
+// func fillLeastSignificant16bitsAVX512(bitmap []uint64, x []uint32, pos int, mask uint32) int
+TEXT ·fillLeastSignificant16bitsAVX512(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 fillLeastSignificant16bitsOffsets<>(SB), Z0
+
+	TESTQ CX, CX
+	JZ fillLeastSignificant16bitsAVX512Done
+fillLeastSignificant16bitsAVX512Loop:
+	MOVQ (SI), AX
+
+	MOVWQZX AX, R11
+	KMOVW R11, K1
+	VPBROADCASTD R9, Z1
+	VPADDD Z1, Z0, Z2
+	VPCOMPRESSD Z2, K1, (DI)
+	POPCNTQ R11, R11
+	ADDQ R11, R8
+	SHLQ $2, R11
+	ADDQ R11, DI
+	ADDQ $16, R9
+	SHRQ $16, AX
+
+	MOVWQZX AX, R11
+	KMOVW R11, K1
+	VPBROADCASTD R9, Z1
+	VPADDD Z1, Z0, Z2
+	VPCOMPRESSD Z2, K1, (DI)
+	POPCNTQ R11, R11
+	ADDQ R11, R8
+	SHLQ $2, R11
+	ADDQ R11, DI
+	ADDQ $16, R9
+	SHRQ $16, AX
+
+	MOVWQZX AX, R11
+	KMOVW R11, K1
+	VPBROADCASTD R9, Z1
+	VPADDD Z1, Z0, Z2
+	VPCOMPRESSD Z2, K1, (DI)
+	POPCNTQ R11, R11
+	ADDQ R11, R8
+	SHLQ $2, R11
+	ADDQ R11, DI
+	ADDQ $16, R9
+	SHRQ $16, AX
+
+	MOVWQZX AX, R11
+	KMOVW R11, K1
+	VPBROADCASTD R9, Z1
+	VPADDD Z1, Z0, Z2
+	VPCOMPRESSD Z2, K1, (DI)
+	POPCNTQ R11, R11
+	ADDQ R11, R8
+	SHLQ $2, R11
+	ADDQ R11, DI
+	ADDQ $16, R9
+
+	ADDQ $8, SI
+	DECQ CX
+	JNZ fillLeastSignificant16bitsAVX512Loop
+fillLeastSignificant16bitsAVX512Done:
+	VZEROUPPER
+	MOVQ R8, ret+64(FP)
+	RET
+
+// func _hasAVX512() bool
+// Reports whether AVX-512 Foundation is available and the OS has enabled the
+// complete state needed by ZMM and opmask registers. The check is done once at
+// package initialization, so unsupported machines never execute the vector
+// decoder.
+TEXT ·_hasAVX512(SB), NOSPLIT, $0-1
+	// CPUID leaf 1: require OSXSAVE (ECX bit 27) and AVX (ECX bit 28).
+	MOVL $1, AX
+	XORL CX, CX
+	CPUID
+	NOTL CX
+	TESTL $0x18000000, CX
+	JNE noavx512
+
+	// XCR0: bits 1, 2, 5, 6, and 7 enable the SSE, YMM, opmask, ZMM_hi256,
+	// and Hi16_ZMM state that AVX-512 instructions use.
+	XORL CX, CX
+	XGETBV
+	NOTL AX
+	TESTL $0xe6, AX
+	JNE noavx512
+
+	// CPUID leaf 7, sub-leaf 0: AVX-512 Foundation is EBX bit 16.
+	MOVL $7, AX
+	XORL CX, CX
+	CPUID
+	NOTL BX
+	TESTL $0x10000, BX
+noavx512:
+	SETEQ ret+0(FP)
+	RET

fillbits_generic.go

@@ -0,0 +1,10 @@
+//go:build !amd64 || appengine
+// +build !amd64 appengine
+
+package roaring
+
+const useAVX512 = false
+
+func fillLeastSignificant16bitsAVX512(bitmap []uint64, x []uint32, pos int, mask uint32) int {
+	return fillLeastSignificant16bitsScalar(bitmap, x, pos, mask)
+}