Vectorized CLZ and CTZ

purplesyringa's blog

8 min read Original article ↗

Reddit Hacker News

clz and ctz are instructions that compute the number of leading (trailing) zero bits in a fixed-size integer. They are natively supported by modern CPUs, though they are not always fast, e.g. tzcnt has a latency of 3 on Arrow Lake.

I used ctz in an FPU emulator I’m working on, but figured out how to avoid it with floating-point trickery, and I just realized that this generalizes to a vectorizable ctz polyfill in a round-about way. I also implemented clz for completeness.

clz

Let’s start with clz, the easier of the two:

fn clz(x: u32) -> u32 {
    let a = 2.0f64.powi(-970);
    32 - ((f64::from_bits(a.to_bits() | x as u64) - a).to_bits() >> 52) as u32
}

The general idea goes like this:

Floating-point exponents are biased logarithms of their values. By substituting a 32-bit number x into the mantissa of some value 2k, we get a double representing 2k(1+2−52x). We can then subtract 2k as a double to get 2k−52x. Extracting the exponent gives k−52+31−clz(x), from which clz can be computed with a bitwise subtraction. From that, k can be chosen such that clz behaves correctly for x=0.

We need 64-bit doubles to handle 32-bit inputs; unfortunately, this means that this trick can’t work for arbitrary 64-bit inputs, only up to 52 bits.

Assuming the inputs and outputs are stored in u64x4, this compiles to:

    vpbroadcastq ymm1, [rip + bias]

    vorpd ymm0, ymm0, [rip + a]
    vsubpd ymm0, ymm0, [rip + a]
    vpsrlq ymm0, ymm0, 52
    vpsubq ymm0, ymm1, ymm0

a:
    .quad 0x350000000000000, 0x350000000000000, 0x350000000000000, 0x350000000000000
bias:
    .quad 32

On my Haswell, this runs at 0.45 ns/iteration, compared to 1 ns for the scalar version. When latency-bound, the numbers rise to 2 ns vs 1 ns (but if you’re latency-bound on vectorized ctz, you’re probably doing something wrong).

Ian Qvist tested this on Alder Lake (thanks!) and got 0.29 ns/iteration, compared to 0.85 ns for the scalar version, and 1.3 ns vs 0.85 ns when latency-bound. On modern Intel CPUs, the numbers should be the same or better.

AMD CPUs make lzcnt so cheap that a scalar version will likely win. Though keep in mind that Zen CPUs support AVX-512, which has vplzcntd, so that’s an option, too.

Added on October 12:

I found a full 8-lane 32-bit clz implementation by aqrit and optimized it further. The idea is to use i32 → f32 conversion to produce a float instead of building one by hand, so that integers wider than 23 bits can fit in f32.

This is tricky due to rounding, e.g. 230−1 might get rounded to 230 during conversion, producing the wrong result. The algorithm on Stack Overflow fixes this by computing (essentially) x &= !(x >> 1) prior to conversion, which clears the bit below the topmost one to ensure rounding doesn’t change the exponent. But I found a better way:

__m256i cleared = _mm256_min_epu8(x, _mm256_set1_epi8(0x80));
__m256i exp = _mm256_srli_epi32(_mm256_castps_si256(_mm256_cvtepi32_ps(cleared)), 23);
return _mm256_min_epi16(_mm256_subs_epu16(_mm256_set1_epi32(0x9e), exp), _mm256_set1_epi32(32));

All we really need to do is avoid having 20-ish ones in a row, which can be done in one instruction by computing per-byte max with 0x80, which reduces 0xff to 0x80 without affecting the clz.

There are some other tricky details that the original algorithm already handles in a clever manner, like sign-extension, but I won’t focus on them here.

This algorithm runs at 0.6 ns on my Haswell (2.7 ns when latency-bound), so it takes just a little more time than the previous one to process twice as much data. If you don’t need to handle clz(0) gracefully, you can drop the final min for more perf.

I’m still not sure how to handle 64-bit clz, since x86 doesn’t support packed i64 → f64 conversion without AVX-512.

ctz

Now for ctz:

fn ctz(x: u32) -> u32 {
    let a = f64::from_bits((0x340000100000001 ^ x as u64) ^ (x as u64 + u32::MAX as u64))
        - f64::from_bits(0x340000000000000);
    (a.to_bits() >> 52) as u32
}

We start with x⊕︎(x−1) to isolate the lowest set bit. ctz equals the logarithm of that value, which we determine by adding 2k bitwise and subtracting 2k as a double, then inspecting the exponent, which with a well-chosen k contains the unbiased ctz. We pre-mix 232 into the mantissa and use x+232−1 instead of x−1 to handle x=0 correctly, and pre-mix 1 into the mantissa to ensure odd x generate a=0 and not a slow subnormal. (Can you imagine how much time I spent arranging this?)

This function compiles to:

    vpxor ymm1, ymm0, [rip + c1]
    vpaddq ymm0, ymm0, [rip + u32_max]
    vpxor ymm0, ymm1, ymm0
    vsubpd ymm0, ymm0, [rip + c2]
    vpsrlq ymm0, ymm0, 52

c1:
    .qword 0x340000100000001, 0x340000100000001, 0x340000100000001, 0x340000100000001
u32_max:
    .qword 0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff
c2:
    .qword 0x340000000000000, 0x340000000000000, 0x340000000000000, 0x340000000000000

