XMM/YMM/ZMM 中最小或最大元素索引的位旋转魔法

Question

是否有指令或有效的无分支指令序列来计算无序（未排序）ZMM 的最大（或最小）元素的索引（而不是其值）？

数据类型并不重要 - 我更感兴趣的是知道是否有已建立的使用模式。

---

与已知解决方案相关的问题是，对于严格排序的 ZMM，可以使用 CMPPS、MOVMSKPS 和 TZCNT 来获取外部元素适合此列表的位置的索引（即 BSEARCH）

Answer 1

在整个向量上广播最小（或最大）元素，比较向量是否相等，使用 movemask 指令转换为位图，然后计算位图中的尾随零。

SSE 向量中 FP32 通道的示例：

uint32_t minpos_ps( __m128 vec )
{
    // Broadcast minimum value in the vector with a few shuffles
    __m128 i = _mm_min_ps( vec, _mm_permute_ps( vec, _MM_SHUFFLE( 1, 0, 3, 2 ) ) );
    i = _mm_min_ps( i, _mm_permute_ps( i, _MM_SHUFFLE( 2, 3, 0, 1 ) ) );

    // Compare lanes for equality with the minimum
    uint32_t mask = (uint32_t)_mm_movemask_ps( _mm_cmpeq_ps( vec, i ) );

    // Return index of the smallest set bit in the mask
    return std::countr_zero( mask );
}

更复杂的示例，对于 32 字节 AVX 向量中的无符号字节：

uint32_t minpos_epu8( __m256i vec )
{
    __m256i i = _mm256_min_epu8( vec, _mm256_permute2x128_si256( vec, vec, 1 ) );
    i = _mm256_min_epu8( i, _mm256_shuffle_epi32( i, _MM_SHUFFLE( 1, 0, 3, 2 ) ) );
    i = _mm256_min_epu8( i, _mm256_shuffle_epi32( i, _MM_SHUFFLE( 2, 3, 0, 1 ) ) );
    // If you calling this in a loop where compiler can preload constant vectors,
    // replace shuffles and shifts below with _mm256_shuffle_epi8
    __m256i tmp = _mm256_shufflehi_epi16( i, _MM_SHUFFLE( 2, 3, 0, 1 ) );
    tmp = _mm256_shufflelo_epi16( tmp, _MM_SHUFFLE( 2, 3, 0, 1 ) );
    i = _mm256_min_epu8( i, tmp );
    tmp = _mm256_or_si256( _mm256_slli_epi16( i, 8 ), _mm256_srli_epi16( i, 8 ) );
    i = _mm256_min_epu8( i, tmp );

    uint32_t mask = (uint32_t)_mm256_movemask_epi8( _mm256_cmpeq_epi8( vec, i ) );

    return std::countr_zero( mask );
}

标准库函数std::countr_zero需要C++/20。

如果您还没有该版本，请根据编译器和目标平台替换为_tzcnt_u32、_BitScanForward或内在函数。__builtin_ctz