c++ - 在 C++ 中用 avx 实现 numpy 的 triu_indices

Question

我想在 c++ 中使用 avx 内在函数实现 numpy.triu_indices(a, 1)(注意第二个参数是 1)。下面的代码片段是我想出的代码的非矢量化版本。这里，a是长度(int)，first和second是两个输出数组

void triu_indices(int a, int* first, int* second)
{
    int index = 0;
    for (int i=0; i < a; i++)
    {

        for(int j = i+1; j < a; j++)
        {
            first[index] = i;
            second[index] = j;
            index++;
        }

    }
}

作为示例输出，如果我给 a=4，则

first = [0,0,0,1,1,2]
second = [1,2,3,2,3,3]

现在，我想在 AVX2 中完全实现它(即以矢量化方式)。最终，该函数将运行整个整数数组，该数组将为函数提供变量 a，并输出数组 first 和 second 将存储在两个父数组中。

能否请您给我一些有用的提示(或代码片段)，说明如何使用显式 AVX2 内参(即，不依赖于编译器自动矢量化)对该函数进行矢量化？对不起，如果这是一个菜鸟问题，因为我最近开始学习 AVX。

Answer 1

首先，确保您确实需要这样做，并且确实希望将索引数组作为最终结果，而不是作为跟踪三角矩阵中数据的一部分。 AVX2 支持 gather，AVX512 支持 scatter，但引入索引数组让 SIMD 变得更糟。

对于三角矩阵的循环，以及 i,j 到线性，参见 algorithm of addressing a triangle matrices memory using assembly . (您可能想要填充索引，以便每一行都从 32 字节对齐的边界开始。即将每一行的长度四舍五入为 8 个浮点元素的倍数，一个完整的 AVX vector 。这也使得循环更容易带有 AVX vector 的矩阵:您可以将垃圾存储到行末尾的填充中，而不是让行的最后一个 vector 包含下一行开头的一些元素。)

对于线性 -> i,j, the closed-form formula包括 sqrt (也是一个 C++ version )，所以数组查找可能会有用，但实际上你应该完全避免这样做。 (例如，如果以打包格式遍历三角矩阵，请跟踪您在 i、j 和线性矩阵中的位置，以便在找到要查找的元素时不需要查找。)

---

如果你确实需要这个:

对于大型数组，它可以很好地分解成整个 vector ，仅在行的末尾变得棘手。

对于最后一个角的特殊情况，您可以使用预定义的 vector 常量，其中在 4 或 8 个 int 元素的相同 vector 中有多个三角形行。

first = [0,0,0,1,1,2]

对于较大的三角形，我们存储相同数字的大量运行(如 memset)，然后是下一个数字的稍短运行，等等。即存储整行 0s 很简单。对于除最后几行以外的所有行，这些运行都大于 1 个 vector 元素。

second = [1,2,3,2,3,3]

同样，在一行中，它是一个简单的矢量化模式。要存储递增序列，从 {1,2,3,4} 的 vector 开始并使用 SIMD 添加 {4,4,4,4} 来增加它，即 _mm_set1_epi32(1)。对于 256 位 AVX2 vector ，_mm256_set1_epi32(8) 将 8 元素 vector 递增 8。

因此，在最内层的循环中，您只需存储一个不变 vector ，然后在另一个 vector 上使用 _mm256_add_epi32 并将其存储到另一个数组。

编译器已经可以自动-相当不错地向量化您的函数，尽管行尾处理比您手动处理要复杂得多。 With your code on the Godbolt compiler explorer (使用 __restrict 告诉编译器输出数组不重叠，使用 __builtin_assume_aligned 告诉编译器它们是对齐的)，我们得到这样的内部循环(来自海湾合作委员会):

.L4:                                       # do {
    movups  XMMWORD PTR [rcx+rax], xmm0      # _mm_store_si128(&second[index], xmm0)
    paddd   xmm0, xmm2                       # _mm_add_epi32
    movups  XMMWORD PTR [r10+rax], xmm1      # _mm_store_si128(&second[index], counter_vec)
    add     rax, 16                          # index += 4  (16 bytes)
    cmp     rax, r9
    jne     .L4                            # }while(index != end_row)

如果我有时间，我可能会写得更详细，包括更好地处理行尾。例如在行尾结束的部分重叠存储通常是好的。或者展开外循环，使内循环具有重复的清理模式。

计算下一个外循环迭代的起始 vector 可以通过类似的方式完成:

vsecond = _mm256_add_epi32(vsecond, _mm256_set1_epi32(1));
vfirst  = _mm256_add_epi32(vfirst, _mm256_set1_epi32(1));

即通过添加全 1 vector 将 {0,0,0,0,...} 变为 {1,1,1,1,...}。并将 {1,2,3,4,5,6,7,8} 转换为 {2,3,4,5,6,7,8,9} 添加一个全部为 1 的 vector 。