I would like to extract the index of the highest set bit in a 256 bit AVX register with 8 bit elements. I could neither find a bsr nor a clz implementation for this.
For clz with 32 bit elements, there is the bithack with a float conversion, but that is probably not possible for 8 bits.
Currently, I am working on a solution, where I check the bits one by one, which I will add later, but I wonder if there is a faster way to do this.
Here is a vpshufb based solution. The idea is to split the input into two halves, make a lookup on both and combine the results:
__m256i clz_epu8(__m256i values)
{
// extract upper nibble:
__m256i hi = _mm256_and_si256(_mm256_srli_epi16(values, 4), _mm256_set1_epi8(0xf));
// this sets the highest bit for values >= 0x10 and otherwise keeps the lower nibble unmodified:
__m256i lo = _mm256_adds_epu8(values, _mm256_set1_epi8(0x70));
// lookup tables for count-leading-zeros (replace this by _mm256_setr_epi8, if this does not get optimized away)
// ideally, this should compile to vbroadcastf128 ...
const __m256i lookup_hi = _mm256_broadcastsi128_si256(_mm_setr_epi8(0, 3, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0));
const __m256i lookup_lo = _mm256_broadcastsi128_si256(_mm_setr_epi8(8, 7, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4));
// look up each half
__m256i clz_hi = _mm256_shuffle_epi8(lookup_hi, hi);
__m256i clz_lo = _mm256_shuffle_epi8(lookup_lo, lo);
// combine results (addition or xor would work as well)
return _mm256_or_si256(clz_hi, clz_lo);
}
godbolt-link with crude test: https://godbolt.org/z/MYq74Wxdh
Normally _mm_shuffle_epi8 needs masking to isolate each nibble to use it as a LUT, because having the high bit set makes that output element 0. But for CLZ, if the high bit is set, the correct result is for the whole byte is 0, and the way we combine means it's fine for lut_lo to produce it.
__m128i ssse3_lzcnt_epi8(__m128i v) {
const __m128i lut_lo = _mm_set_epi8(4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 7, 8);
const __m128i lut_hi = _mm_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 8);
__m128i t;
t = _mm_and_si128(_mm_srli_epi16(v, 4), _mm_set1_epi8(0x0F));
t = _mm_shuffle_epi8(lut_hi, t);
v = _mm_shuffle_epi8(lut_lo, v);
v = _mm_min_epu8(v, t);
return v;
}
This saves an instruction vs. using _mm_adds_epu8 and combining LUT results with or.
AVX512 solution, didn't try it, but I think the idea should work:
// Form four 32-bit vectors with high bytes from the source
__m256i a0 = _mm256_or_si256(_mm256_slli_si256(a, 3), _mm256_set1_epi32(0x00FF'FFFF));
__m256i a1 = _mm256_or_si256(_mm256_slli_si256(a, 2), _mm256_set1_epi32(0x00FF'FFFF));
__m256i a2 = _mm256_or_si256(_mm256_slli_si256(a, 1), _mm256_set1_epi32(0x00FF'FFFF));
__m256i a3 = _mm256_or_si256( a, _mm256_set1_epi32(0x00FF'FFFF));
// Count lead bits and shift according to bit position
__m256i c0 = _mm256_lzcnt_epi32(a0);
__m256i c1 = _mm256_slli_si256(_mm256_lzcnt_epi32(a1), 1);
__m256i c2 = _mm256_slli_si256(_mm256_lzcnt_epi32(a2), 2);
__m256i c3 = _mm256_slli_si256(_mm256_lzcnt_epi32(a3), 3);
//Gather the result
__m256i r = _mm256_or_si256(_mm256_or_si256(c0,c1),_mm256_or_si256(c2,c3));
Not sure if it is faster than checking one by one
Given a target AVX register _a, this works. Let me know (or edit directly) if there is something to optimize.
__m256i _a;
__m256i _old_mask = _mm256_set1_epi8(-1);
__m256i _extract_bitmask, _extracted_bit, _mask;
for (int i = 7; i >= 0; i--)
{
// bitmask to extract bit from _a at position i
_extract_bitmask = _mm256_set1_epi8(1 << i);
// the extracted bit
_extracted_bit = _mm256_and_si256(_a, _extract_bitmask);
// check if bit at position i is set and if was not set before
_mask = _mm256_cmpeq_epi8(_extract_bitmask, _extracted_bit);
_mask = _mm256_and_si256(_mask, _old_mask);
// update mask
_old_mask = _mm256_andnot_si256(_mask, _old_mask);
// update result according to _mask
_result = _mm256_blendv_epi8(_result, _mm256_set1_epi8(i), _mask);
}