找到具有最大值的元素的索引。使用 AVX512 指令的绝对值

Question

我是使用 AVX512 指令编码的新手。我的机器是 Intel KNL 7250。我正在尝试使用 AVX512 指令来查找具有最大绝对值的元素的索引，它是数组的双精度和大小 % 8 = 0。但它每次都会打印一个输出索引 = 0。不知道哪里有问题，请帮帮我。 另外，__m512i 类型如何使用 printf？

谢谢。

代码：

void main()
{
    int i;
    int N=160;
    double vec[N];

    for(i=0;i<N;i++)
    {
        vec[i]=(double)(-i) ;
        if(i==10)
        {
            vec[i] = -1127;
        }
    }

    int max = avxmax_(N, vec);
    printf("maxindex=%d\n", max);
}

int avxmax_(int   N,     double *X )
{
    // return the index of element having maximum absolute value.
    int maxindex, ix, i, M;
    register __m512i increment, indices, maxindices, maxvalues, absmax, max_values_tmp, abs_max_tmp, tmp;
    register __mmask8 gt;
    double values_X[8];
    double indices_X[8];
    double maxvalue;
    maxindex = 1;
    if( N == 1) return(maxindex);


    M = N % 8;
    if( M == 0)
    {
        increment =  _mm512_set1_epi64(8); // [8,8,8,8,8,8,8,8]
        indices = _mm512_setr_epi64(0, 1, 2, 3, 4, 5, 6, 7);
        maxindices = indices;
        maxvalues =  _mm512_loadu_si512(&X[0]);
        absmax = _mm512_abs_epi64(maxvalues); 

        for( i = 8; i < N; i += 8)
        {
            // advance scalar indices: indices[0] + 8, indices[1] + 8,...,indices[7] + 8
            indices = _mm512_add_epi64(indices, increment);

            // compare
            max_values_tmp  = _mm512_loadu_si512(&X[i]);
            abs_max_tmp = _mm512_abs_epi64(max_values_tmp);
            gt           = _mm512_cmpgt_epi64_mask(abs_max_tmp, absmax);

            // update
            maxindices = _mm512_mask_blend_epi64(gt, maxindices, indices);
            absmax     = _mm512_max_epi64(absmax, abs_max_tmp);
        }

        // scalar part
        _mm512_storeu_si512((__m512i*)values_X, absmax);
        _mm512_storeu_si512((__m512i*)indices_X, maxindices);
        maxindex = indices_X[0];
        maxvalue = values_X[0];
        for(i = 1; i < 8; i++)
        {
            if(values_X[i] > maxvalue)
            {
                maxvalue = values_X[i];
                maxindex = indices_X[i];
            }
            
        }
        return(maxindex);
    }
}

Answer 1

您的函数返回 0，因为您将 int64 索引视为 double 的位模式，并将该（微小）数字转换为整数。 double indices_X[8]; 是错误;应该是uint64_t。 还有其他错误，见下文。

如果您在使用变量时声明变量（C99 样式，而不是过时的 C89 样式），则更容易发现此错误。

你 _mm512_storeu_si512 将 int64_t 的向量索引到 double indices_X[8]，将其键入双关，然后在纯 C 中执行 int maxindex = indices_X[0];。这是隐式类型转换，将次正规的double 转换为整数。

（我注意到 asm https://godbolt.org/z/zsfc36 中有一个神秘的vcvttsd2si FP->int 转换，同时将代码转换为初始化器旁边的 C99 样式变量声明。这是一个线索：在这个函数。我注意到大约在同一时间我将double indices_X[8]; 声明向下移动到使用它的块中，并注意到它的类型为double。）

---

实际上可以在 FP 位模式上使用整数运算

但前提是你使用正确的！ IEEE754 指数偏差意味着编码/位模式可以作为符号/幅度整数进行比较。所以你可以做 abs / min / max 并对其进行比较，但当然不能进行整数加 / 减（除非您正在实施nextafter）。

_mm512_abs_epi64 是 2 的补码绝对值，不是符号幅度。相反，您必须屏蔽符号位。然后，您就可以将结果视为无符号整数或有符号 2 的补码。 （任何一个都有效，因为高位清晰。）

使用整数 max 有一个有趣的特性，即 NaN 将比较高于任何其他值，Inf 低于该值，然后是有限值。所以我们基本上免费获得了一个 NaN 传播的 max-finder。

