Optimize AVX2 popcount inner loop and CPU feature detection - #546
Open
lemire wants to merge 1 commit into
Open
Conversation
David Sparks reviewed popcnt_avx2_amd64.s and suggested the improvements
implemented here; the credit for all of them is his.
The main win is in COUNTBLOCK. VPSADBW computes |a-b| per byte and sums
each group of 8, so it can absorb the per-byte add that VPADDB was doing:
feeding it the two nibble counts directly yields (B + lo) - (B - hi) =
lo + hi, removing VPADDB and its latency from the hot loop. That needs
two lookup tables, one biased up by B and one subtracted from B. The bias
must satisfy 4 <= B <= 251 so that neither table wraps as unsigned bytes
and a >= b always holds, making the absolute value a no-op; B = 15 is
free because Ymask already holds 15 in every byte.
The rest:
- Generic (VEX-encoded) AVX instructions accept an unaligned memory
source, so the second input of the And/Or/Xor/Mask loops is read
straight out of memory rather than loaded into a register first. Only
VPANDN's non-negated operand may come from memory, so that loop loads
m into the register and reads s from memory.
- Ydata/Yhi/Yc2 and Ylo/Yc1 have disjoint live ranges and now share
registers; six architectural registers suffice instead of ten.
- lutmask shrinks from 64 to 24 bytes: VBROADCASTI128 duplicates the
16-byte table into both 128-bit lanes and VPBROADCASTB splats the
nibble mask from a single byte.
- SHRQ and ANDQ already set ZF, so the TESTQ instructions that followed
them were redundant.
- SETUP moves below the branch that skips the vector loop entirely.
- The 64-bit tail loops use POPCNTQ with a memory source and a 32-bit
loop counter (32-bit avoids the partial-register merge a DECB needs).
- XORL rather than XORQ for zeroing: no REX prefix, and some cores do
not recognize XORQ as a zeroing idiom.
- _hasAVX2 complements each feature word and TESTs it instead of
AND/CMP, saving a large immediate per check, and all three checks now
fall through to a single SETEQ that stores ZF into the bool result.
Measured on an Intel Xeon Gold 6548N with go1.26.3. Microbenchmarks,
count=10:
PopcntSlice1024AVX2 205.0n -> 180.0n -12.20%
PopcntAndSlice1024AVX2 241.6n -> 225.7n -6.58%
Popcount 14.66n -> 13.73n -6.35%
And through the public API on dense bitmaps (32 bitmap containers,
count=8), which is what the popcount slice routines actually serve:
Bitmap.AndCardinality 9.404µ -> 8.655µ -7.96%
Bitmap.OrCardinality 128.6µ -> 120.8µ -6.02%
The pure-Go fallbacks are unchanged and measure flat, as expected. The
assembly text shrinks from 987 to 914 bytes and the rodata blob from 64
to 24. Verified against the Go reference for all five routines over the
existing differential test plus 20000 randomized rounds with unaligned
sub-slices and all-ones/all-zeros inputs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
David Sparks reviewed
popcnt_avx2_amd64.sand sent a list of improvementsto it. This PR implements all of them; the credit is his.
The main change
VPSADBWcomputes|a-b|per byte and sums each group of 8, so it canabsorb the per-byte add that
VPADDBwas doing. Feeding it the two nibblecounts directly gives
(B + lo) - (B - hi) = lo + hi, which removesVPADDBand its latency from the hot loop:This needs two lookup tables, one biased up by
Band one subtracted fromB. The bias must satisfy4 <= B <= 251so that neither table wraps asunsigned bytes (max nibble popcount is 4) and
a >= balways holds, makingVPSADBW's absolute value a no-op.B = 15is free, becauseYmaskalready holds 15 in every byte — so
SETUPbuilds both tables with oneVPSUBBand oneVPADDBand dropsYzeroentirely.The rest
so the second input of the And/Or/Xor/Mask loops is read straight out of
memory instead of being loaded into a register first. Only
VPANDN'snon-negated operand may come from memory, so that loop loads
mintothe register and reads
sfrom memory.Ydata/Yhi/Yc2andYlo/Yc1have disjoint live ranges and nowshare registers: six architectural registers instead of ten.
lutmaskshrinks from 64 to 24 bytes.VBROADCASTI128duplicates the16-byte table into both 128-bit lanes and
VPBROADCASTBsplats thenibble mask from a single byte, so neither needs to be stored twice.
SHRQandANDQalready set ZF, so theTESTQinstructions thatfollowed them were redundant.
SETUPmoves below the branch that skips the vector loop entirely.POPCNTQwith a memory source and a 32-bitloop counter (32-bit rather than 8-bit avoids a partial-register merge).
XORLrather thanXORQfor zeroing: no REX prefix, and some cores donot recognize
XORQas a zeroing idiom._hasAVX2complements each feature word andTESTs it instead ofAND/CMP, saving a large immediate per check, and all three checks nowfall through to a single
SETEQthat stores ZF into the bool result.The inner loop goes from 13 instructions to 10. Assembly text shrinks from
987 to 914 bytes and the rodata blob from 64 to 24.
Benchmarks
Intel Xeon Gold 6548N, go1.26.3, linux/amd64. Microbenchmarks,
count=10:The pure-Go fallbacks are untouched and measure flat, as a control.
Through the public API on dense bitmaps (32 bitmap containers), which is
what these routines actually serve,
count=8:Testing
Full test suite passes on amd64.
TestAVX2PopcntDifferentialruns ratherthan skipping, which confirms the rewritten
_hasAVX2still returns true;I also disassembled it and the mask loop to check the encodings landed as
intended. Beyond the existing differential test, I ran 20000 randomized
rounds over all five routines with unaligned sub-slices, random lengths,
and all-ones/all-zeros inputs — the numeric extremes of the bias trick —
all matching the Go reference.