使用SIMD指令，有没有一种高效的方法获取SIMD寄存器中第一个非零元素？

Question

使用SIMD指令，有没有一种高效的方法获取SIMD寄存器中第一个非零元素？

x86bit-manipulationsimdintrinsicsavx

10

如题所述，如果一个256位SIMD寄存器是：

0 | 0 | 1 | 0 | 0 | 1 | 0 | 0 |

如何高效地获取第一个非零元素的索引（即第一个1的索引2）？最直接的方法是逐个存储到内存并检查，但这可能会代价太大。是否有什么巧妙的想法可以解决这个问题？

- MarZzz

2个回答

2

最近我一直在编写“获取X的索引”SIMD算法。到目前为止，从比较掩码中提取索引的最通用方法是通过水平索引最小值。

这里是（无符号）整数的水平最小值：

int horizontal_min(__m128i x) {
    x = _mm_min_epu32(x, _mm_shuffle_epi32(x, 0b01001110));
    x = _mm_min_epu32(x, _mm_shuffle_epi32(x, 0b11100001));
    return _mm_extract_epi32(x,0);
}

现在请按以下步骤操作：

int equal_first(__m128i x, __m128i y) {
    const __m128i index = _mm_set_epi32(0,1,2,3);
    // Compute mask
    __m128i mask = _mm_cmpeq_epi32(x,y);
    // Select indices.
    mask = _mm_blendv_epi8(_mm_set1_epi32(-1), index, mask);
    // mask = index | (~mask);
    // pick smallest indice.
    return horizontal_min(mask);
}

这段代码的优点在于不需要任何位扫描指令，所有操作都在FPU上完成。

提示：如果使用phminposuw128指令计算最小索引，则在使用16位索引时效率非常高。

编辑：Peter的分析表明，除非需要在SIMD寄存器中获取结果，否则我的解决方案较慢。

另一种情况是减少循环中想要数组中某个元素的索引。在循环中，您可以在SIMD寄存器中累加例如最小/最大元素的索引。现在无序的索引可能指向源数组中的任何位置。现在您必须使用horizontal_min()来确定最小/最大元素的位置。

- JATothrim

blendv在大多数CPU上占用2个uops。更好的选择是使用带有比较的_mm_or_si128，其中在要保留的元素中为false（0），在要拒绝的元素中为true（0xFFFFFFFF）。就像在这种情况下寻找任何非零元素一样，_mm_cmpeq_epi32(x，_mm_setzero_si128()); _mm_or_si128(index，mask); 然后您也可以放弃set1(4)向量常量。（它本来就可以是-1，最高的无符号数字，这对于编译器来说更便宜，可以使用pcmpeqd xmm0，xmm0实现。）pminud也是SSE4.1，因此避免blendv不能仅使用SSE2完成此工作。 - Peter Cordes

无论如何，对于像dword或qword这样的宽元素来说，这是一个有趣的想法，其中水平缩减不需要太多的洗牌/最小步骤，但仍然不够快。我不明白避免位扫描有什么好处，除非你实际上想要将结果作为向量，例如作为vpermilps的洗牌控制。 - Peter Cordes

如果您的 equal_first 函数支持 tzcnt，那么可以使用 cmp/movemask/bit-scan 在 3 条单uop指令内实现（https://godbolt.org/z/Ks46E81Gh）；否则，如果它可能在只有 bsf 的 CPU 上运行，则需要使用4个指令。当一个数上的低位都未被设置时，__builtin_ctz(bitmask | (1 << 4)) 将返回 4。 bsf 在 AMD CPU 上速度较慢，但是 rep bsf 可以在支持它的 CPU 上作为 tzcnt 运行，在不支持它的 CPU 上作为 bsf 运行，并且对于非零输入给出相同的结果。因此，使用 OR 技巧确保它始终为非零值。经过优化后，使用 set1(-1)，你的函数需要8条指令和9个uops。 - Peter Cordes

对于吞吐量而言，我的表现要好得多（前端和后端uops更少）。对于延迟，你的SIMD指令有7个周期的延迟，加上movd大约是2或3个周期的延迟，就像pmovmskb一样。因此大约需要9个周期。我的是1（pcmpeqd）+ 3（pmovmskb）+ 1（or）+ Intel上的3（tzcnt / bsf）= Intel上的8个周期。或者在AMD Zen上为7个周期，其中tzcnt是2个uop，2c延迟（可能是位反转和lzcnt）。（在AMD上，你的速度也快一个周期，其中vpblendvb是单uop 1c延迟。） - Peter Cordes

@PeterCordes 分析得很好！是的，这段代码可能对AMD处理器有偏向性，因为这就是我所拥有的。 - JATothrim

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Peter Cordes · Accepted Answer

