diff --git a/popcnt_avx2_amd64.s b/popcnt_avx2_amd64.s index 1c61a237..52e2c989 100644 --- a/popcnt_avx2_amd64.s +++ b/popcnt_avx2_amd64.s @@ -13,10 +13,10 @@ // AVX2 has no single "popcount a whole vector" instruction, so each byte's // popcount is taken from a 16-entry lookup table indexed by a 4-bit nibble: // a byte is split into its low and high nibble, each nibble is looked up (one -// VPSHUFB performs all 32 lookups in a 256-bit register at once), the two -// results are added to give a per-byte popcount, and VPSADBW then sums each -// group of 8 byte-counts into a 64-bit lane total that is accumulated. After -// the loop the four lane totals are summed (HSUM) into a scalar register. +// VPSHUFB performs all 32 lookups in a 256-bit register at once), and VPSADBW +// then combines the two results and sums each group of 8 byte-counts into a +// 64-bit lane total that is accumulated (COUNTBLOCK). +// After the loop the four lane totals are summed (HSUM) into a scalar register. // Each iteration handles 256 bits (4 uint64); a scalar POPCNTQ tail handles // the trailing len%4 words, so any slice length is counted correctly. // @@ -31,85 +31,102 @@ // (e.g. ret+24(FP) for one slice arg, ret+48(FP) for two). // - Every routine is a leaf (makes no calls): NOSPLIT with a $0 local frame. // - Loads/stores use VMOVDQU (unaligned): container slices are only 8-byte -// aligned, not 32. VZEROUPPER precedes every RET to avoid the AVX<->SSE -// transition penalty in any non-VEX SSE code that runs afterwards. +// aligned, not 32. Generic (VEX-encoded) AVX instructions impose no +// alignment requirement on a memory source either, so the second input of +// the two-operand loops is read straight out of memory by VPAND/VPOR/ +// VPXOR/VPANDN instead of being loaded into a register first. +// VZEROUPPER precedes every RET to avoid the AVX<->SSE transition penalty +// in any non-VEX SSE code that runs afterwards. -// lutmask is a 64-byte read-only blob holding two constants used by every -// routine: -// bytes 0..31 - the nibble popcount table, i.e. table[i] = number of set -// bits in the 4-bit value i. VPSHUFB indexes within each -// 128-bit lane independently, so the 16-entry table is stored -// twice (once per lane). Read low-byte-first, the first qword -// 0x0302020102010100 is the bytes {0,1,1,2,1,2,2,3} for +// lutmask is a 17-byte read-only blob (the linker pads it out to whatever its +// alignment requires) holding the two constants used by every routine: +// bytes 0..15 - the nibble popcount table, i.e. table[i] = number of set +// bits in the 4-bit value i. Read low-byte-first, the first +// qword 0x0302020102010100 is the bytes {0,1,1,2,1,2,2,3} for // nibbles 0..7, and 0x0403030203020201 is {1,2,2,3,2,3,3,4} -// for nibbles 8..15. -// bytes 32..63 - 0x0F in every byte: a mask that isolates the low nibble of -// each byte. +// for nibbles 8..15. VPSHUFB indexes within each 128-bit lane +// independently and so needs the table in both lanes, but +// VBROADCASTI128 duplicates the 16 bytes at load time; only +// one copy has to be stored. +// byte 16 - 0x0F, the mask that isolates the low nibble of each byte, +// splatted to all 32 bytes by VPBROADCASTB. // RODATA|NOPTR marks it read-only and pointer-free (so the GC ignores it). DATA lutmask<>+0(SB)/8, $0x0302020102010100 DATA lutmask<>+8(SB)/8, $0x0403030203020201 -DATA lutmask<>+16(SB)/8, $0x0302020102010100 -DATA lutmask<>+24(SB)/8, $0x0403030203020201 -DATA lutmask<>+32(SB)/8, $0x0f0f0f0f0f0f0f0f -DATA lutmask<>+40(SB)/8, $0x0f0f0f0f0f0f0f0f -DATA lutmask<>+48(SB)/8, $0x0f0f0f0f0f0f0f0f -DATA lutmask<>+56(SB)/8, $0x0f0f0f0f0f0f0f0f -GLOBL lutmask<>(SB), RODATA|NOPTR, $64 +DATA lutmask<>+16(SB)/1, $0x0f +GLOBL lutmask<>(SB), RODATA|NOPTR, $17 -// Register aliases. Ylut/Ymask/Yzero are constants set up once per call (see -// SETUP); Yacc is the running accumulator of lane totals; Ydata/Yb hold the -// current input vector(s); Ylo/Yhi/Yc1/Yc2 are scratch used by COUNTBLOCK. -#define Ylut Y0 -#define Ymask Y1 -#define Yzero Y2 +// Register aliases. Ylut1/Ylut2/Ymask are constants set up once per call (see +// SETUP); Yacc is the running accumulator of lane totals; Ydata holds the +// current input vector; Ylo/Yhi are scratch used by COUNTBLOCK. Ydata is dead +// once its nibbles have been extracted, so Yhi shares its register: only five +// architectural registers are needed. +#define Ylut1 Y0 +#define Ylut2 Y1 +#define Ymask Y2 #define Yacc Y3 #define Ydata Y4 -#define Yb Y5 -#define Ylo Y6 -#define Yhi Y7 -#define Yc1 Y8 -#define Yc2 Y9 +#define Yhi Y4 +#define Ylo Y5 + +// Low 128-bit halves of Yacc and Ylo, used as scratch by HSUM. +#define Xacc X3 +#define Xtmp X5 // COUNTBLOCK folds the popcount of the 32 bytes currently in Ydata into the // accumulator Yacc. Line by line: -// VPAND Ymask,Ydata,Ylo : Ylo = low nibble of every byte -// VPSRLW $4,Ydata,Yhi : shift each 16-bit lane right by 4... -// VPAND Ymask,Yhi,Yhi : ...then mask, leaving the high nibble of each byte -// VPSHUFB Ylo,Ylut,Yc1 : Yc1[b] = popcount(low nibble of byte b) -// VPSHUFB Yhi,Ylut,Yc2 : Yc2[b] = popcount(high nibble of byte b) -// VPADDB Yc2,Yc1,Yc1 : Yc1[b] = popcount(byte b) (0..8 each) -// VPSADBW Yzero,Yc1,Yc1 : sum each group of 8 bytes -> 4 lane totals (0..512) -// VPADDQ Yc1,Yacc,Yacc : add the 4 lane totals into the accumulator +// VPAND Ymask,Ydata,Ylo : Ylo = low nibble of every byte +// VPSRLW $4,Ydata,Yhi : shift each 16-bit lane right by 4... +// VPAND Ymask,Yhi,Yhi : ...then mask, leaving the high nibble of each byte +// VPSHUFB Ylo,Ylut1,Ylo : Ylo[b] = B + popcount(low nibble of byte b) +// VPSHUFB Yhi,Ylut2,Yhi : Yhi[b] = B - popcount(high nibble of byte b) +// VPSADBW Ylo,Yhi,Ylo : sum each group of 8 bytes -> 4 lane totals +// VPADDQ Ylo,Yacc,Yacc : add the 4 lane totals into the accumulator +// +// VPSADBW computes |a-b| per byte and sums each group of 8, so it can do the +// work of the per-byte add as well: feeding it the two nibble counts directly +// yields (B + lo) - (B - hi) = lo + hi, and the separate VPADDB the naive +// version needs (with a zero second VPSADBW operand) disappears along with its +// latency. That is why SETUP builds two tables, one biased up by B and one +// subtracted from B. The bias must satisfy 4 <= B <= 251 so that neither +// table's entries (max nibble popcount is 4) wrap around as unsigned bytes and +// so that a >= b always holds, making the absolute value a no-op; B = 15 is +// used simply because Ymask already holds 15 in every byte. +// // Per-byte counts max at 8 and lane totals at 512, so accumulating across the // whole loop never overflows the 64-bit lanes. #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 \ - VPADDQ Yc1, Yacc, Yacc + VPSHUFB Ylo, Ylut1, Ylo \ + VPSHUFB Yhi, Ylut2, Yhi \ + VPSADBW Ylo, Yhi, Ylo \ + VPADDQ Ylo, Yacc, Yacc -// SETUP loads the lookup table and nibble mask and zeroes Yzero (the VPSADBW -// addend) and Yacc (the accumulator). Run once at the top of each routine. +// SETUP builds the two biased lookup tables and the nibble mask, and zeroes +// Yacc (the accumulator). Run once per routine, after the check that the +// vector loop runs at least one iteration. Ylut1 first holds the raw table, so +// the VPSUBB must come before the VPADDB that overwrites it. #define SETUP \ - VMOVDQU lutmask<>+0(SB), Ylut \ - VMOVDQU lutmask<>+32(SB), Ymask \ - VPXOR Yzero, Yzero, Yzero \ - VPXOR Yacc, Yacc, Yacc + VBROADCASTI128 lutmask<>+0(SB), Ylut1 \ + VPBROADCASTB lutmask<>+16(SB), Ymask \ + VPXOR Yacc, Yacc, Yacc \ + VPSUBB Ylut1, Ymask, Ylut2 \ + VPADDB Ylut1, Ymask, Ylut1 -// HSUM reduces Yacc's four 64-bit lane totals to a single sum in AX. X3 is the -// low 128 bits of Yacc (Y3); VEXTRACTI128 pulls the high 128 bits into X5, the -// two halves are added (giving two qwords in X3), and those two qwords are then -// added into AX. +// HSUM reduces Yacc's four 64-bit lane totals to a single sum in AX. +// VEXTRACTI128 pulls Yacc's high 128 bits into Xtmp and the two halves are +// added, leaving two qwords in Xacc; VPSHUFD $0x4e then swaps those two qwords +// so a second VPADDQ puts their total in the low qword, which one MOVQ moves +// out. Finishing the reduction in SIMD avoids VPEXTRQ, which is 2 uops on both +// AMD Zen and Intel, against 1 each for VPSHUFD and VPADDQ. #define HSUM \ - VEXTRACTI128 $1, Yacc, X5 \ - VPADDQ X5, X3, X3 \ - VPEXTRQ $1, X3, DX \ - MOVQ X3, R9 \ - ADDQ R9, AX \ + VEXTRACTI128 $1, Yacc, Xtmp \ + VPADDQ Xtmp, Xacc, Xacc \ + VPSHUFD $0x4e, Xacc, Xtmp \ + VPADDQ Xtmp, Xacc, Xacc \ + MOVQ Xacc, DX \ ADDQ DX, AX // func _popcntSliceAVX2(s []uint64) uint64 @@ -119,12 +136,11 @@ GLOBL lutmask<>(SB), RODATA|NOPTR, $64 TEXT ·_popcntSliceAVX2(SB), NOSPLIT, $0-32 MOVQ s_base+0(FP), SI // SI = &s[0] MOVQ s_len+8(FP), CX // CX = len(s), in 64-bit words - XORQ AX, AX // AX = running result - SETUP // load table/mask; zero Yzero and Yacc + XORL AX, AX // AX = running result MOVQ CX, R8 SHRQ $2, R8 // R8 = len/4 = number of full 256-bit blocks - TESTQ R8, R8 JZ slicetail // fewer than 4 words: skip the vector loop + SETUP // load tables/mask; zero Yacc sliceloop: VMOVDQU (SI), Ydata // load 4 words (32 bytes) COUNTBLOCK // Yacc += popcount(those 32 bytes) @@ -133,15 +149,13 @@ sliceloop: JNZ sliceloop HSUM // AX += sum of Yacc's lane totals slicetail: - ANDQ $3, CX // CX = len % 4 = leftover words (0..3) - TESTQ CX, CX + ANDL $3, CX // CX = len % 4 = leftover words (0..3) JZ slicedone slicetailloop: - MOVQ (SI), DX - POPCNTQ DX, DX // scalar popcount of one word + POPCNTQ (SI), DX // scalar popcount of one word ADDQ DX, AX ADDQ $8, SI - DECQ CX + DECL CX JNZ slicetailloop slicedone: VZEROUPPER // clear upper YMM state before returning @@ -149,23 +163,21 @@ slicedone: RET // func _popcntAndSliceAVX2(s, m []uint64) uint64 -// Returns the sum of popcount(s[i] & m[i]). Mirrors _popcntSliceAVX2 but loads -// a vector from each of s and m and ANDs them before counting. s and m are +// Returns the sum of popcount(s[i] & m[i]). Mirrors _popcntSliceAVX2 but ANDs +// a vector of s with the matching bytes of m before counting. s and m are // assumed to have equal length. TEXT ·_popcntAndSliceAVX2(SB), NOSPLIT, $0-56 MOVQ s_base+0(FP), SI // SI = &s[0] MOVQ m_base+24(FP), DI // DI = &m[0] MOVQ s_len+8(FP), CX // CX = len - XORQ AX, AX - SETUP + XORL AX, AX MOVQ CX, R8 SHRQ $2, R8 - TESTQ R8, R8 JZ andtail + SETUP andloop: VMOVDQU (SI), Ydata - VMOVDQU (DI), Yb - VPAND Yb, Ydata, Ydata // Ydata = s & m + VPAND (DI), Ydata, Ydata // Ydata = s & m COUNTBLOCK ADDQ $32, SI ADDQ $32, DI @@ -173,8 +185,7 @@ andloop: JNZ andloop HSUM andtail: - ANDQ $3, CX - TESTQ CX, CX + ANDL $3, CX JZ anddone andtailloop: MOVQ (SI), DX @@ -183,7 +194,7 @@ andtailloop: ADDQ DX, AX ADDQ $8, SI ADDQ $8, DI - DECQ CX + DECL CX JNZ andtailloop anddone: VZEROUPPER @@ -197,16 +208,14 @@ TEXT ·_popcntOrSliceAVX2(SB), NOSPLIT, $0-56 MOVQ s_base+0(FP), SI MOVQ m_base+24(FP), DI MOVQ s_len+8(FP), CX - XORQ AX, AX - SETUP + XORL AX, AX MOVQ CX, R8 SHRQ $2, R8 - TESTQ R8, R8 JZ ortail + SETUP orloop: VMOVDQU (SI), Ydata - VMOVDQU (DI), Yb - VPOR Yb, Ydata, Ydata // Ydata = s | m + VPOR (DI), Ydata, Ydata // Ydata = s | m COUNTBLOCK ADDQ $32, SI ADDQ $32, DI @@ -214,8 +223,7 @@ orloop: JNZ orloop HSUM ortail: - ANDQ $3, CX - TESTQ CX, CX + ANDL $3, CX JZ ordone ortailloop: MOVQ (SI), DX @@ -224,7 +232,7 @@ ortailloop: ADDQ DX, AX ADDQ $8, SI ADDQ $8, DI - DECQ CX + DECL CX JNZ ortailloop ordone: VZEROUPPER @@ -238,16 +246,14 @@ TEXT ·_popcntXorSliceAVX2(SB), NOSPLIT, $0-56 MOVQ s_base+0(FP), SI MOVQ m_base+24(FP), DI MOVQ s_len+8(FP), CX - XORQ AX, AX - SETUP + XORL AX, AX MOVQ CX, R8 SHRQ $2, R8 - TESTQ R8, R8 JZ xortail + SETUP xorloop: VMOVDQU (SI), Ydata - VMOVDQU (DI), Yb - VPXOR Yb, Ydata, Ydata // Ydata = s ^ m + VPXOR (DI), Ydata, Ydata // Ydata = s ^ m COUNTBLOCK ADDQ $32, SI ADDQ $32, DI @@ -255,8 +261,7 @@ xorloop: JNZ xorloop HSUM xortail: - ANDQ $3, CX - TESTQ CX, CX + ANDL $3, CX JZ xordone xortailloop: MOVQ (SI), DX @@ -265,7 +270,7 @@ xortailloop: ADDQ DX, AX ADDQ $8, SI ADDQ $8, DI - DECQ CX + DECL CX JNZ xortailloop xordone: VZEROUPPER @@ -275,21 +280,22 @@ xordone: // func _popcntMaskSliceAVX2(s, m []uint64) uint64 // Returns the sum of popcount(s[i] &^ m[i]) == popcount(s & ~m). Same structure // as _popcntAndSliceAVX2; the combine is VPANDN, which computes (NOT first) AND -// second, i.e. VPANDN Ydata, Yb, Ydata -> Ydata = (NOT Yb) AND Ydata = s &^ m. +// second. Only VPANDN's second source may come from memory, and it is the +// operand that is *not* negated, so here it is m that is loaded into a register +// and s that is read straight out of memory: +// "VPANDN (SI), Ydata, Ydata" with Ydata = m gives (NOT m) AND s = s &^ m. TEXT ·_popcntMaskSliceAVX2(SB), NOSPLIT, $0-56 MOVQ s_base+0(FP), SI MOVQ m_base+24(FP), DI MOVQ s_len+8(FP), CX - XORQ AX, AX - SETUP + XORL AX, AX MOVQ CX, R8 SHRQ $2, R8 - TESTQ R8, R8 JZ masktail + SETUP maskloop: - VMOVDQU (SI), Ydata - VMOVDQU (DI), Yb - VPANDN Ydata, Yb, Ydata // Ydata = s &^ m (= (NOT m) AND s) + VMOVDQU (DI), Ydata // Ydata = m + VPANDN (SI), Ydata, Ydata // Ydata = s &^ m (= (NOT m) AND s) COUNTBLOCK ADDQ $32, SI ADDQ $32, DI @@ -297,8 +303,7 @@ maskloop: JNZ maskloop HSUM masktail: - ANDQ $3, CX - TESTQ CX, CX + ANDL $3, CX JZ maskdone masktailloop: MOVQ (DI), R10 @@ -309,7 +314,7 @@ masktailloop: ADDQ DX, AX ADDQ $8, SI ADDQ $8, DI - DECQ CX + DECL CX JNZ masktailloop maskdone: VZEROUPPER @@ -320,33 +325,36 @@ maskdone: // Reports whether the CPU supports AVX2 and the OS has enabled the wide (YMM) // register state. All three checks must pass; otherwise the Go wrappers fall // back to the scalar implementation. Note CPUID clobbers AX/BX/CX/DX. +// +// Each check complements the feature word and then TESTs the bits of interest: +// ZF is set exactly when every required bit was set in the original value. That +// needs only one large immediate instead of the two an AND/CMP pair would +// encode, and it lets all three checks converge on a single SETEQ, which stores +// the final ZF straight into the bool result. Since this runs once per process, +// code size matters more here than the (negligible) speed difference. TEXT ·_hasAVX2(SB), NOSPLIT, $0-1 - // CPUID leaf 1: require OSXSAVE (ECX bit 27) and AVX (ECX bit 28). Both must - // be set, so mask and compare against the combined bit pattern. + // CPUID leaf 1: require OSXSAVE (ECX bit 27) and AVX (ECX bit 28). MOVL $1, AX XORL CX, CX CPUID - ANDL $0x18000000, CX - CMPL CX, $0x18000000 + NOTL CX + TESTL $0x18000000, CX JNE noavx2 // XGETBV(0): the OS must have enabled saving of SSE and AVX/YMM state, i.e. // XCR0 bits 1 and 2. Without this the YMM registers would be corrupted // across a context switch even though the CPU supports the instructions. XORL CX, CX XGETBV - ANDL $0x6, AX - CMPL AX, $0x6 + NOTL AX + TESTL $0x6, AX JNE noavx2 // CPUID leaf 7, sub-leaf 0: require AVX2 itself (EBX bit 5). The sub-leaf is // selected via ECX, which must be 0. MOVL $7, AX XORL CX, CX CPUID - ANDL $0x20, BX - CMPL BX, $0x20 - JNE noavx2 - MOVB $1, ret+0(FP) - RET + NOTL BX + TESTL $0x20, BX noavx2: - MOVB $0, ret+0(FP) + SETEQ ret+0(FP) // ZF is still set by whichever TESTL ran last RET