双YMM寄存器中10个uint32_t的有限移位实现优化咨询

我需要用两个256位YMM寄存器处理10个uint32_t值，当前存储方式是前8个值存入左侧寄存器，剩余2个左对齐存入右侧寄存器，形式如下：
[0, 1, 2, 3, 4, 5, 6, 7], [8, 9, x, x, x, x, x, x]
其中x为可忽略的无效值。我要实现一个移位操作，支持向左或向右最多移位2个位置：

• 左移1位结果：
  [1, 2, 3, 4, 5, 6, 7, 8], [9, Z, x, x, x, x, x, x]
• 右移2位结果：
  [Z, Z, 0, 1, 2, 3, 4, 5], [6, 7, x, x, x, x, x, x]
  标记为Z的位置需要填充0。以下是我当前的实现代码：

__m256i ind[5] = {
    _mm256_setr_epi32(2, 3, 4, 5, 6, 7, 0, 1),
    _mm256_setr_epi32(1, 2, 3, 4, 5, 6, 7, 0),
    _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7),
    _mm256_setr_epi32(7, 0, 1, 2, 3, 4, 5, 6),
    _mm256_setr_epi32(6, 7, 0, 1, 2, 3, 4, 5)
};

__m256i mask[5][2] = {
    {_mm256_setr_epi32( 0, 0, 0, 0, 0, 0, 1, 1), _mm256_setr_epi32(-1,-1,-1,-1,-1,-1,-1,-1)},
    {_mm256_setr_epi32( 0, 0, 0, 0, 0, 0, 0, 1), _mm256_setr_epi32( 1,-1,-1,-1,-1,-1,-1,-1)},
    {_mm256_setr_epi32( 0, 0, 0, 0, 0, 0, 0, 0), _mm256_setr_epi32(-1,-1,-1,-1,-1,-1,-1,-1)},
    {_mm256_setr_epi32(-1, 0, 0, 0, 0, 0, 0, 0), _mm256_setr_epi32( 0,-1,-1,-1,-1,-1,-1,-1)},
    {_mm256_setr_epi32(-1,-1, 0, 0, 0, 0, 0, 0), _mm256_setr_epi32( 0, 0,-1,-1,-1,-1,-1,-1)}
};

__m256i zeroes[5] = {
    _mm256_setr_epi32( 0, 0, 0, 0, 0, 0, 0, 0),
    _mm256_setr_epi32(-1, 0, 0, 0, 0, 0, 0, 0),
    _mm256_setr_epi32(-1,-1,-1,-1,-1,-1,-1,-1),
    _mm256_setr_epi32( 0,-1,-1,-1,-1,-1,-1,-1),
    _mm256_setr_epi32( 0, 0,-1,-1,-1,-1,-1,-1)
};

void shift(__m256i& ymml, __m256i& ymmr, int x){
    if (x == 0)
        return;
    
    __m256i resl = _mm256_permutevar8x32_epi32(ymml, ind[x+2]);
    __m256i resr = _mm256_permutevar8x32_epi32(ymmr, ind[x+2]);
    ymml = _mm256_blendv_epi8(resl, resr, mask[x+2][0]);
    ymmr = _mm256_blendv_epi8(resl, resr, mask[x+2][1]);

    if (x > 0) {
        ymml = _mm256_and_si256(ymml, zeroes[x+2]);
    } else {
        ymmr = _mm256_and_si256(ymmr, zeroes[x+2]);
    }
}

请问有没有更优的实现方式？

你的当前实现依赖预定义的置换、掩码数组，虽然逻辑清晰，但可以通过直接利用AVX2的移位/提取/插入指令减少内存开销，同时降低指令延迟，因为不需要加载预定义的常量数组。以下是针对固定移位范围（±1/±2）的优化实现：

优化后的代码 #

#include <immintrin.h>

void shift_opt(__m256i& ymml, __m256i& ymmr, int x) {
    if (x == 0) return;

    const __m256i zero = _mm256_setzero_si256();
    __m256i tmp_l, tmp_r;

    if (x == 1) { // 左移1位
        // 左寄存器左移1个uint32_t（4字节），空出的位置补右寄存器第一个元素
        tmp_l = _mm256_slli_si256(ymml, 4);
        const uint32_t r0 = _mm256_extract_epi32(ymmr, 0);
        tmp_l = _mm256_insert_epi32(tmp_l, r0, 7);
        // 右寄存器左移1个uint32_t，空出位置补0
        tmp_r = _mm256_slli_si256(ymmr, 4);
        tmp_r = _mm256_blend_epi32(tmp_r, zero, 0x01);
        ymml = tmp_l;
        ymmr = tmp_r;
    } else if (x == 2) { // 左移2位
        tmp_l = _mm256_slli_si256(ymml, 8);
        const uint32_t r0 = _mm256_extract_epi32(ymmr, 0);
        const uint32_t r1 = _mm256_extract_epi32(ymmr, 1);
        tmp_l = _mm256_insert_epi32(tmp_l, r0, 6);
        tmp_l = _mm256_insert_epi32(tmp_l, r1, 7);
        tmp_r = zero; // 右寄存器全部清零
        ymml = tmp_l;
        ymmr = tmp_r;
    } else if (x == -1) { // 右移1位
        // 右寄存器右移1个uint32_t，空出位置补左寄存器最后一个元素
        tmp_r = _mm256_srli_si256(ymmr, 4);
        const uint32_t l7 = _mm256_extract_epi32(ymml, 7);
        tmp_r = _mm256_insert_epi32(tmp_r, l7, 0);
        // 左寄存器右移1个uint32_t，空出位置补0
        tmp_l = _mm256_srli_si256(ymml, 4);
        tmp_l = _mm256_blend_epi32(tmp_l, zero, 0x80);
        ymml = tmp_l;
        ymmr = tmp_r;
    } else if (x == -2) { // 右移2位
        tmp_r = _mm256_srli_si256(ymmr, 8);
        const uint32_t l6 = _mm256_extract_epi32(ymml, 6);
        const uint32_t l7 = _mm256_extract_epi32(ymml, 7);
        tmp_r = _mm256_insert_epi32(tmp_r, l6, 0);
        tmp_r = _mm256_insert_epi32(tmp_r, l7, 1);
        // 左寄存器右移2个uint32_t，空出位置补0
        tmp_l = _mm256_srli_si256(ymml, 8);
        tmp_l = _mm256_blend_epi32(tmp_l, zero, 0xC0);
        ymml = tmp_l;
        ymmr = tmp_r;
    }
}

优化点说明 #

1. 消除内存依赖：去掉了预定义的常量数组，所有操作直接在寄存器内完成，避免了常量加载的额外开销。
2. 指令效率提升：使用_mm256_slli_si256/_mm256_srli_si256进行字节级移位，比permutevar8x32延迟更低；extract/insert_epi32为单周期指令，跨寄存器数据传递更直接。
3. 精准清零操作：用_mm256_blend_epi32直接清零需要填充0的位置，逻辑更直观且性能与and_si256相当。
4. 分支预测友好：移位范围固定为±1/±2，分支数量极少，现代CPU的分支预测可做到零开销。

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