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

Question

我正在考虑为 uint32_t 实现 8 元堆排序。为此，我需要一个函数来选择 8 元素向量中最大元素的索引，以便我可以将其与父元素进行比较，并有条件地执行交换和进一步的 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： 这不是链接问题的重复 - 请注意，我要求的是最大元素的 index，而不仅仅是元素值。

Answer 1

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

• 找到最大值（通常是几个移位/置换和最大值操作）
• 在第二个向量的所有 8 个元素上求最大值（可以与之前的操作结合使用）
• 比较原始向量和最大向量 (_mm256_cmpeq_epi32)
• 提取标量掩码 (_mm256_movemask_epi8)
• 将标量掩码转换为索引

这是我刚刚汇总的 AVX2 实现的第一次通过 - 我在 2.6 GHz Haswell 上对其进行了测试和基准测试，它以大约 1.7 ns/向量的速度运行（包括加载向量和存储结果索引）：

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;
}

Answer 2

在 n 路 SIMD 向量上执行水平操作（点积、求和、最大索引等）的最有效方法是一次执行 n 个操作，方法是转置它们并改用垂直操作。某些 SIMD 架构对水平操作有更好的支持，但总的来说，分块转置方法会更加灵活和高效。

Answer 3

使用Vc 库我会写：

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

使用内在函数应该是这样的（这是没有 AVX2 的 AVX - 使用 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