如何使用非立即输入进行类似于 _mm_extract_epi8 的操作？

Posted 2023-02-16

【中文标题】如何使用非立即输入进行类似于 _mm_extract_epi8 的操作？

【英文标题】：How to make an operation similar to _mm_extract_epi8 with non-immediate input?

【发布时间】：2021-12-25 19:11:54

【问题描述】：

我想要的是使用可变标量索引从向量中提取一个值。

类似于_mm_extract_epi8 / _mm256_extract_epi8，但具有非即时输入。

（向量中有一些结果，找到给定索引的结果是真正的结果，其余的被丢弃）

【参考方案1】：

特别是，如果index 在 GPR 中，最简单的方法可能是将 val 存储到内存中，然后将 movzx 存储到另一个 GPR 中。使用 C 的示例实现：

uint8_t extract_epu8var(__m256i val, int index) 
    union 
        __m256i m256;
        uint8_t array[32];
     tmp;
    tmp.m256 = val;
    return tmp.array[index];

Godbolt 转换（请注意，堆栈对齐会产生大量开销——如果您没有对齐的临时存储区域，则可以只使用 vmovdqu 而不是 vmovdqa）：https://godbolt.org/z/Gj6Eadq9r

【讨论】：

这里不使用联合，您可以在本地 alignas(32) uint8_t array[32]; 上使用 _mm256_store_si256（或 storeu）。或alignas(__m256i)。一些编译器会选择过度对齐数组，即使你不要求它，当他们看到这个存储/重新加载时。我认为所有支持 Intel 内在函数的 C++ 编译器都允许联合类型双关语，尽管我记得听说 SunCC 没有。无论如何，至少在我看来，内部存储比工会成员分配更惯用且更容易识别。

但是是的，在 asm 中，存储/重新加载只有大约 6 个周期的延迟 IIRC（所以并不比在 Skylake 上的车道交叉洗牌加上 pshufb 差多少），而且在 uop 方面相当便宜吞吐量，尤其是对于前端。如果编译器想要将其优化为具有某些未来指令集或已知常数索引的随机播放，它可以。

【参考方案2】：

到目前为止，最好的选择似乎是将_mm_shuffle_epi8 用于 SSE

uint8_t extract_epu8var(__m128i val, int index) 
    return (uint8_t)_mm_cvtsi128_si32(
        _mm_shuffle_epi8(val, _mm_cvtsi32_si128(index)));

不幸的是，这不适用于 AVX。 vpshufb 不会跨车道移动。有一个跨车道洗牌_mm256_permutevar8x32_epi32，但结果似乎很复杂：

uint8_t extract_epu8var(__m256i val, int index) 
    int index_low = index & 0x3;
    int index_high = (index >> 2);
    return (uint8_t)(_mm256_cvtsi256_si32(_mm256_permutevar8x32_epi32(
      val, _mm256_zextsi128_si256(_mm_cvtsi32_si128(index_high))))
      >> (index_low << 3));