Skip to content

Optimize AVX2 popcount inner loop and CPU feature detection - #546

Open
lemire wants to merge 1 commit into
masterfrom
avx2-popcnt-sparks
Open

Optimize AVX2 popcount inner loop and CPU feature detection#546
lemire wants to merge 1 commit into
masterfrom
avx2-popcnt-sparks

Conversation

@lemire

@lemire lemire commented Aug 12, 2026

Copy link
Copy Markdown
Member

David Sparks reviewed popcnt_avx2_amd64.s and sent a list of improvements
to it. This PR implements all of them; the credit is his.

The main change

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 gives (B + lo) - (B - hi) = lo + hi, which removes
VPADDB and its latency from the hot loop:

 #define COUNTBLOCK \
 	VPAND Ymask, Ydata, Ylo \
 	VPSRLW $4, Ydata, Yhi \
 	VPAND Ymask, Yhi, Yhi \
-	VPSHUFB Ylo, Ylut, Yc1 \
-	VPSHUFB Yhi, Ylut, Yc2 \
-	VPADDB Yc2, Yc1, Yc1 \
-	VPSADBW Yzero, Yc1, Yc1 \
+	VPSHUFB Ylo, Ylut1, Yc1 \
+	VPSHUFB Yhi, Ylut2, Yc2 \
+	VPSADBW Yc1, Yc2, Yc1 \
 	VPADDQ Yc1, Yacc, Yacc

This 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 (max nibble popcount is 4) and a >= b always holds, making
VPSADBW's absolute value a no-op. B = 15 is free, because Ymask
already holds 15 in every byte — so SETUP builds both tables with one
VPSUBB and one VPADDB and drops Yzero entirely.

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 instead of being 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 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, so neither needs to be stored twice.
  • 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 rather than 8-bit avoids a partial-register merge).
  • 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.

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:

                          │    base     │             new              │
PopcntSlice1024AVX2-128     205.0n ± 0%   180.0n ± 0%  -12.20% (p=0.000)
PopcntAndSlice1024AVX2-128  241.6n ± 1%   225.7n ± 1%   -6.58% (p=0.000)
Popcount-128                14.66n ± 1%   13.73n ± 1%   -6.35% (p=0.000)
PopcntSlice1024Go-128       305.2n ± 0%   305.1n ± 1%        ~ (p=0.811)
PopcntAndSlice1024Go-128    449.0n ± 1%   448.9n ± 0%        ~ (p=0.616)

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:

                          │    base     │             new              │
DenseAndCardinality-128     9.404µ ± 1%   8.655µ ± 0%   -7.96% (p=0.000)
DenseOrCardinality-128      128.6µ ± 3%   120.8µ ± 4%   -6.02% (p=0.001)

Testing

Full test suite passes on amd64. TestAVX2PopcntDifferential runs rather
than skipping, which confirms the rewritten _hasAVX2 still 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.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant