查找x86 SIMD向量中最大元素的索引

Find index of maximum element in x86 SIMD vector

本文关键字：元素 索引 x86 SIMD 向量 查找      更新时间：2023-10-16

我正在考虑为uint32_t的实现8-ary堆排序。要做到这一点，我需要一个函数来选择8个元素向量中最大元素的索引，以便我可以将其与父元素进行比较，并有条件地执行swap和进一步的siftDown步骤。

(8 uint32_ts可以更改为16 uint32_ts或8 uint64_t或任何x86 SIMD可以有效支持)。

我有一些关于如何做到这一点的想法，但我正在寻找比非矢量化代码更快的东西，特别是我正在寻找能够让我做快速堆排序的东西。

我有clang++ 3.3和Core i7-4670，所以我应该能够使用甚至是最新的x86 SIMD的东西。

(顺便说一句:这是一个更大的项目的一部分:https://github.com/tarsa/SortingAlgorithmsBenchmark，例如，第四元堆排序，所以在实现SIMD堆排序后，我可以立即比较它们)

重复- 问题是:在x86 SIMD向量中计算最大元素索引的最有效方法是什么?

PS:它不是链接问题的副本-请注意，我要求的是一个最大元素的索引，而不仅仅是元素值。

水平操作对于SIMD来说是个坏消息，特别是对于AVX，其中大多数256位指令实际上被分解为两个独立的128位操作。话虽如此，如果你真的必须在8个元素之间设置一个32位的水平最大值那么我认为一般的方法是:

  • 查找最大值(通常是几个shift/permute和max操作)
  • 在第二个矢量的所有8个元素上绘制最大值(可以与之前的操作结合)
  • 比较原始矢量和最大矢量(_mm256_cmpeq_epi32)
  • 提取标量掩码(_mm256_movemask_epi8)
  • 将标量掩码转换为索引

这是我刚刚放在一起的AVX2实现的第一次通过-我对它进行了测试，并在2.6 GHz的Haswell上对其进行了基准测试，它的运行速度约为1.7 ns/vector(包括加载vector和存储结果索引):

uint8_t _mm256_hmax_index(const __m256i v)
{
    __m256i vmax = v;
    vmax = _mm256_max_epu32(vmax, _mm256_alignr_epi8(vmax, vmax, 4));
    vmax = _mm256_max_epu32(vmax, _mm256_alignr_epi8(vmax, vmax, 8));
    vmax = _mm256_max_epu32(vmax, _mm256_permute2x128_si256(vmax, vmax, 0x01));
    __m256i vcmp = _mm256_cmpeq_epi32(v, vmax);
    uint32_t mask = _mm256_movemask_epi8(vcmp);
    return __builtin_ctz(mask) >> 2;
}

在n路SIMD向量上执行水平操作(点积，求和，max-index等)的最有效方法是一次执行n个操作，通过将它们转置并使用垂直操作来代替。某些SIMD架构对水平操作有更好的支持，但一般来说，块转置方法将更加灵活和高效。

使用Vc库我会写:

size_t maximumIndex(Vc::uint_v vec) {
  const unsigned int max = vec.max();
  return (max == vec).firstOne();
}

对于intrinsic，它应该是沿着这些行(这是AVX没有AVX2 - AVX2变得稍微容易):

size_t maximumIndex(_mm256i vec) {
  __m128i lo = _mm256_castsi256_si128(vec);
  __m128i hi = _mm256_extractf128_si256(vec, 1);
  __m128i tmp = _mm_max_epu32(lo, hi);
  tmp = _mm_max_epu32(tmp, _mm_shuffle_epi32(tmp, _MM_SHUFFLE(1, 0, 3, 2)));
  tmp = _mm_max_epu32(tmp, _mm_shufflelo_epi16(tmp, _MM_SHUFFLE(1, 0, 3, 2))); // using lo_epi16 for speed here
  const int max = _mm_cvtsi128_si32(tmp);
  tmp = _mm_packs_epi16(_mm_packs_epi32(_mm_cmpeq_epi32(_mm_set1_epi32(max), lo),
                                        _mm_cmpeq_epi32(_mm_set1_epi32(max), hi)),
                        _mm_setzero_si128());
  return _bit_scan_forward(_mm_movemask_epi8(tmp));
}

顺便说一句，如果你想从SIMDized合并排序中得到一些灵感，看看这里:http://code.compeng.uni-frankfurt.de/projects/vc/repository/revisions/master/entry/src/avx_sorthelper.cpp

相关文章：