在 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 的线性化，请参见。（您可能想要填充索引，以便每一行都从 32 字节对齐的边界开始。即将每一行的长度四舍五入为 8 个浮点元素的倍数，一个完整的 AVX 向量。这也使得循环更容易带有 AVX 向量的矩阵：您可以将垃圾存储到行末尾的填充中，而不是让行的最后一个向量包含下一行开头的一些元素。）

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

如果你确实需要这个：

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

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

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

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

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

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

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

编译器已经可以自动-相当不错地向量化你的函数，尽管行尾处理比你可以用手做。 With your code on the Godbolt compiler explorer（__restrict 告诉编译器输出数组不重叠，__builtin_assume_aligned 告诉编译器它们是对齐的），我们得到一个这样的内部循环（来自gcc):

.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)

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

计算下一个外循环迭代的起始向量可以用这样的方法来完成：

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

即通过添加全 1 向量将 {0,0,0,0,...} 变成 {1,1,1,1,...}。并通过添加全 1 向量将 {1,2,3,4,5,6,7,8} 变为 {2,3,4,5,6,7,8,9}。

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

Implement numpy's triu_indices with avx in c++

c++

sse

numpy

vectorization

avx

如果你确实需要这个：