Параллельная префикс (суммарная) сумма с SSE

Я ищу несколько советов о том, как сделать параллельную префиксную сумму с SSE. Я заинтересован в том, чтобы делать это в массиве int, float или doubleles.

Я придумал два решения. Частный случай и общий случай. В обоих случаях решение проходит через массив через два прохода параллельно с OpenMP. Для специального случая я использую SSE на обоих проходах. В общем случае я использую его только на втором проходе.

Мой главный вопрос: как я могу использовать SSE на первом проходе в общем случае? Следующая ссылка simd-prefix-sum-on-intel-cpu показывает улучшение для байтов, но не для 32-битных типов данных.

Причина, по которой особый случай называется special, заключается в том, что он требует, чтобы массив находился в специальном формате. Например, допустим, что было только 16 элементов массива a для float. Затем, если массив был перестроен следующим образом (массив структур для структуры массивов):

a[0] a[1] ...a[15] -> a[0] a[4] a[8] a[12] a[1] a[5] a[9] a[13]...a[3] a[7] a[11] a[15]

На обоих проходах можно использовать вертикальные суммы SSE. Однако это было бы эффективно только в том случае, если массивы уже были в специальном формате, и выход можно было использовать в специальном формате. В противном случае дорогостоящий перерасчет должен был бы выполняться как на входе, так и на выходе, что сделало бы его намного медленнее, чем в общем случае.

Возможно, мне стоит рассмотреть другой алгоритм для суммы префикса (например, двоичное дерево)?

Код для общего случая:

void prefix_sum_omp_sse(double a[], double s[], int n) {
    double *suma;
    #pragma omp parallel
    {
        const int ithread = omp_get_thread_num();
        const int nthreads = omp_get_num_threads();
        #pragma omp single
        {
            suma = new double[nthreads + 1];
            suma[0] = 0;
        }
        double sum = 0;
        #pragma omp for schedule(static) nowait //first parallel pass
        for (int i = 0; i<n; i++) {
            sum += a[i];
            s[i] = sum;
        }
        suma[ithread + 1] = sum;
        #pragma omp barrier
        #pragma omp single
        {
            double tmp = 0;
            for (int i = 0; i<(nthreads + 1); i++) {
                tmp += suma[i];
                suma[i] = tmp;
            }
        }
        __m128d offset = _mm_set1_pd(suma[ithread]);
        #pragma omp for schedule(static) //second parallel pass with SSE as well
        for (int i = 0; i<n/4; i++) {       
            __m128d tmp1 = _mm_load_pd(&s[4*i]);
            tmp1 = _mm_add_pd(tmp1, offset);    
            __m128d tmp2 = _mm_load_pd(&s[4*i+2]);
            tmp2 = _mm_add_pd(tmp2, offset);
            _mm_store_pd(&s[4*i], tmp1);
            _mm_store_pd(&s[4*i+2], tmp2);
        }
    }
    delete[] suma;
}

Ответы

Ответ 1

Это первый раз, когда я отвечаю на свой вопрос, но это кажется уместным. Основано на hirschhornsalz ответ на сумму префикса на 16 байт simd-prefix-sum-on-intel-cpu Я придумал решение для использования SIMD на первом проходе для 4, 8 и 16 32-битных слов.

Общая теория гласит следующее. Для последовательного сканирования слов n требуется n добавления (n-1 для сканирования n слов и еще одно добавление, перенесенное из предыдущего набора сканированных слов). Однако использование SIMD n слов может быть отсканировано в добавлениях log 2 (n) и равном числе сдвигов плюс еще одно добавление и широковещание для переноса из предыдущего SIMD-сканирования. Поэтому при некотором значении n метод SIMD победит.

Посмотрите на 32-битные слова с SSE, AVX и AVX-512:

4 32-bit words (SSE):      2 shifts, 3 adds, 1 broadcast       sequential: 4 adds
8 32-bit words (AVX):      3 shifts, 4 adds, 1 broadcast       sequential: 8 adds
16 32 bit-words (AVX-512): 4 shifts, 5 adds, 1 broadcast       sequential: 16 adds

