Реализуйте triu_indices numpy с avx в C++

Прежде всего, убедитесь, что вам В самом деле нужно это сделать, и вам действительно нужны массивы индексов как конечный результат, а не как часть отслеживания данных в треугольной матрице. AVX2 имеет поддержку сборки, а AVX512 имеет поддержку разброса, но введение массива индексов значительно ухудшает SIMD.

Для перехода по треугольным матрицам и для i, j к линейному, см. алгоритм адресации памяти треугольных матриц с помощью сборки. (Возможно, вы захотите дополнить индексирование, чтобы каждая строка начиналась с 32-байтовой выровненной границы. То есть округлите длину каждой строки до числа, кратного 8 элементам с плавающей запятой, всего вектора AVX. Это также упрощает цикл по матрица с векторами AVX: вы можете хранить мусор в заполнении в конце строки вместо того, чтобы последний вектор строки включал некоторые элементы с начала следующей строки.)

Для linear -> i, j формула в закрытой форме включает sqrt (также Версия C++), поэтому вполне возможно, что поиск в массиве может быть полезен, но на самом деле вам следует избегать этого вообще. (например, если перебирать треугольную матрицу в упакованном формате, отслеживайте, где вы находитесь, в i, j, а также в линейном, чтобы вам не понадобился поиск, когда вы найдете элемент, который ищете.)

Если вам это нужно:

Для больших массивов он довольно хорошо разбивается на целые векторы, усложняется только на концах строк.

Вы можете использовать предопределенную векторную константу для особого случая последнего угла, когда у вас есть несколько строк треугольников в одном и том же векторе из 4 или 8 элементов int.

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

В случае большего треугольника мы сохраняем большие серии с одним и тем же номером (например, memset), затем несколько более короткую серию следующего номера и т. д., То есть хранить целый ряд 0 легко. Для всех строк, кроме последней пары, эти прогоны превышают 1 элемент вектора.

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

Опять же, в пределах одной строки это простой шаблон для векторизации. Чтобы сохранить возрастающую последовательность, начните с вектора {1,2,3,4} и увеличивайте его с помощью SIMD, добавьте с помощью {4,4,4,4}, то есть _mm_set1_epi32(1). Для 256-битных векторов AVX2 _mm256_set1_epi32(8) увеличивает 8-элементный вектор на 8.

Итак, в самом внутреннем цикле вы просто сохраняете один инвариантный вектор и используете _mm256_add_epi32 для другого и сохраняете его в другом массиве.

Компиляторы уже могут авто-векторизовать вашу функцию довольно прилично, хотя обработка конца строки намного сложнее, чем вы могли бы сделать вручную. С вашим кодом в проводнике компилятора Godbolt (с __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}, добавив вектор всех единиц.