为什么 _mm512_store_pd 在这个矩阵乘法代码中超级慢？

Question

我正在使用 avx512 和矩阵乘法，但我一定做错了，因为当我尝试使用 _mm512_store_pd 存储我的结果时，我的表现很糟糕。

以下是相关的sn-ps代码，首先是我正在使用的数据结构以及我如何初始化它：

typedef struct {
        double* values;
        int nb_l;
        int nb_c;
} matrix;

matrix* alloc_matrix(int nb_l, int nb_c){
        matrix* tmp_matrix = (matrix*)malloc(sizeof(matrix));
        tmp_matrix->values = (double*)aligned_alloc(64, sizeof(double) * nb_l * nb_c);
        tmp_matrix->nb_l = nb_l;
        tmp_matrix->nb_c = nb_c;
        return tmp_matrix;
}

这就是我如何尝试将在我的代码中其他地方初始化的两个矩阵相乘：

matrix* mult_matrix(matrix* A, matrix* B){
        /* avx512 */
        matrix* res_matrix = zero_matrix(A->nb_l, B->nb_c);
        double* res_ptr; // start index of the current line in res_matrix
        double* B_ptr; // start index of the current line in B

        __m512d A_broadcast, B_l_8, res_ptr_8;
        for (unsigned int idx_A = 0; idx_A < A->nb_l * A-> nb_c; idx_A++){
                // broadcast current value of A  eight times
                A_broadcast = _mm512_set1_pd(A->values[idx_A]);
                res_ptr = res_matrix->values + (idx_A / A->nb_c) * B->nb_c;
                B_ptr = B->values + (idx_A % A->nb_c) * B->nb_c;
                for (unsigned int offset_B = 0; offset_B < B->nb_c; offset_B+=8){
                        B_l_8 = _mm512_load_pd(&B_ptr[offset_B]);
                        res_ptr_8 = _mm512_load_pd(&res_ptr[offset_B]);
                        _mm512_store_pd(
                                &res_ptr[offset_B] , 
                                _mm512_fmadd_pd(A_broadcast, B_l_8, res_ptr_8)
                                );
                }
        }
        return res_matrix;

结果还可以，但 _mm512_store_pd 占用了约 90% 的执行时间，实际上这个 avx512 代码几乎不比它的非 avx 版本快。

我已经尝试了所有我能想到的方法，但我不知道为什么我的代码表现如此令人失望。你有什么想法吗？

谢谢。

编辑 1

这里是非avx代码

        matrix* res_matrix = zero_matrix(A->nb_l, B->nb_c);
        double* res_ptr; // start index of the current line in res_matrix
        double* B_ptr; // start index of the current line in B

        for (unsigned int idx_A = 0; idx_A < A->nb_l * A-> nb_c; idx_A++){
                res_ptr = res_matrix->values + (idx_A / A->nb_c) * B->nb_c; 
                B_ptr = B->values + (idx_A % A->nb_c) * B->nb_c; 
                for (unsigned int offset_B = 0; offset_B < B->nb_c; offset_B++){                    
                        res_ptr[offset_B] += A->values[idx_A] * B_ptr[offset_B];
                }
        }
        return res_matrix;

所有矩阵都是512x512的随机矩阵，每次乘法重复50次，平均运行时间。

最后，为了测试我的代码的 avx 和 non_avx 版本，下面的 sn-p 应该没问题。我使用以下选项使用 gcc 8.3.0 编译它： gcc -Ofast -mavx -mavx512f -m64 -mfpmath=sse -mfma -flto -funroll-loops matrix_minimal.c

#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
#include <time.h>
#include <string.h>
#include <immintrin.h>

typedef struct {
        double* values;
        int nb_l;
        int nb_c;
} matrix;

matrix* alloc_matrix(int nb_l, int nb_c){
        matrix* tmp_matrix = (matrix*)malloc(sizeof(matrix));
        tmp_matrix->values = (double*)aligned_alloc(64, sizeof(double) * nb_l * nb_c);
        tmp_matrix->nb_l = nb_l;
        tmp_matrix->nb_c = nb_c;
        return tmp_matrix;
}

void free_matrix(matrix** to_free){
        free((*to_free)->values);
        free(*to_free);
}

matrix* zero_matrix(int nb_l, int nb_c){
        matrix* z_matrix;
        z_matrix = alloc_matrix(nb_l, nb_c);
        for (int idx=0; idx < nb_l * nb_c; idx++){
                z_matrix->values[idx] = 0.0;
        }
        return z_matrix;
}
matrix* rand_matrix(int nb_l, int nb_c, double max_abs_val){
        static struct timeval seed; //static variables are zeroed at initialization
        matrix* rnd_matrix;
        rnd_matrix = alloc_matrix(nb_l, nb_c);

        if (seed.tv_sec == 0){ //ts_sec will never be zero after gettimeofday, whereas tv_usec could
                gettimeofday(&seed, NULL);
                srand((unsigned) seed.tv_usec);
        }
        for (int idx=0; idx < nb_l * nb_c; idx++){
                rnd_matrix->values[idx] = max_abs_val * ((double)rand() / RAND_MAX * 2.0 - 1.0);
        }

        return rnd_matrix;
}

matrix* mult_matrix_avx(matrix* A, matrix* B){
        /* pas trop mal en avx512 */
        matrix* res_matrix = zero_matrix(A->nb_l, B->nb_c);
        double* res_ptr; // start index of the current line in res_matrix
        double* B_ptr; // start index of the current line in B

        __m512d A_broadcast, B_l_8, res_ptr_8;
        for (unsigned int idx_A = 0; idx_A < A->nb_l * A-> nb_c; idx_A++){
                A_broadcast = _mm512_set1_pd(A->values[idx_A]); // broadcast current value of A eight times
                res_ptr = res_matrix->values + (idx_A / A->nb_c) * B->nb_c;
                B_ptr = B->values + (idx_A % A->nb_c) * B->nb_c;
                for (unsigned int offset_B = 0; offset_B < B->nb_c; offset_B+=8){
                        B_l_8 = _mm512_load_pd(&B_ptr[offset_B]);
                        res_ptr_8 = _mm512_load_pd(&res_ptr[offset_B]);
                        _mm512_store_pd(&res_ptr[offset_B] , _mm512_fmadd_pd(A_broadcast, B_l_8, res_ptr_8));
                }
        }
        return res_matrix;
}

matrix* mult_matrix(matrix* A, matrix* B){
        /* non avx512 */
        matrix* res_matrix = zero_matrix(A->nb_l, B->nb_c);
        double* res_ptr; // start index of the current line in res_matrix
        double* B_ptr; // start index of the current line in B

        for (unsigned int idx_A = 0; idx_A < A->nb_l * A-> nb_c; idx_A++){
                res_ptr = res_matrix->values + (idx_A / A->nb_c) * B->nb_c;
                B_ptr = B->values + (idx_A % A->nb_c) * B->nb_c;
                for (unsigned int offset_B = 0; offset_B < B->nb_c; offset_B++){
                        res_ptr[offset_B] += A->values[idx_A] * B_ptr[offset_B];
                }
        }
        return res_matrix;
}
int main(int argc, char *argv[]){
        struct timeval before;
        struct timeval after;

        matrix* A = rand_matrix(512, 512, 5);
        matrix* B = rand_matrix(512, 512, 5);
        matrix *C;
        gettimeofday(&before, NULL);
        for (int j=0; j<50;j++){
                C = mult_matrix_avx(A, B);
                free_matrix(&C); // we will measure the same overhead here and in the non avx version
        }
        gettimeofday(&after, NULL);
        double delta = ((after.tv_sec - before.tv_sec) * 1000000 +
                (after.tv_usec - before.tv_usec))/50;
        printf("avx %lf ms\n", delta);
        gettimeofday(&before, NULL);
        for (int j=0; j<50;j++){
                C = mult_matrix(A, B);
                free_matrix(&C); 
        }
        gettimeofday(&after, NULL);
        delta = ((after.tv_sec - before.tv_sec) * 1000000 +
                (after.tv_usec - before.tv_usec))/50;
        printf("non avx %lf ms\n", delta);

        free_matrix(&A);
        free_matrix(&B);
        return 0;
}

Answer 1

@chtz 指出了显而易见的答案，我过于频繁地更新结果，而不是依赖 mm512 寄存器。

如果我比较这个非 avx 版本：

        matrix* res_matrix = alloc_matrix(A->nb_l, B->nb_c);
        double tmp;
        for (unsigned int res_l=0; res_l < res_matrix->nb_l; res_l++){
                for (unsigned int res_c=0; res_c < res_matrix->nb_c; res_c++){
                        tmp = 0.0;
                        for (int offset = 0; offset < A->nb_c; offset++){
                                tmp += A->values[res_l * A->nb_c + offset] *
                                        B->values[offset * A->nb_c + res_c];
                        }
                        res_matrix->values[res_l * res_matrix->nb_c + res_c] = tmp;
                }

        }
        return res_matrix;

与其精确的 AVX 对应物：

        matrix* res_matrix = alloc_matrix(A->nb_l, B->nb_c);
        __m512d res_ptr_8;

        for (unsigned int res_l=0; res_l < res_matrix->nb_l; res_l++){
                for (unsigned int res_c=0; res_c < res_matrix->nb_c; res_c+=8){
                        // compute values from res_matrix[res_l, res_c] to [res_l, res_c+7]
                        res_ptr_8 = _mm512_set1_pd(0.0);
                        for (unsigned int offset_A_c = 0; offset_A_c < A->nb_c; offset_A_c++){
                                // on the res_l th line of A pick values one at a time
                                // at coordinates A[res_l, offset_A_c].
                                // Broadcast this value eight times into a mm512 vector
                                // and perform a dot product with the 8 values found in
                                // B from coordinates [offset_A_c, res_c] to [offset_A_c, res_c + 7]
                                res_ptr_8 = _mm512_fmadd_pd(
                                        _mm512_set1_pd(A->values[res_l * A->nb_c + offset_A_c]),
                                        _mm512_load_pd(&B->values[offset_A_c * B->nb_c + res_c]),
                                        res_ptr_8);
                        }
                        _mm512_store_pd(&res_matrix->values[res_l*res_matrix->nb_c + res_c] , res_ptr_8);
                }
        }

AVX 代码的运行时间比我预期的非 AVX 代码快约 7 倍。

但是，如果我将这个新的 AVX 代码与粘贴在问题中的非 AVX 代码进行比较，速度增益仅为 3 左右，这对我来说仍然足够好。

编辑

这些速度增益是在小块上测量的，以适应缓存，而不是在原始帖子中的 512*512 矩阵上。

编辑 2 感谢 Peter Cordes（见下文，编辑：看起来我完全误解了他所说的话）这是一个更新的 AVX 代码，在 48x48 块上比上面的代码快两倍。令人惊讶的是，它甚至比同一矩阵上的 numpy/openblas 还要快。

        #define NB_L_STRIDE 8
        __m512d res_ptr_8[NB_L_STRIDE], B_ptr_8;
        for (unsigned int res_l=0; res_l < res_matrix->nb_l; res_l+=NB_L_STRIDE){
                for (unsigned int res_c=0; res_c < res_matrix->nb_c; res_c+=8){
                        for(unsigned int i=0; i<NB_L_STRIDE; i++)
                                res_ptr_8[i] = _mm512_setzero_pd();
                        for (unsigned int offset_A_c = 0; offset_A_c < A->nb_c; offset_A_c++){
                        // compute values from res_matrix[res_l, res_c] to [res_l, res_c+7]
                                // on the res_l th line of A pick values one at a time
                                // at coordinates A[res_l, offset_A_c].
                                // Broadcast this value eight times into a mm512 vector
                                // and perform a dot product with the 8 values found in
                                // B from coordinates [offset_A_c, res_c] to [offset_A_c, res_c + 7]
                                B_ptr_8 = _mm512_load_pd(&B->values[offset_A_c * B->nb_c + res_c]);
                                for(unsigned int i=0; i<NB_L_STRIDE; i++)
                                        res_ptr_8[i] = _mm512_fmadd_pd(
                                                _mm512_set1_pd(A->values[(res_l +i) * A->nb_c + offset_A_c]),
                                                B_ptr_8,
                                                res_ptr_8[i]);

                        }
                        for(unsigned int i=0; i<NB_L_STRIDE; i++)
                                _mm512_store_pd(
                                        &res_matrix->values[(res_l + i)*res_matrix->nb_c + res_c] ,
                                        res_ptr_8[i]);
                }
        }