Исходя из этого, SIMD не будет полезен для сканирования 32-разрядных слов до AVX-512. Это также предполагает, что сдвиги и широковещание могут выполняться только с одной инструкцией. Это справедливо для SSE, но не для AVX и, возможно, даже для AVX2.

В любом случае я собрал некоторый рабочий и проверенный код, который использует префиксную сумму, используя SSE.

inline __m128 scan_SSE(__m128 x) {
    x = _mm_add_ps(x, _mm_castsi128_ps(_mm_slli_si128(_mm_castps_si128(x), 4))); 
    x = _mm_add_ps(x, _mm_castsi128_ps(_mm_slli_si128(_mm_castps_si128(x), 8)));
    return x;
}

void prefix_sum_SSE(float *a, float *s, const int n) {
__m128 offset = _mm_setzero_ps();
for (int i = 0; i < n; i+=4) {
    __m128 x = _mm_load_ps(&a[i]);
    __m128 out = scan_SSE(x);
    out = _mm_add_ps(out, offset);
    _mm_store_ps(&s[i], out);
    offset = _mm_shuffle_ps(out, out, _MM_SHUFFLE(3, 3, 3, 3)); 
}

Обратите внимание, что функция scan_SSE имеет два дополнения (_mm_add_ps) и две смены (_mm_slli_si128). Броски используются только для того, чтобы сделать компилятор счастливым и не преобразовываться в инструкции. Затем внутри основного цикла над массивом в prefix_sum_SSE используется другое добавление и один случайный перетасовка. Всего 6 операций по сравнению с 4 дополнениями с последовательной суммой.

Вот работающее решение для AVX:

inline __m256 scan_AVX(__m256 x) {
    __m256 t0, t1;
    //shift1_AVX + add
    t0 = _mm256_permute_ps(x, _MM_SHUFFLE(2, 1, 0, 3));
    t1 = _mm256_permute2f128_ps(t0, t0, 41);
    x = _mm256_add_ps(x, _mm256_blend_ps(t0, t1, 0x11));
    //shift2_AVX + add
    t0 = _mm256_permute_ps(x, _MM_SHUFFLE(1, 0, 3, 2));
    t1 = _mm256_permute2f128_ps(t0, t0, 41);
    x = _mm256_add_ps(x, _mm256_blend_ps(t0, t1, 0x33));
    //shift3_AVX + add
    x = _mm256_add_ps(x,_mm256_permute2f128_ps(x, x, 41));
    return x;
}

void prefix_sum_AVX(float *a, float *s, const int n) {
    __m256 offset = _mm256_setzero_ps();
    for (int i = 0; i < n; i += 8) {
        __m256 x = _mm256_loadu_ps(&a[i]);
        __m256 out = scan_AVX(x);
        out = _mm256_add_ps(out, offset);
        _mm256_storeu_ps(&s[i], out);
        //broadcast last element
        __m256 t0 = _mm256_permute2f128_ps(out, out, 0x11);
        offset = _mm256_permute_ps(t0, 0xff);
    }   
}

Для трех сдвигов требуется 7 встроенных функций. Трансляция требует 2 встроенных функций. Таким образом, с 4 дополнениями, что 13 интрисики. Для AVX2 для сдвигов требуется всего 5 встроенных функций, поэтому всего 11 intrinsics total. Для последовательной суммы требуется только 8 дополнений. Поэтому, вероятно, ни AVX, ни AVX2 не будут полезны для первого прохода.

Edit:

Итак, я, наконец, сравнил это, и результаты неожиданны. Код SSE и AVX примерно в два раза быстрее, чем следующий последовательный код:

void scan(float a[], float s[], int n) {
    float sum = 0;
    for (int i = 0; i<n; i++) {
        sum += a[i];
        s[i] = sum;
    }
}

Я предполагаю, что это связано с паралеллизмом уровня инструкций.

Так что я отвечаю на свой вопрос. Мне удалось использовать SIMD для pass1 в общем случае. Когда я совмещаю это с OpenMP на моей 4-жильной мостовой системе плюща, общая скорость составляет около семи для float 512k.