On Haswell, this runs at 0.49 ns/iteration and 2.3 ns when latency-bound. On Alder Lake, it’s 0.35 ns/iteration and 1.3 ns when latency-bound. The slowdown compared to clz is due to using one more instruction. It can be avoided by using vpternlogq if AVX-512 is present, but at that point you might as well run vpopcntd on (x - 1) & !x. The scalar version behaves no differently from clz.

Added on October 10:

Nikolay Malkovsky pointed out that de Bruijn sequences offer another vectorizable approach. After some testing, I arrived at the following code:

const char table[32] = {
    0, 4, 5, 6, 11, 9, 7, 12, 15, 3, 10, 8, 14, 2, 13, 1,
    0, 4, 5, 6, 11, 9, 7, 12, 15, 3, 10, 8, 14, 2, 13, 1,
};
__m256i bit = _mm256_andnot_si256(x, _mm256_sub_epi32(x, _mm256_set1_epi32(1)));
__m256i high = _mm256_madd_epi16(
    _mm256_cmpeq_epi16(bit, _mm256_set1_epi16(-1)),
    _mm256_set1_epi16(-16)
);
__m256i index = _mm256_srli_epi32(_mm256_mullo_epi32(bit, _mm256_set1_epi32(0xf0a6f0a7)), 28);
__m256i low = _mm256_shuffle_epi8(_mm256_loadu_si256((__m256i*)table), index);
return _mm256_add_epi32(low, high);

We can’t use a true 32-byte LUT because vpshufb cannot cross 16-byte lanes. The approach I used instead is tricky to explain, but essentially we use a 16-bit de Bruijn sequence repeated twice to compute bits 0-3 of the ctz, and then add 16 or 32 depending on which halves are zeroes. Six seven 0xf0a6f0a7 is one of only four magic constants that make this work.

This takes 1 ns on Haswell (0.7 ns on Alder Lake) and 3.8 ns when latency-bound, but has twice the throughput, so it may be a little faster than the FP-based approach if it helps avoid shuffling.

If you don’t need to deal with x=0 (or want ctz(0) to be 0 and not 32), using

__m256i high = _mm256_and_si256(
    _mm256_cmpgt_epi32(bit, _mm256_set1_epi32(0x7fff)),
    _mm256_set1_epi32(16)
);

brings the time down to 0.82 ns.

Added on October 11:

arya dradjica and I brainstormed a 64-bit version of the algorithm and managed to find what I think is the optimal algorithm:

const char table[32] = { // x -> 4 - popcnt(x)
    4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0,
    4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0,
};
__m256i lut = _mm256_loadu_si256((__m256i*)table);

__m256i bit = _mm256_andnot_si256(x, _mm256_sub_epi64(x, _mm256_set1_epi64x(1)));
__m256i low = _mm256_shuffle_epi8(lut, bit);
__m256i high = _mm256_shuffle_epi8(lut, _mm256_srli_epi16(bit, 4));
return _mm256_sad_epu8(_mm256_add_epi8(low, high), _mm256_set1_epi8(8));

This correctly returns ctz(0)=64 and takes 0.63 ns on my Haswell (2.8 ns when latency-bound). It computes inverted popcounts of the two nibbles of each byte of (x - 1) & !x separately (inverted because of some idiosyncrasies of pshufb) and then uses vpsadbw to perform the horizontal sum and correct the inversion in one go.

vpsadbw can only horizontally sum exactly 8 bytes, which is why this algorithm is specialized to 64 bits and doesn’t work with 32-bit inputs. The approach based on the de Bruijn sequence is still the fastest in that scenario.

Added on October 12:

More progress! u/UnalignedAxis111 mentioned that floating point-based approaches can be extended to the full 32/64-bit lengths by using int-to-float conversion instructions instead of bit twiddling. This can be slower, but it might be faster than de Bruijn. There are quite a few nuances:

  • It introduces rounding issues, but for ctz that’s not a problem, because x & -x has only one bit set.
  • x86 supports signed-int-to-float casts but not unsigned-int-to-float, which causes problems for x=231.
  • It doesn’t allow tweaking the exponent bias.

After some brainstorming, I got this for 32-bit ctz:

__m256i bit = _mm256_andnot_si256(_mm256_sub_epi32(x, _mm256_set1_epi32(1)), x);
__m256i exp = _mm256_srli_epi32(_mm256_castps_si256(_mm256_cvtepi32_ps(bit)), 23);
return _mm256_min_epu8(_mm256_sub_epi32(exp, _mm256_set1_epi32(0x7f)), _mm256_set1_epi32(32));

This starts by extracting the biased exponent exp of !x & (x - 1), which is:

  • 0x7f when ctz(x)=0
  • 0x80 when ctz(x)=1
  • 0x81 when ctz(x)=2
  • …
  • 0x9d when ctz(x)=30
  • 0x19e when ctz(x)=31 (since x=231 is interpreted as x=−231!)
  • 0x00 when ctz(x)=32 (because that corresponds to x=0)

I managed to map this to ctz in just two instructions:

  1. exp - 0x7f maps most inputs to the correct ctz, but gives 0x11f instead of 0x1f and -0x7f (0xffffff81) instead of 32.
  2. Per-byte min with 32 = 00 00 00 1f clears the second byte of 0x11f and maps the bottom byte of -0x7f to 32 at the same time.

On my Haswell, this takes 0.6 ns when throughput-bound and 3 ns when latency-bound, beating the de Bruijn sequence-based approach.

This can’t be extended to the 64-bit case because pre-AVX-512 x86 lacks packed i64 ↔ f64 conversions. For that, use the vpsadbw approach from yesterday.

P.S.: I find it funny that the 32-bit and 64-bit implementations now have pretty much the same performance characteristics despite using completely different algorithms.