如何优化基础矩阵逐元运算以达到Linpack级FLOPS性能？

问题背景 #

核心需求是优化基础矩阵逐元加/乘（Hadamard）运算的性能，这类运算在计算中高频且关键。BLAS/LAPACK不提供这类基础逐元操作，当前性能仅达Linpack水平的12%（乘加运算时为25%），目标是达到Linpack性能的60%-80%。

参考性能指标 #

• 理论峰值：i5-8259u为 4核×3.8GHz×16 FLOPS=240 GFlops
• Linpack实测：i5-8259u双精度运算可达140~160 GFlops

平台环境 #

• 设备：Macbook Pro 2018（Monterey系统）
• CPU：i5-8259u（4核8线程）
• 内存：8GB
• 编译器：gcc 11.3.0
• 编译选项：-mavx2 -mfma -fopenmp -O3

现有实现与测试结果 #

FLOPS计算方式 #

double time = stop - start;
double ops = 1.0 * Nx * Ny * iterNum; //2.0 for complex numbers
double flops = ops / time;
double gFlops = flops / 1E9;

运行结果（实数运算，复数结果相近） #

测试参数：Nx = Ny = 2048, iterNum = 10000（业务场景典型矩阵尺寸与迭代深度）

threads =    1:  1 GFlops
threads =    2:  2 GFlops
threads =    4:  3 GFlops
threads =    8:  4 GFlops
threads =   16:  9 GFlops
threads =   32: 11 GFlops
threads =   64: 15 GFlops
threads =  128: 18 GFlops
threads =  256: 19 GFlops
threads =  512: 21 GFlops
threads = 1024: 20 GFlops
threads = 2048: 40 GFlops // wrong answer

内存分配（矩阵展平为行优先向量） #

// for real numbers
x = (double *)_mm_malloc(Nx * Ny * sizeof(double), 32);
y = (double *)_mm_malloc(Nx * Ny * sizeof(double), 32);
z = (double *)_mm_malloc(Nx * Ny * sizeof(double), 32);
sum = (double *)_mm_malloc(Nx * Ny * sizeof(double), 32);
// for complex numbers
x = (double *)_mm_malloc(Nx * Ny * sizeof(double complex), 32);
y = (double *)_mm_malloc(Nx * Ny * sizeof(double complex), 32);
z = (double *)_mm_malloc(Nx * Ny * sizeof(double complex), 32);
sum = (double *)_mm_malloc(Nx * Ny * sizeof(double complex), 32);

OpenMP并行实现（加法运算） #

double start = omp_get_wtime();

#pragma omp parallel private(shift)
{
    for (int tds = omp_get_thread_num(); tds < threads; tds = tds + threads)
    {
        shift = Nx * Ny / threads * tds;
        for (int i = 0; i < iterNum; i++)
        {
            AddComplex(sum+shift, sum+shift, z+shift, Nx/threads, Ny);
        }
    }
}

double stop = omp_get_wtime();

AVX显式向量化代码 #

//real matrix addition
void AddReal(double *summation, const double *summand, const double *addend, int Nx, int Ny)
{
    int nBlock = Nx * Ny / realPackSize;
    int nRem = Nx * Ny % realPackSize;

    register __m256d packSummand, packAddend, packSum;
    const double *px = summand;
    const double *py = addend;
    double *pSum = summation;
    
    for (int i = 0; i < nBlock; i++)
    {
        packSummand = _mm256_load_pd(px);
        packAddend = _mm256_load_pd(py);
        packSum = _mm256_add_pd(packSummand, packAddend);
        _mm256_store_pd(pSum, packSum);
        px = px + realPackSize;
        py = py + realPackSize;
        pSum = pSum + realPackSize;
    }
    for (int i = 0; i < nRem; i++)
    {
        pSum[i] = px[i] + py[i];
    }

    px = NULL;
    py = NULL;
    pSum = NULL;

    return;
}
//Complex matrix addition
void AddComplex(double complex *summation, const double complex *summand, const double complex *addend, int Nx, int Ny)
{
    int nBlock = Nx * Ny / complexPackSize;
    int nRem = Nx * Ny % complexPackSize;

    register __m256d packSummand, packAddend, packSum;
    const double complex *px = summand;
    const double complex *py = addend;
    double complex *pSum = summation;

    for (int i = 0; i < nBlock; i++)
    {
        packSummand = _mm256_load_pd(px);
        packAddend = _mm256_load_pd(py);
        packSum = _mm256_add_pd(packSummand, packAddend);
        _mm256_store_pd(pSum, packSum);
        px = px + complexPackSize;
        py = py + complexPackSize;
        pSum = pSum + complexPackSize;
    }
    for (int i = 0; i < nRem; i++)
    {
        pSum[i] = px[i] + py[i];
    }
    
    px = NULL;
    py = NULL;
    pSum = NULL;

    return;
}

内容的提问来源于stack exchange，提问作者Tippsie