Реализовать triu_indices numpy с помощью avx в c ++ - PullRequest
Новая
       35

Реализовать triu_indices numpy с помощью avx в c ++

0 голосов
/ 25 мая 2018

Я хотел бы реализовать numpy.triu_indices (a, 1) (обратите внимание, что вторым аргументом является 1) в c ++ со встроенной функцией avx.Приведенный ниже фрагмент кода представляет собой невекторизованную версию кода, который я придумал.Здесь 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 (то есть в векторизации).В конечном счете, функция будет выполняться по всему массиву целых чисел, который будет передавать в функцию переменную a, а выходные массивы first и second будут храниться в двух родительских массивах.

Не могли бы вы дать мне несколько полезных советов (или фрагмент кода) о том, как векторизовать эту функцию с явными инстинктами AVX2 (то есть, не зависящими от автоматической векторизации компилятора)?Извините, если это нубский вопрос, так как я недавно начал изучать AVX.

1 Ответ

0 голосов
/ 25 мая 2018

Прежде всего, убедитесь, что вам действительно нужно это сделать, и вы действительно хотите, чтобы массивы индексов были конечным результатом, а не частью отслеживания данных в треугольной матрице.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), затем немного более короткий цикл следующего числа и т. д., т.е. легко сохранить целый ряд 0 s.Для всех строк, кроме последней пары, эти серии больше 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 для другогои сохраняя его в другом массиве.

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

...