Как оптимизировать мой код AVX

Я попытался перевести следующий код в встроенные функции AVX для повышения производительности:

for (int alpha = 0; alpha < 4; alpha++) {
for (int k = 0; k < 3; k++) {
for (int beta = 0; beta < 4; beta++) {
for (int l = 0; l < 4 ; l++) {
d2_phi[(alpha*3+k)*16 + beta*4+l] =
-   (d2_phi[(alpha*3+k)*16 + beta*dim+l]

+   b[k] * (  lam_12[ beta][alpha] *   a[l]
+ lam_22[alpha][ beta] *   b[l]
+ lam_23[alpha][ beta] * rjk[l]  )

+ rjk[k] * (  lam_13[ beta][alpha] *   a[l]
+ lam_23[ beta][alpha] *   b[l]
+ lam_33[alpha][ beta] * rjk[l]  )
) / sqrt_gamma;
}
}
}
}

и попробовал это следующим образом:

// load sqrt_gamma, because it is constant
__m256d ymm7 = _mm256_broadcast_sd(&sqrt_gamma);

for (int alpha=0; alpha < 4; alpha++) {
for (int k=0; k < 3; k++) {
// Load values that are only dependent on k
__m256d ymm9 = _mm256_broadcast_sd(b+k);   // all   b[k]
__m256d ymm8 = _mm256_broadcast_sd(rjk+k); // all rjk[k]

for (int beta=0; beta < 4; beta++) {
// Load the lambdas, because they will stay the same for nine iterations
__m256d ymm15 = _mm256_broadcast_sd(lam_12_p + 4*beta + alpha);   // all lam_12[ beta][alpha]
__m256d ymm14 = _mm256_broadcast_sd(lam_22_p + 4*alpha + beta);   // all lam_22[alpha][ beta]
__m256d ymm13 = _mm256_broadcast_sd(lam_23_p + 4*alpha + beta);   // all lam_23[alpha][ beta]
__m256d ymm12 = _mm256_broadcast_sd(lam_13_p + 4*beta + alpha);   // all lam_13[ beta][alpha]
__m256d ymm11 = _mm256_broadcast_sd(lam_23_p + 4*beta + alpha);   // all lam_23[ beta][alpha]
__m256d ymm10 = _mm256_broadcast_sd(lam_33_p + 4*alpha + beta);   //     lam_33[alpha][ beta]

// Load the values that depend on the innermost loop, which is removed do to AVX
__m256d ymm6 =_mm256_load_pd(a);   //   a[i] until   a[l+3]
__m256d ymm5 =_mm256_load_pd(b);   //   b[i] until   b[l+3]
__m256d ymm4 =_mm256_load_pd(rjk); // rjk[i] until rjk[l+3]
//__m256d ymm3 =_mm256_load_pd(d2_phi_p + (alpha*3+k)*16  + beta*dim); // d2_phi[(alpha*3+k)*12 + beta*dim] until d2_phi[(alpha*3+k)*12 + beta*dim +3]
__m256d ymm3 =_mm256_load_pd(d2_phi_p + 4*s);
// Block that is later on multiplied with b[k]
__m256d ymm2 = _mm256_mul_pd(ymm15, ymm6); // lam_12[ beta][alpha] * a[l]
__m256d ymm1 = _mm256_mul_pd(ymm14, ymm5); // lam_22[alpha][ beta] * b[l];

__m256d ymm0 = _mm256_add_pd(ymm2, ymm1);  // lam_12[ beta][alpha] * a[l] + lam_22[alpha][ beta]*b[l];

ymm2 = _mm256_mul_pd(ymm13, ymm4);         // lam_23[alpha][ beta] * rjk[l]
ymm0 = _mm256_add_pd(ymm2, ymm0);          // lam_12[ beta][alpha] * a[l] + lam_22[alpha][ beta]*b[l] + lam_23[alpha][ beta] * b[i];

ymm0 = _mm256_mul_pd(ymm9, ymm0);          // b[k] * (first sum of three)// Block that is later on multiplied with rjk[k]
ymm2 = _mm256_mul_pd(ymm12, ymm6); // lam_13[ beta][alpha] *  a[l]
ymm1 = _mm256_mul_pd(ymm11, ymm5); // lam_23[ beta][alpha] *  b[l]

ymm2 = _mm256_add_pd(ymm2, ymm1);  // lam_13[ beta][alpha] *  a[l] + lam_22[alpha][ beta]*b[l];

ymm1 = _mm256_mul_pd(ymm10, ymm4); // lam_33[alpha][ beta] * rjk[l]
ymm2 = _mm256_add_pd(ymm2, ymm1);  // lam_13[ beta][alpha] *  a[l] + lam_22[alpha][ beta]*b[l] + lam_33[alpha][ beta] *rjk[l]

ymm2 = _mm256_mul_pd(ymm2, ymm8);  // rjk[k] * (second sum of three)
ymm0 = _mm256_add_pd(ymm0, ymm2);  // add to temporal result in ymm0
ymm0 = _mm256_add_pd(ymm3, ymm0);  // Old value of d2 Phi;

ymm0 = _mm256_div_pd(ymm0, ymm7);   // all divided by sqrt_gamma

_mm256_store_pd(d2_phi_p + (alpha*3+k)*16  + beta*dim, ymm0);
}
}
}