使用PCMPEQB/W/D/Q指令与一个全为0的寄存器进行比较，得到一个向量，其中元素都是对应位置为0时为1，对应位置为1时为0。
使用PMOVMSKB指令将这个向量转换成整型位掩码。（或者使用movmskps或pd指令，每个dword或qword得到1个位而不是每个byte，如果这样做可以使您的位扫描 -> 索引计算更加高效，例如如果您想要一个元素偏移而不是字节偏移。）
取反（使用C中的~运算符，汇编中使用NOT指令），得到位图中非零元素的值为1。
使用TZCNT或BSF指令查找第一个（最低）置位的位。注意，如果BSF的输入全部为0，则会出现问题。但幸运的是，当输入为~bitmask时，这不是一个问题——高16个零位变成了1。（使用vpmovmskb ymm填充一个完整的uint32_t的AVX2版本可能需要使用~(uint64_t)bitmask，或者只使用tzcnt，因为AVX2 CPU也具有BMI1。）

例如使用内部函数：

int first_nonzero_byte(__m128i v){
    //__m128i v = _mm_loadu_si128((const __m128i*)p);  // for a pointer arg
    __m128i vcmp = _mm_cmpeq_epi8(v, _mm_setzero_si128());
    unsigned bitmask = _mm_movemask_epi8(vcmp);
#ifdef __GNUC__
    return __builtin_ctz(~bitmask);
#else
    return _tzcnt_u32( ~bitmask );
#endif
   // returns 16 if v was all zero so ~bitmask is 0xFFFF0000
}

在https://godbolt.org/z/Y8vYbsW69上编译为

# GCC11.2 -O3 -msse4.1
        movdqa  xmm1, xmm0      # missed optimization, should zero XMM1 instead
        pxor    xmm0, xmm0
        pcmpeqb xmm0, xmm1
        pmovmskb        eax, xmm0
        not     eax
        rep bsf eax, eax        # tzcnt on new CPUs, BSF on old
        ret

在GNU C中，如果没有-march=haswell或其他编译选项，_tzcnt_u32将无法编译，我们使用__builtin_ctz。正如我所说，~bitmask保证是非零的。tzcnt被编码为rep bsf；旧的CPU将执行它作为bsf，对于非零输入产生相同的结果。新的CPU将执行它作为tzcnt，在AMD上更有效（2个uops而不是7个）。英特尔将执行单个uop。如果您没有告诉GCC要针对特定的CPU进行优化，它将使用rep bsf，也就是tzcnt。

对于像JATothrim的答案中所示的相关函数，仅使用4个单uop指令（实际上在AMD上tzcnt只需要2个uop）而不是8个指令，包括一个pblendvb（Intel上需要2个uop），可以获得更好的效果。当您实际想要一个标量int时，该答案中的洗牌/水平缩减思路可能有用，作为vpermilps的洗牌控制向量的元素索引，但似乎不太优化。

int equal_first_dword_bitscan(__m128i x, __m128i y)
{
    __m128i vcmp = _mm_cmpeq_epi32(x,y);
    unsigned bitmask = _mm_movemask_ps(_mm_castsi128_ps(vcmp));
    bitmask |= 1<<4;    // return 4 if the low 4 bits are all 0
#ifdef __GNUC__
    return __builtin_ctz(bitmask);
#else
    return  _tzcnt_u32( bitmask );  // runs as BSF on old CPUs, don't skip the OR
#endif
}

MSVC没有__builtin_ctz，但即使您没有告诉它目标CPU支持BMI1，它也会编译_tzcnt_u32。如果您确定只在支持BMI1的CPU上运行，则可以省略bitmask |= 1 << 4;，以便返回32表示未找到。

如果您在多个函数中使用尾随零计数，则最好将ifdef代码包装在一个帮助函数中，而不是在每个用例中都进行。

如果只有一个可能的非零值（例如1），则对该向量进行PCMPEQB操作，这样您就不需要在之后反转它。

如果是这种情况，请考虑将数据存储在位图中，以减少缓存占用量8倍。然后，只需对数组的64位块进行TZCNT操作。

或者对于更大的数据数组，使用SIMD搜索第一个非零向量，然后TZCNT其第一个非零元素，如果您希望在第一个设置位之前存在多个qword的零。就像memcmp查找不匹配字节位置一样。
请参见Efficiently find least significant set bit in a large array?和How to find the first nonzero in an array efficiently?。

顺便提一下，汇编指令参考手册在每个条目底部列出了相关的C内嵌函数，您可以通过汇编助记符搜索Intel的内嵌函数查找器。(有关链接，请参见x86标签维基)。