diff --git a/bitmapcontainer.go b/bitmapcontainer.go index 416da732..271ab846 100644 --- a/bitmapcontainer.go +++ b/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++ diff --git a/bitmapcontainer_bench_test.go b/bitmapcontainer_bench_test.go index 530772ac..48cb9a96 100644 --- a/bitmapcontainer_bench_test.go +++ b/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) { diff --git a/fillbits_avx512_amd64.go b/fillbits_avx512_amd64.go new file mode 100644 index 00000000..c9dcef52 --- /dev/null +++ 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() diff --git a/fillbits_avx512_amd64.s b/fillbits_avx512_amd64.s new file mode 100644 index 00000000..e4b5c240 --- /dev/null +++ b/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 diff --git a/fillbits_generic.go b/fillbits_generic.go new file mode 100644 index 00000000..742f4d97 --- /dev/null +++ b/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) +}