在 KNL（Knight's Landing）上，FP vmaxpd 和 vcmppd 具有与其等效整数相同的性能：2 个周期延迟，0.5c 吞吐量。 (https://agner.org/optimize/)。所以你的方式在 KNL 上的优势为零，但对于主流英特尔来说，这是一个巧妙的技巧，比如 Skylake-X 和 IceLake。

---

修正优化版：

• 使用size_t 作为返回类型和循环计数器/索引来处理潜在的巨大数组，而不是int 和64 位向量元素的随机混合。 （uint64_t 用于收集水平最大值的临时数组：即使在具有 32 位指针/size_t 的构建中，它也始终是 64 位。）

• 修正：在 N==1 上返回 0，而不是 1：唯一元素的索引是 0。

• 修正：在 N%8 != 0 上返回 -1，而不是从非 void 函数的末尾脱落。 （如果调用者在 C 中使用结果，或者在 C++ 中执行结束时，则为未定义的行为。

• 修正：FP 值的绝对值 = 清除符号位，而不是位模式上的 2 的补码绝对值

• 某种错误修复：使用无符号整数比较和最大值，因此它适用于带有 _mm512_abs_epi64 的 2 的补码整数（这会产生无符号结果；请记住，如果您继续将其视为有符号，-LONG_MIN 会溢出到 LONG_MIN )。

• 风格改进：if (N%8 != 0) return -1;，而不是将大部分主体放在 if 块中。

• 风格改进：在第一次使用时声明变量，并删除了一些未使用的纯噪声变量。这是自 20 多年前标准化的 C99 以来 C 的惯用语。

• 样式改进：对仅保存加载结果的 tmp 向量变量使用更简单的名称。有时您只需要一个 tmp var，因为内在名称太长以至于您不想键入 _mm...load... 作为另一个内在函数的 arg。像v 这样的名称限定为内部循环是一个明确的标志，它只是一个占位符，以后不使用。 （当您在初始化时声明它时，这种样式效果最好，因此很容易看出它不能在外部范围中使用。）

• 优化：使用 SIMD 循环后减少 8 -> 4 个元素：提取高半部分，与现有的低半部分结合。 （与您对 sum 或 max 等更简单的水平缩减所做的相同）。当我们需要只有AVX512有的指令，而KNL没有AVX512VL的时候，很不方便，所以我们必须使用512位版本，忽略高垃圾。但是 KNL 确实有 AVX1 / AVX2，所以我们仍然可以存储 256 位向量并做一些事情。

  使用合并掩码 _mm512_mask_extracti64x4_epi64 提取直接混合同一向量的高半部分和低半部分是一个很酷的技巧，如果您使用 512 位掩码混合，编译器不会发现它。 :P

• 某种错误修复：在 C 中，main 在托管实现（在操作系统下运行）中的返回类型为 int。

#include <immintrin.h>
#include <stdio.h> 
#include <stdint.h>
#include <stdlib.h>

