怎么用向量指令计算多个元素尾部 0 的个数?

大家都知道 gcc 提供了__builtin_ctz 这种内建函数计算整数尾部 0 的个数,它在大部分情况下会转译成 CPU 的内建指令,比如 x86 平台上的 tzcnt 指令。

但如果你想同时计算多个元素尾部 0 的个数,就会发现有些指令集,比如 SSE/AVX,并没有提供类似的向量指令。而这种指令有时候在写向量化代码时又是需要的,我自己针对 8-bits 单字节整数做了一些粗糙的实现,如果你通过搜索引擎来到这里,希望对你能有所帮助。

函数功能

输入一个整数向量(16 或者 32 个 uint8),返回一个同类型整数向量,其中输出的每一个整数表示对应位置输入整数末尾 0 的个数。

SSE 版本

#include <smmintrin.h>
__m128i tzcnt_epi8(__m128i a) {
    __m128i ctz_table = _mm_setr_epi8(
        4, 0, 1, 0, 2, 0, 1, 0,
        3, 0, 1, 0, 2, 0, 1, 0
    );
    __m128i shuffle_mask = _mm_set1_epi8(0x0F);
    __m128i v_4 = _mm_set1_epi8(0x4);
    __m128i a_low = _mm_and_si128(a, shuffle_mask);
    __m128i ctz_low = _mm_shuffle_epi8(ctz_table, a_low);
    __m128i a_high = _mm_and_si128(_mm_srli_epi16(a, 4), shuffle_mask);
    __m128i ctz_high = _mm_shuffle_epi8(ctz_table, a_high);
    ctz_high = _mm_and_si128(ctz_high, _mm_cmpeq_epi8(ctz_low, v_4));
    return _mm_add_epi8(ctz_low, ctz_high);
}

AVX 版本

#include <immintrin.h>
__m256i tzcnt_epi8(__m256i a) {
    __m256i ctz_table = _mm256_setr_epi8(
        4, 0, 1, 0, 2, 0, 1, 0,
        3, 0, 1, 0, 2, 0, 1, 0,
        4, 0, 1, 0, 2, 0, 1, 0,
        3, 0, 1, 0, 2, 0, 1, 0
    );
    __m256i shuffle_mask = _mm256_set1_epi8(0x0F);
    __m256i v_4 = _mm256_set1_epi8(0x4);
    __m256i a_low = _mm256_and_si256(a, shuffle_mask);
    __m256i ctz_low = _mm256_shuffle_epi8(ctz_table, a_low);
    __m256i a_high = _mm256_and_si256(_mm256_srli_epi16(a, 4), shuffle_mask);
    __m256i ctz_high = _mm256_shuffle_epi8(ctz_table, a_high);
    ctz_high = _mm256_and_si256(ctz_high, _mm256_cmpeq_epi8(ctz_low, v_4));
    return _mm256_add_epi8(ctz_low, ctz_high);
}

NEON 版本

#include <arm_neon.h>
uint8x16_t vctzq_u8(uint8x16_t a) {
    return vclzq_u8(vrbitq_u8(a));
}

发表回复

您的电子邮箱地址不会被公开。 必填项已用 * 标注