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

Question

Я попытался перевести следующий код в 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, чтобы посмотреть на сборку и потраченное время.)

Спасибо за помощь!

Answer 1

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

структура Оригинальный цикл:

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

Answer 2

Обратите внимание, что 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