// bugfix: indices can be larger than an int
size_t avxmax_(size_t N,  double *X )
{
    // return the index of element having maximum absolute value.
    if( N == 1)
        return 0;    // bugfix: 0 is the only valid element in this case, not 1
    if( N % 8 != 0)      // [[unlikely]] // C++20
        return -1;   // bugfix: don't fall off the end of the function in this case

    const __m512i fp_absmask = _mm512_set1_epi64(0x7FFFFFFFFFFFFFFF);
    __m512i indices = _mm512_setr_epi64(0, 1, 2, 3, 4, 5, 6, 7);
    __m512i maxindices = indices;
    __m512i v =  _mm512_loadu_si512(&X[0]);
    __m512i absmax = _mm512_and_si512(v, fp_absmask);
    for(size_t i = 8; i < N; i += 8) // [[likely]]  // C++20
    {
        // advance indices by 8 each.
        indices = _mm512_add_epi64(indices, _mm512_set1_epi64(8));
        // compare
                v    = _mm512_loadu_si512(&X[i]);
        __m512i vabs = _mm512_and_si512(v, fp_absmask);
             // vabs = _mm512_abs_epi64(max_values_tmp);  // for actual integers, not FP bit patterns
        __mmask8 gt  = _mm512_cmpgt_epu64_mask(vabs, absmax);
        // update
        maxindices = _mm512_mask_blend_epi64(gt, maxindices, indices);
        absmax     = _mm512_max_epu64(absmax, vabs);
    }

    // reduce 512->256; KNL doesn't have AVX512VL so some ops require 512-bit vectors
    __m256i absmax_hi = _mm512_extracti64x4_epi64(absmax, 1);
    __m512i absmax_hi512 = _mm512_castsi256_si512(absmax_hi); // free
    __mmask8 gt = _mm512_cmpgt_epu64_mask(absmax_hi512, absmax);
    __m256i abs256 = _mm512_castsi512_si256(_mm512_max_epu64(absmax_hi512, absmax));  // reduced to low 4 elements

    // extract with merge-masking = blend
    __m256i maxindices256 = _mm512_mask_extracti64x4_epi64(
                           _mm512_castsi512_si256(maxindices), gt, maxindices, 1);

    // scalar part
    double values_X[4];
    uint64_t indices_X[4];
    _mm256_storeu_si256((__m256i*)values_X, abs256);
    _mm256_storeu_si256((__m256i*)indices_X, maxindices256);

    size_t maxindex = indices_X[0];
    double  maxvalue = values_X[0];
    for(int i = 1; i < 4; i++)
    {
        if(values_X[i] > maxvalue)
        {
            maxvalue = values_X[i];
            maxindex = indices_X[i];
        }
        
    }
    return maxindex;
}

On Godbolt：GCC10.2 -O3 -march=knl 的主循环是8条指令。因此，即使（最好的情况）KNL 可以解码并以 2/clock 运行它，每个向量仍然需要 4 个周期。您可以在 Godbolt 上运行该程序；它在 Skylake-X 服务器上运行，因此可以运行 AVX512 代码。你可以看到它打印出10。

.L4:
        vpandd  zmm2, zmm5, ZMMWORD PTR [rsi+rax*8]   # load, folded into AND
        add     rax, 8
        vpcmpuq k1, zmm2, zmm0, 6
        vpaddq  zmm1, zmm1, zmm4                    # increment current indices
        cmp     rdi, rax
        vmovdqa64       zmm3{k1}, zmm1              # blend maxidx using merge-masking
        vpmaxuq zmm0, zmm0, zmm2
        ja      .L4

        vmovapd zmm1, zmm3                        # silly missed optimization related to the case where the loop runs 0 times.
.L3:
        vextracti64x4   ymm2, zmm0, 0x1           # high half of absmax
        vpcmpuq k1, zmm2, zmm0, 6                 # compare high and low
        vpmaxuq zmm0, zmm0, zmm2
      #  vunpckhpd       xmm2, xmm0, xmm0  # setting up for unrolled scalar loop
        vextracti64x4   ymm1{k1}, zmm3, 0x1       # masked extract of indices

循环的另一个选项是屏蔽vpbroadcastq zmm3{k1}, rax，在循环之后添加[0..7] 每个元素的偏移量。这实际上会将vpaddq 保存在循环中，如果GCC 无论如何都要使用索引寻址模式，我们在寄存器中有正确的i。 （这在 Skylake-X 上不好；破坏了内存源 vpandd 的微融合。）

Agner Fog 没有列出 GP->vector 广播的性能，但希望它至少在 KNL 上只是单 uop。 （并且https://uops.info/ 没有 KNL 或 KNM 结果）。

---

分支策略：当一个新的最大值非常罕见时

如果您希望找到一个新的最大值非常罕见（例如，数组很大且分布均匀，或者至少没有向上趋势），则广播当前最大值并在找到任何更大的向量元素时分支可能会更快。

找到一个新的最大值意味着跳出循环（这可能会预测错误，所以这很慢）并广播该元素（可能使用tzcnt 来查找元素索引，然后是广播加载，并更新索引） .

特别是使用 KNL 的 4 路 SMT 来隐藏分支未命中成本，这可能是大型阵列的整体吞吐量胜利；平均每个元素的指令更少。

但对于确实呈上升趋势的输入可能会更糟，因此平均会找到 O(n) 次的新最大值，而不是 sqrt(n) 或 log(n) 或任何统一的频率分发会给我们。

---

PS：打印向量，存储到数组并重新加载元素。 print a __m128i variable

或者使用调试器向您展示他们的元素。