如何在C++中交错3个AVX寄存器的字节？

针对你的需求，我们可以利用AVX2的_mm256_shuffle_epi8（PSHUFB）指令，通过构造掩码分别提取三个输入数组的元素到目标位置，再合并得到预期的升序交错序列。以下是完整的实现代码：

#include <immintrin.h>
#include <stdint.h>
#include <stdio.h>

int main() {
    uint8_t a[] = {0, 3, 6, 9, 12, 15, 18, 21};
    uint8_t b[] = {1, 4, 7, 10, 13, 16, 19, 22};
    uint8_t c[] = {2, 5, 8, 11, 14, 17, 20, 23};
    uint8_t result[32] = {0};

    __m256i as = _mm256_loadu_epi8(a);
    __m256i bs = _mm256_loadu_epi8(b);
    __m256i cs = _mm256_loadu_epi8(c);

    // 构造掩码：将a的元素放到结果的0,3,6,...,21位置，其余置0
    __m256i mask_a = _mm256_setr_epi8(
        0x00, 0x80, 0x80, 0x01, 0x80, 0x80, 0x02, 0x80, 0x80, 0x03, 0x80, 0x80,
        0x04, 0x80, 0x80, 0x05, 0x80, 0x80, 0x06, 0x80, 0x80, 0x07, 0x80, 0x80,
        0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80
    );
    // 构造掩码：将b的元素放到结果的1,4,7,...,22位置，其余置0
    __m256i mask_b = _mm256_setr_epi8(
        0x80, 0x00, 0x80, 0x80, 0x01, 0x80, 0x80, 0x02, 0x80, 0x80, 0x03, 0x80,
        0x80, 0x04, 0x80, 0x80, 0x05, 0x80, 0x80, 0x06, 0x80, 0x80, 0x07, 0x80,
        0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80
    );
    // 构造掩码：将c的元素放到结果的2,5,8,...,23位置，其余置0
    __m256i mask_c = _mm256_setr_epi8(
        0x80, 0x80, 0x00, 0x80, 0x80, 0x01, 0x80, 0x80, 0x02, 0x80, 0x80, 0x03,
        0x80, 0x80, 0x04, 0x80, 0x80, 0x05, 0x80, 0x80, 0x06, 0x80, 0x80, 0x07,
        0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80
    );

    // 分别提取三个数组的元素到目标位置
    __m256i res_a = _mm256_shuffle_epi8(as, mask_a);
    __m256i res_b = _mm256_shuffle_epi8(bs, mask_b);
    __m256i res_c = _mm256_shuffle_epi8(cs, mask_c);

    // 合并三个中间结果
    __m256i result_vec = _mm256_or_si256(_mm256_or_si256(res_a, res_b), res_c);

    // 存储结果到数组
    _mm256_storeu_epi8(result, result_vec);

    // 验证输出
    for (int i = 0; i < 24; ++i) {
        printf("%d ", result[i]);
    }
    printf("\n");

    return 0;
}

关键逻辑解释 #

1. PSHUFB指令作用：_mm256_shuffle_epi8是AVX2中用于字节级重排的指令，它根据掩码参数的每个字节来决定目标寄存器对应位置的内容：
   • 掩码字节的最高位为1时，目标字节置0；
   • 最高位为0时，低4位表示源寄存器中对应128位lane内的字节索引，高4位选择128位lane（这里我们只用低128位，所以高4位为0）。
2. 掩码构造：三个掩码分别将a、b、c数组的第i个元素映射到结果数组的3i、3i+1、3i+2位置，其余位置通过掩码最高位置1来置0。
3. 结果合并：由于三个中间结果的非零位置互不重叠，直接通过按位或操作即可合并得到最终的交错序列。

运行这段代码后，结果数组的前24个元素会是0,1,2,3,4,5,...,23，完全符合你的预期。

内容的提问来源于stack exchange，提问作者Steffen Roeber