如何低成本提取__m256i中16x16位矩阵的四条边？

要实现这个需求，我们可以避开高成本的全位转置操作，直接通过向量掩码、位提取和少量指令完成，核心思路如下：

• 顶部（第0行）和底部（第15行）直接提取向量的首尾16位元素
• 第0列和第15列通过_mm256_movemask_epi8提取对应位，再用位操作压缩成16位值

实现代码（依赖BMI2指令集，最优性能） #

#include <immintrin.h>
#include <x86intrin.h>

__m256i get_edges(__m256i m) {
    // 提取顶部（第0行）和底部（第15行）
    uint16_t top = _mm256_extract_epi16(m, 0);
    uint16_t bottom = _mm256_extract_epi16(m, 15);

    // 提取第15列：每个16位元素的最高位（bit15）
    int col15_mask = _mm256_movemask_epi8(m);
    // 提取掩码中的奇数位（对应每个16位元素的高字节bit7，即原bit15）
    uint16_t col15 = _pext_u32(col15_mask, 0xAAAAAAAA);

    // 提取第0列：每个16位元素的最低位（bit0）
    __m256i col0_bits = _mm256_and_si256(m, _mm256_set1_epi16(0x0001));
    col0_bits = _mm256_slli_epi16(col0_bits, 7); // 将bit0移到低字节的bit7位置
    int col0_mask = _mm256_movemask_epi8(col0_bits);
    // 提取掩码中的偶数位（对应每个低字节的bit7，即原bit0）
    uint16_t col0 = _pext_u32(col0_mask, 0x55555555);

    // 将四个值打包到__m256i的低四个16位通道（其余位置为0，可按需调整顺序）
    return _mm256_set_epi16(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, bottom, col15, top, col0);
}

无BMI2指令集的替代方案 #

如果无法使用BMI2的_pext_u32，可以用手动位操作提取对应位：

#include <immintrin.h>

// 提取掩码中的奇数位，对应第15列
uint16_t extract_col15(int mask) {
    uint16_t res = 0;
    for (int i = 0; i < 16; ++i) {
        res |= ((mask >> (1 + 2*i)) & 1) << i;
    }
    return res;
}

// 提取掩码中的偶数位，对应第0列
uint16_t extract_col0(int mask) {
    uint16_t res = 0;
    for (int i = 0; i < 16; ++i) {
        res |= ((mask >> (2*i)) & 1) << i;
    }
    return res;
}

__m256i get_edges_no_bmi2(__m256i m) {
    uint16_t top = _mm256_extract_epi16(m, 0);
    uint16_t bottom = _mm256_extract_epi16(m, 15);

    int col15_mask = _mm256_movemask_epi8(m);
    uint16_t col15 = extract_col15(col15_mask);

    __m256i col0_bits = _mm256_and_si256(m, _mm256_set1_epi16(0x0001));
    col0_bits = _mm256_slli_epi16(col0_bits, 7);
    int col0_mask = _mm256_movemask_epi8(col0_bits);
    uint16_t col0 = extract_col0(col0_mask);

    return _mm256_set_epi16(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, bottom, col15, top, col0);
}

方案优势 #

• 完全避免了全位转置的高成本操作，整体指令数大幅减少
• 核心操作仅用到向量掩码、移位、movemask和位提取，均为低延迟指令
• 输出位置可灵活调整，示例中优先将结果放在寄存器低半部分，符合需求偏好

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