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

Question

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

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中完全实现这一点(这是一种矢量化的方式)。最终，该函数将运行在整个in数组上，该数组将向函数提供变量a，输出数组first和second将存储在两个父数组中。

请您给我一些有用的提示(或代码片段)，作为如何使用显式AVX2工具(也就是说，不依赖于编译器自动向量化)来向量化这个函数吗？对不起，如果这是一个菜鸟的问题，因为我已经开始学习AVX最近。

Answer 1

首先，确保您确实需要这样做，并且实际上希望索引数组作为最终结果，而不是作为在三角矩阵中跟踪数据的一部分。AVX2有集合支持，AVX512有分散支持，但是引入一系列索引会使SIMD更加糟糕。

对于三角矩阵上的循环，对于i，j到线性，请参见基于装配的三角形矩阵存储器寻址算法。(您可能希望取消索引，以便每行从32字节对齐边界开始。即将每一行的长度乘以8个浮点元素，这是一个完整的AVX向量。这也使得使用AVX向量循环矩阵变得更容易:您可以将垃圾存储到一行末尾的填充中，而不是让一行的最后一个向量包含下一行开始时的一些元素。)

对于线性-> i，j，闭形公式包括sqrt (也是C++版本)，因此数组查找可能很有用，但实际上您应该避免这样做。(例如，如果以填充格式在三角矩阵上循环，请跟踪在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}的向量开始，然后用{4,4,4,4}添加它，即_mm_set1_epi32(1)。对于256位AVX2向量，_mm256_set1_epi32(8)将8元矢量增量为8.

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

尽管行尾处理比手工处理要复杂得多，但auto-vectorize编译器已经可以很好地您的函数了。在戈德波特编译器浏览器上使用您的代码。 ( __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));

即，通过添加所有向量将{0,0,0,0,...}转换为{1,1,1,1,...}。并通过添加所有向量将{1,2,3,4,5,6,7,8}转换为{2,3,4,5,6,7,8,9}。