Vectorize the bitmap-container word operations, fusing the cardinality pass

bitmapcontainer.go

@@ -562,10 +562,7 @@ func (bc *bitmapContainer) orArrayCardinality(value2 *arrayContainer) int {
 
 func (bc *bitmapContainer) orBitmap(value2 *bitmapContainer) container {
 	answer := newBitmapContainer()
-	for k := 0; k < len(answer.bitmap); k++ {
-		answer.bitmap[k] = bc.bitmap[k] | value2.bitmap[k]
-	}
-	answer.computeCardinality()
+	answer.cardinality = int(orCardSlice(answer.bitmap, bc.bitmap, value2.bitmap))
 	if answer.isFull() {
 		return newRunContainer16Range(0, MaxUint16)
 	}
@@ -601,11 +598,7 @@ func (bc *bitmapContainer) iorArray(ac *arrayContainer) container {
 
 func (bc *bitmapContainer) iorBitmap(value2 *bitmapContainer) container {
 	answer := bc
-	answer.cardinality = 0
-	for k := 0; k < len(answer.bitmap); k++ {
-		answer.bitmap[k] = bc.bitmap[k] | value2.bitmap[k]
-	}
-	answer.computeCardinality()
+	answer.cardinality = int(orCardSlice(answer.bitmap, bc.bitmap, value2.bitmap))
 	if bc.isFull() {
 		return newRunContainer16Range(0, MaxUint16)
 	}
@@ -654,16 +647,7 @@ func (bc *bitmapContainer) lazyIORBitmap(value2 *bitmapContainer) container {
 	bitmap := answer.bitmap
 	other := value2.bitmap
 
-	// Bitmap containers always span bitmapContainerSize words. Prove the
-	// bounds once so the compiler can eliminate the checks in the unrolled loop.
-	_ = bitmap[bitmapContainerSize-1]
-	_ = other[bitmapContainerSize-1]
-	for k := 0; k < bitmapContainerSize; k += 4 {
-		bitmap[k] |= other[k]
-		bitmap[k+1] |= other[k+1]
-		bitmap[k+2] |= other[k+2]
-		bitmap[k+3] |= other[k+3]
-	}
+	orSlice(bitmap, bitmap, other)
 	answer.cardinality = invalidCardinality
 	return answer
 }
@@ -674,17 +658,7 @@ func (bc *bitmapContainer) lazyORBitmap(value2 *bitmapContainer) container {
 	left := bc.bitmap
 	right := value2.bitmap
 
-	// Bitmap containers always span bitmapContainerSize words. Prove the
-	// bounds once so the compiler can eliminate the checks in the unrolled loop.
-	_ = bitmap[bitmapContainerSize-1]
-	_ = left[bitmapContainerSize-1]
-	_ = right[bitmapContainerSize-1]
-	for k := 0; k < bitmapContainerSize; k += 4 {
-		bitmap[k] = left[k] | right[k]
-		bitmap[k+1] = left[k+1] | right[k+1]
-		bitmap[k+2] = left[k+2] | right[k+2]
-		bitmap[k+3] = left[k+3] | right[k+3]
-	}
+	orSlice(bitmap, left, right)
 	answer.cardinality = invalidCardinality
 	return answer
 }
@@ -744,9 +718,7 @@ func (bc *bitmapContainer) xorBitmap(value2 *bitmapContainer) container {
 
 	if newCardinality > arrayDefaultMaxSize {
 		answer := newBitmapContainer()
-		for k := 0; k < len(answer.bitmap); k++ {
-			answer.bitmap[k] = bc.bitmap[k] ^ value2.bitmap[k]
-		}
+		xorSlice(answer.bitmap, bc.bitmap, value2.bitmap)
 		answer.cardinality = newCardinality
 		if answer.isFull() {
 			return newRunContainer16Range(0, MaxUint16)
@@ -868,9 +840,7 @@ func (bc *bitmapContainer) andBitmap(value2 *bitmapContainer) container {
 	newcardinality := int(popcntAndSlice(bc.bitmap, value2.bitmap))
 	if newcardinality > arrayDefaultMaxSize {
 		answer := newBitmapContainer()
-		for k := 0; k < len(answer.bitmap); k++ {
-			answer.bitmap[k] = bc.bitmap[k] & value2.bitmap[k]
-		}
+		andSlice(answer.bitmap, bc.bitmap, value2.bitmap)
 		answer.cardinality = newcardinality
 		return answer
 	}
@@ -901,10 +871,7 @@ func (bc *bitmapContainer) intersectsBitmap(value2 *bitmapContainer) bool {
 }
 
 func (bc *bitmapContainer) iandBitmap(value2 *bitmapContainer) container {
-	newcardinality := int(popcntAndSlice(bc.bitmap, value2.bitmap))
-	for k := 0; k < len(bc.bitmap); k++ {
-		bc.bitmap[k] = bc.bitmap[k] & value2.bitmap[k]
-	}
+	newcardinality := int(andCardSlice(bc.bitmap, bc.bitmap, value2.bitmap))
 	bc.cardinality = newcardinality
 
 	if newcardinality <= arrayDefaultMaxSize {
@@ -938,9 +905,7 @@ func (bc *bitmapContainer) ixorRun16(value2 *runContainer16) container {
 func (bc *bitmapContainer) ixorBitmap(value2 *bitmapContainer) container {
 	newCardinality := int(popcntXorSlice(bc.bitmap, value2.bitmap))
 	if newCardinality > arrayDefaultMaxSize {
-		for k := 0; k < len(bc.bitmap); k++ {
-			bc.bitmap[k] = bc.bitmap[k] ^ value2.bitmap[k]
-		}
+		xorSlice(bc.bitmap, bc.bitmap, value2.bitmap)
 		bc.cardinality = newCardinality
 		return bc
 	}
@@ -1064,9 +1029,7 @@ func (bc *bitmapContainer) andNotBitmap(value2 *bitmapContainer) container {
 	newCardinality := int(popcntMaskSlice(bc.bitmap, value2.bitmap))
 	if newCardinality > arrayDefaultMaxSize {
 		answer := newBitmapContainer()
-		for k := 0; k < len(answer.bitmap); k++ {
-			answer.bitmap[k] = bc.bitmap[k] &^ value2.bitmap[k]
-		}
+		andNotSlice(answer.bitmap, bc.bitmap, value2.bitmap)
 		answer.cardinality = newCardinality
 		return answer
 	}
@@ -1077,9 +1040,7 @@ func (bc *bitmapContainer) andNotBitmap(value2 *bitmapContainer) container {
 
 func (bc *bitmapContainer) iandNotBitmapSurely(value2 *bitmapContainer) container {
 	newCardinality := int(popcntMaskSlice(bc.bitmap, value2.bitmap))
-	for k := 0; k < len(bc.bitmap); k++ {
-		bc.bitmap[k] = bc.bitmap[k] &^ value2.bitmap[k]
-	}
+	andNotSlice(bc.bitmap, bc.bitmap, value2.bitmap)
 	bc.cardinality = newCardinality
 	if bc.getCardinality() <= arrayDefaultMaxSize {
 		return bc.toArrayContainer()

bitsetops.go

@@ -0,0 +1,53 @@
+package roaring
+
+import "math/bits"
+
+// Portable implementations of the bitmap-container word operations. Each writes
+// dst[i] = a[i] op b[i]; the Card variants also return the population count of
+// the result, which callers would otherwise obtain with a second pass.
+//
+// dst may alias a or b: every element is read before it is written.
+
+func orSliceGo(dst, a, b []uint64) {
+	for i := range dst {
+		dst[i] = a[i] | b[i]
+	}
+}
+
+func andSliceGo(dst, a, b []uint64) {
+	for i := range dst {
+		dst[i] = a[i] & b[i]
+	}
+}
+
+func xorSliceGo(dst, a, b []uint64) {
+	for i := range dst {
+		dst[i] = a[i] ^ b[i]
+	}
+}
+
+func andNotSliceGo(dst, a, b []uint64) {
+	for i := range dst {
+		dst[i] = a[i] &^ b[i]
+	}
+}
+
+func orCardSliceGo(dst, a, b []uint64) uint64 {
+	card := 0
+	for i := range dst {
+		v := a[i] | b[i]
+		dst[i] = v
+		card += bits.OnesCount64(v)
+	}
+	return uint64(card)
+}
+
+func andCardSliceGo(dst, a, b []uint64) uint64 {
+	card := 0
+	for i := range dst {
+		v := a[i] & b[i]
+		dst[i] = v
+		card += bits.OnesCount64(v)
+	}
+	return uint64(card)
+}

bitsetops_avx512_amd64.go

@@ -0,0 +1,85 @@
+//go:build amd64 && !appengine
+// +build amd64,!appengine
+
+package roaring
+
+import "golang.org/x/sys/cpu"
+
+// The functions below are implemented in bitsetops_avx512_amd64.s. The Card
+// variants fuse the write with the population count of the result, so a
+// bitmap-container operation that needs both makes one pass over the container
+// instead of two.
+
+//go:noescape
+func orSliceAVX512(dst, a, b []uint64)
+
+//go:noescape
+func andSliceAVX512(dst, a, b []uint64)
+
+//go:noescape
+func xorSliceAVX512(dst, a, b []uint64)
+
+//go:noescape
+func andNotSliceAVX512(dst, a, b []uint64)
+
+//go:noescape
+func orCardSliceAVX512(dst, a, b []uint64) uint64
+
+//go:noescape
+func andCardSliceAVX512(dst, a, b []uint64) uint64
+
+// useAVX512BitsetOps requires AVX512_VPOPCNTDQ because the fused kernels use
+// VPOPCNTQ; the plain writes only need AVX512F, but they are gated together so
+// a single flag governs the whole file. x/sys/cpu verifies operating-system
+// support for the ZMM state and honors GODEBUG=cpu.avx512vpopcntdq=off.
+var useAVX512BitsetOps = cpu.X86.HasAVX512VPOPCNTDQ
+
+func orSlice(dst, a, b []uint64) {
+	if useAVX512BitsetOps {
+		orSliceAVX512(dst, a, b)
+		return
+	}
+	orSliceGo(dst, a, b)
+}
+
+func andSlice(dst, a, b []uint64) {
+	if useAVX512BitsetOps {
+		andSliceAVX512(dst, a, b)
+		return
+	}
+	andSliceGo(dst, a, b)
+}
+
+func xorSlice(dst, a, b []uint64) {
+	if useAVX512BitsetOps {
+		xorSliceAVX512(dst, a, b)
+		return
+	}
+	xorSliceGo(dst, a, b)
+}
+
+func andNotSlice(dst, a, b []uint64) {
+	if useAVX512BitsetOps {
+		andNotSliceAVX512(dst, a, b)
+		return
+	}
+	andNotSliceGo(dst, a, b)
+}
+
+func orCardSlice(dst, a, b []uint64) uint64 {
+	if useAVX512BitsetOps {
+		return orCardSliceAVX512(dst, a, b)
+	}
+	// Without a vector population count the fused loop is no faster than the
+	// two passes the callers used before, and popcntSlice may itself be AVX2.
+	orSliceGo(dst, a, b)
+	return popcntSlice(dst)
+}
+
+func andCardSlice(dst, a, b []uint64) uint64 {
+	if useAVX512BitsetOps {
+		return andCardSliceAVX512(dst, a, b)
+	}
+	andSliceGo(dst, a, b)
+	return popcntSlice(dst)
+}