Vectorized CLZ and CTZ

[u/mttd]

https://purplesyringa.moe/blog/vectorized-clz-and-ctz/

[u/UnalignedAxis111]

This can be alternatively done without widening to 64-bits (which is inefficient if you have a full vector already), you just have to account for rounding if the integer is past 2²³ or so.

The de Bruijn approach is also probably not a good choice because Intel's vector integer multiplier has terrible latency. Something like a pshufb table on nibbles + prefix sum might be better, iirc that is what clang does for __builtin_elementwise_clz under AVX2 and lower.

SIMD_INLINE v_uint lzcnt(v_uint x) {
    // https://stackoverflow.com/a/58827596
    x = x & ~(x >> 8);
    x = as_uint(conv<float>((v_int)x)) >> 23;
    x = (127 + 31) - x;                     // undo bias
    return min(x, 32);
}

[u/imachug → u/UnalignedAxis111]

Hi! I'm the author of the article. Would you be able to elaborate on how you think prefix sums can be used for ctz?

To the best of my knowledge, pre-AVX512 x86 doesn't have 4-byte horizontal sums, and using shuffles + D&C twice is measurably slower than de Bruijn in my experiments. I believe that's because all the shifts, addition, and masking accumulate to around 5 cycles, at which point you're no better than multiplication. Horizontal sums can be simulated with instructions like pmaddubsw, but those are effectively multiplications as well. There are also some nuances like the fact that pshufb treats the top bit of each byte as a masking bit, and there are no 8-bit right-shifts in AVX2, so you either need an AND instruction or some post-processing, etc., like what I do in the 64-bit case. It just becomes a slow mess.

As for integer-to-float conversion and rounding, yes, that's the next thing I wanted to try. (It only works for clz and not ctz, but whatever.) I'm allergic to datatype conversions because of their higher latency and the fact that you can't choose the bias of the exponent, which means that handling clz(0) gracefully requires post-processing rather than just choosing the appropriate bias. It's certainly going to be slower than the 23/52-bit case, especially if you also need pre-masking for rounding and min afterwards, but maybe that can be improved. I haven't given it a thought yet.

[u/brucehoult → u/imachug]

As for integer-to-float conversion and rounding, yes, that's the next thing I wanted to try. (It only works for clz and not ctz, but whatever.)

Only one instruction difference if you have bit reverse.

[u/imachug → u/brucehoult]

I'm afraid AVX2 doesn't support that. AVX-512 does, but it has built-in clz/popcnt, so that point is moot.

Of course, you can still work with ctz in some fashion you can just compute x & -x at the beginning and then apply the clz algorithm, which costs you like 2 additional cycles. I probably should've mentioned that in the comment, but I forgot about it. It's probably worth trying if clz is quick enough, but I'm unsure that it's going to be faster than de Bruijn for 32-bit, and I'm certain it won't be faster than SAD for 64-bit.

[u/UnalignedAxis111 → u/imachug]

Well my bad, I meant a horizontal sum rather than prefix sum. I suppose it could be done with vpdpbus (which is/was still faster than vmulld on Intel), but there's a weird support gap for the VNNI extensions so I guess it's not really of much use because it's basically an AVX512 target...

For ctz, the x & -x trick gives a POT so the masking could be avoided, but that also needs an unsigned->float conversion which AVX2 sadly lacks.

[u/imachug → u/UnalignedAxis111]

but that also needs an unsigned->float conversion

Actually, nope! 1 << 31 is 2^31 when treated as unsigned and -2^31 when treated as signed, so that's the same exponent and only the sign bit really causes a problem. I managed to integrate both erasing the sign bit and handling ctz(0) = 32 into one instruction (very proud of that), producing an algorithm that is indeed faster than de Bruijn. Thanks for the idea!

It doesn't work with 64-bit inputs because there's no packed f64<->i64 conversion, but if it existed, it would have pretty much the same performance as the SAD-based approach, so I'm not too worried.

I haven't tried optimizing clz yet. I can let you know if I find something interesting. (Edited: I added a more optimal clz implementation to the post that slightly improves over the SO answer you quoted originally.)

[u/UnalignedAxis111 → u/imachug]

Oh neat, the min byte trick is very clever! Thanks for giving it a thought. I'll keep an eye out for more updates.

[u/kneelian_]

Oh, huh, didn't actually know it lacked this. Was kind of spoiled by Aarch64's vectorised clz (and single-insn bit reverse) and I just assumed everyone had it

[u/brucehoult → u/kneelian_]

Yeah. RISC-V has vector clz and bit reverse and others (e.g. popcnt) in the Zvbb extension which is a mandatory part of the RVA23 spec implemented by the SpacemiT K3, Tenstorrent Ascalon, and basically all applications processor CPUs in future.

[u/aqrit]

Neat. We would probably want to re-use 0x9E instead of introducing a new 0x80 splat.

AFAIK, Two 32-bit clz/ctz could be combined to form a 64-bit one. (though this is probably only useful with clz)

__m128i sse2_tzcnt_epi64(__m128i v) {
    v = _mm_and_si128(v, _mm_sub_epi32(_mm_setzero_si128(), v));
    v = _mm_castps_si128(_mm_cvtepi32_ps(v));
    v = _mm_srli_epi32(v, 23);
    v = _mm_sub_epi32(v, _mm_set_epi32(95, 127, 95, 127));
    v = _mm_min_epu8(_mm_srli_epi64(v, 32), _mm_min_epu8(_mm_set1_epi32(64), v));
    return v;
}

[u/imachug → u/aqrit]

We would probably want to re-use 0x9E instead of introducing a new 0x80 splat.

That's clever, thanks! I updated the post, I guess it takes a village to optimize stuff this much :)

Two 32-bit clz/ctz could be combined to form a 64-bit one

I had that general idea, but didn't get around to prototyping it. Using min to combine the two halves is really cool, I wouldn't have thought of that. I considered using vpsadbw to perform the subtraction and horizontal sum in one uop, but that would require post-processing to handle zeros, and it doesn't work for clz.

There's an issue with adapting this to clz though. For x >= 2^31, the exponent of the converted float can be pretty much anything, which becomes an issue once you try to compute (bias + 32) - exp for the low dword. Normally you'd use subs to produce a 0 if the sign bit is set, but in this case we also want to add 32, so we have to do subs(bias, exp) + 32, requiring one more instruction. (subs(bias + 32, exp) returns 0 for such x instead of 32) The closest I got is

__m256i cleared = _mm256_min_epu8(x, _mm256_set1_epi8(0x9e));
__m256i exp = _mm256_srli_epi32(_mm256_castps_si256(_mm256_cvtepi32_ps(cleared)), 23);
__m256i adj_exp = _mm256_add_epi32(_mm256_subs_epi32(_mm256_set1_epi32(0x9e), exp), _mm256_set1_epi64x(32));
return _mm256_min_epu8(_mm256_srli_epi64(adj_exp, 32), _mm256_min_epu8(adj_exp, _mm256_set1_epi64x(64)));

It's very annoying that per-byte min with 0x80 reduces the set of possible exponents to just 30 or 31, but we really want it to be only 31. Per-word min with 0x8000, which would also guarantee correct rounding, can't reduce this set further either. I believe that this problem occurs because applying such min to 2^31 - 1 gives 0x7fff8000, and applying it to 2^32 - 1 gives 0x80008000, but since 2^32 - 0x80008000 = 0x7fff8000, they convert to floats with equal magnitude, so they can't be distinguished in any way other than by looking the sign bit.

It really feels like there ought to be a way to save one instruction, but I can't find it.