Но производительность плохая. Это даже медленнее, чем автоматически векторизованный код, сгенерированный компилятором Intel. Я пробовал следующие вещи:

  • Все массивы данных выровнены по 64 байта __declspec(align(64))
  • Магазин в конце был заменен потоковым магазином _mm256_stream_pd

Когда я смотрю на созданный код сборки, я вижу, что авто-код выбирает все параметры на каждой итерации (и не так, как я, только в циклах, к которым они принадлежат). Он также содержит больше арифметических операций. Как последний пункт, магазину в конце нужна только половина моего времени (я повторяю фрагмент кода 1000000 раз), и я не вижу причины для этого. (Я использовал Intel VTune Amplifier, чтобы посмотреть на сборку и потраченное время.)

Заранее всем спасибо за помощь!

3

Решение

Я помещу это здесь как второй ответ, так как это другая и намного более обширная оптимизация. Ключ к изменить порядок циклов уменьшить количество избыточных операций, вытащив множество нагрузок и арифметических операций из самого внутреннего цикла.

Оригинальная структура петли:

for (int alpha=0; alpha < 4; alpha++) {
for (int k=0; k < 3; k++) {
for (int beta=0; beta < 4; beta++) {
for (int l=0; l < 4 ; l++) {

Новая структура цикла:

for (int alpha=0; alpha < 4; alpha++) {
for (int beta=0; beta < 4; beta++) {
for (int k=0; k < 3; k++) {
for (int l=0; l < 4 ; l++) {

Полная проверенная и оптимизированная реализация вашей функции:

static void foo(
double lam_11[4][4],
double lam_12[4][4],
double lam_13[4][4],
double lam_22[4][4],
double lam_23[4][4],
double lam_33[4][4],
const double rjk[4],
const double a[4],
const double b[4],
const double sqrt_gamma,
const double SPab,
const double d1_phi[16],
double d2_phi[192])
{
const double re_sqrt_gamma = 1.0 / sqrt_gamma;

memset(d2_phi, 0.0, 192*sizeof(double));

const __m256d ymm6 = _mm256_load_pd(a); // load the whole 4-vector 'a' into register

{
// load SPab, because it is constant
const __m256d ymm0 = _mm256_broadcast_sd(&SPab);
const __m256d ymm7 = _mm256_load_pd(b); // load the whole 4-vector 'b' into register
const __m256d ymm8 = _mm256_load_pd(rjk); // load the whole 4-vector 'rjk' into register

for (int alpha=0; alpha < 4; alpha++)
{
for (int beta=0; beta < 4; beta++)
{
// Load the three lambdas to all
const __m256d ymm3 = _mm256_broadcast_sd(&lam_11[alpha][beta]);
const __m256d ymm4 = _mm256_broadcast_sd(&lam_12[alpha][beta]);
const __m256d ymm5 = _mm256_broadcast_sd(&lam_13[alpha][beta]);

const __m256d ymm9 = _mm256_load_pd(d1_phi + beta*4);

// Do the three Multiplications
const __m256d ymm13 = _mm256_mul_pd(ymm4,ymm7); // lam_12[alpha][ beta] *  b[l] = PROD2
const __m256d ymm14 = _mm256_mul_pd(ymm5,ymm8); // lam_13[alpha][ beta] * rjk[l] = PROD3
const __m256d ymm15 = _mm256_mul_pd(ymm3,ymm6); // lam_11[alpha][ beta] *  a[l] = PROD1
__m256d ymm12 = _mm256_add_pd(ymm15, ymm13); // PROD1 + PROD2 = PROD12
ymm12 = _mm256_add_pd(ymm12, ymm14); // PROD12 + PROD3 = PROD123

double* addr = d2_phi + alpha*3*16  + beta*dim;

for (int k=0; k < 3; k++)
{
const __m256d ymm1 = _mm256_broadcast_sd(&d1_phi[alpha*dim + k]); // load d1_phi[alpha*dim+k] to all
const __m256d ymm2 = _mm256_broadcast_sd(&a[k]); // load a[k] to all
const __m256d ymm10 = _mm256_mul_pd(ymm0, ymm1); // SPab * d1_phi[alpha*dim+k] = PRE
const __m256d ymm11 = _mm256_mul_pd(ymm10, ymm9); // PRE * d1_phi[beta*dim+l] = SUM1

__m256d ymm12t = _mm256_mul_pd(ymm12, ymm2); // a[k] * PROD123 = SUM2
ymm12t = _mm256_add_pd(ymm11, ymm12t); // SUM1 + SUM2

_mm256_store_pd(addr, ymm12t);

addr+=16;
}
}
}
}

{
const __m256d ymm4 =_mm256_load_pd(rjk); // rjk[i] until rjk[l+3]
const __m256d ymm5 =_mm256_load_pd(b); // b[l] until b[l+3]

// load sqrt_gamma, because it is constant
const __m256d ymm7 = _mm256_broadcast_sd(&re_sqrt_gamma);

for (int alpha=0; alpha < 4; alpha++)
{
for (int beta=0; beta < 4; beta++)
{
// Load the lambdas, because they will stay the same for nine iterations
const __m256d ymm15 = _mm256_broadcast_sd(&lam_12[beta][alpha]);   // all lam_12[ beta][alpha]
const __m256d ymm14 = _mm256_broadcast_sd(&lam_22[alpha][beta]);   // all lam_22[alpha][ beta]
const __m256d ymm13 = _mm256_broadcast_sd(&lam_23[alpha][beta]);   // all lam_23[alpha][ beta]
const __m256d ymm12 = _mm256_broadcast_sd(&lam_13[beta][alpha]);   // all lam_13[ beta][alpha]
const __m256d ymm11 = _mm256_broadcast_sd(&lam_23[beta][alpha]); // all lam_23[ beta][alpha]
const __m256d ymm10 = _mm256_broadcast_sd(&lam_33[alpha][beta]); // lam_33[alpha][ beta]

__m256d ymm0, ymm1, ymm2;

// Block that is later on multiplied with b[k]
ymm2 = _mm256_mul_pd(ymm15 , ymm6); // lam_12[ beta][alpha] *  a[l]
ymm1 = _mm256_mul_pd(ymm14 , ymm5); // lam_22[alpha][ beta] * b[l];
ymm0 = _mm256_add_pd(ymm2, ymm1);   // lam_12[ beta][alpha]* a[l] + lam_22[alpha][ beta]*b[l];
ymm2 = _mm256_mul_pd(ymm13 , ymm4); // lam_23[alpha][ beta] * rjk[l]
ymm0 = _mm256_add_pd(ymm2, ymm0);   // lam_12[ beta][alpha]* a[l] + lam_22[alpha][ beta]*b[l] + lam_23[alpha][ beta] * b[i];

// Block that is later on multiplied with rjk[k]
ymm2 = _mm256_mul_pd(ymm12 , ymm6); // lam_13[ beta][alpha] *  a[l]
ymm1 = _mm256_mul_pd(ymm11 , ymm5); // lam_23[ beta][alpha] *  b[l]
ymm2 = _mm256_add_pd(ymm2, ymm1);   // lam_13[ beta][alpha] *  a[l] + lam_22[alpha][ beta]*b[l];
ymm1 = _mm256_mul_pd(ymm10 , ymm4); // lam_33[alpha][ beta] * rjk[l]
ymm2 = _mm256_add_pd(ymm2 , ymm1);  // lam_13[ beta][alpha] *  a[l] + lam_22[alpha][ beta]*b[l] + lam_33[alpha][ beta] *rjk[l]

double* addr = d2_phi + alpha*3*16  + beta*dim;

for (int k=0; k < 3; k++)
{
// Load values that are only dependent on k
const __m256d ymm9 = _mm256_broadcast_sd(b+k); // all b[k]
const __m256d ymm8 = _mm256_broadcast_sd(rjk+k); // all rjk[k]

// Load the values that depend on the innermost loop, which is removed do to AVX

const __m256d ymm3 =_mm256_load_pd(addr);

__m256d ymm0t, ymm1t, ymm2t;

// Block that is later on multiplied with b[k]

ymm0t = _mm256_mul_pd(ymm9 , ymm0); // b[k] * (first sum of three)

// Block that is later on multiplied with rjk[k]

ymm1t = _mm256_mul_pd(ymm2 , ymm8); // rjk[k] * (second sum of three)
ymm2t = _mm256_add_pd(ymm0t, ymm1t); // add to temporal result in ymm0
ymm1t = _mm256_add_pd(ymm3, ymm2t);  // Old value of d2 Phi;

ymm2t = _mm256_mul_pd(ymm1t, ymm7); // all divided by sqrt_gamma
ymm1t = _mm256_xor_pd(ymm2t, SIGNMASK);

_mm256_store_pd(addr, ymm1t);

addr += 16;
}
}
}
}
}

Исходный код AVX работал с вашим тестовым проводом около 500 мс, новая версия — около 200 мс, так что это увеличило пропускную способность в 2,5 раза.

Обновленная версия вашего тестового набора с оригинальным кодом и оптимизированным кодом здесь: http://pastebin.com/yMPbYPjb

3

Другие решения

Обратите внимание, что VDIVPD стоит дорого — имеет типичную задержку / пропускную способность порядка 20 — 40 циклов (точные цифры зависят от процессора). Итак, одно немедленное изменение, которое я бы предложил, — преобразовать деление на константу в умножение, так как VMULPD имеет задержку всего несколько циклов и пропускную способность одного цикла:

// load 1 / sqrt_gamma, because it is constant
const double re_sqrt_gamma = 1.0 / sqrt_gamma;
__m256d ymm7 = _mm256_broadcast_sd(&re_sqrt_gamma);

ymm0 = _mm256_mul_pd(ymm0, ymm7);   // all divided by sqrt_gamma
3