OpenMP嵌套线程下多插槽SpGEMM性能不达预期问题排查

我之前在双插槽Skylake系统上做过类似的NUMA优化工作，你的场景里性能远低于预期的问题，大概率是嵌套并行的线程亲和性、内存本地性、线程调度开销这几个点没处理好，下面是具体的排查方向和解决思路：

1. 内层线程未正确绑定到本地插槽核心 #

你设置了OMP_PROC_BIND=spread,close，外层2线程会被spread到两个socket，但内层的24线程是否真的限制在当前socket的核心上？GCC的OpenMP默认可能不会严格继承外层的亲和性范围，导致内层线程跨socket调度，直接拉低性能。

解决办法： #

在SpGEMM函数的并行循环里显式指定线程数和绑定策略，确保内层线程只在当前socket的24核运行：

void SpGEMM(Matrix A, Matrix B, Matrix C, int num_threads) {
    // 其他预处理逻辑...
    #pragma omp parallel for num_threads(num_threads) proc_bind(close) schedule(static)
    for (int k = 0; k < A.ncols; ++k) {
        // 基于外积的SpGEMM计算逻辑
    }
}

proc_bind(close)会让内层线程紧密绑定到当前外层线程所在的NUMA节点核心，避免跨socket调度。

2. 共享矩阵B的跨NUMA访问瓶颈 #

你的方案里B是全局共享的，单插槽时B在本地内存所以快，但双插槽分区后，其中一个socket的线程访问B会走远端NUMA链路，延迟是本地的3-5倍，这是性能暴跌的核心原因之一。

解决办法： #

给每个socket复制一份B的本地副本，让每个分区的SpGEMM只访问本地内存里的B：

// 引入NUMA库头文件
#include <numa.h>

// 在每个section里分配本地B副本
#pragma omp section {
    // 获取当前线程所在的NUMA节点
    int local_node = numa_node_of_cpu(sched_getcpu());
    // 分配本地内存存储B副本
    Matrix B_local;
    B_local.nrows = B.nrows;
    B_local.ncols = B.ncols;
    B_local.data = numa_alloc_onnode(sizeof(double)*B.nrows*B.ncols, local_node);
    // 复制数据到本地内存
    memcpy(B_local.data, B.data, sizeof(double)*B.nrows*B.ncols);
    
    // 调用SpGEMM时使用本地B副本
    SpGEMM(A_upper, B_local, C_upper, 24);
    
    // 释放本地内存
    numa_free(B_local.data, sizeof(double)*B.nrows*B.ncols);
}

这样两个socket的线程都只访问自己本地的B副本，彻底消除跨NUMA访问的开销。

3. 循环内嵌套并行的线程创建开销过大 #

你把omp parallel sections放在迭代循环内部，每次迭代都要创建/销毁外层2线程和内层48线程，频繁的线程调度开销会严重吞噬计算时间，这也是性能只有400 MFLOPS的关键因素。

解决办法： #

把并行区域移到迭代循环外面，让线程只初始化一次，循环在每个section内部执行：

omp_set_nested(1);
omp_set_dynamic(0);
omp_set_num_threads(2);

double total_ave = 0;

#pragma omp parallel sections reduction(+:total_ave)
{
    #pragma omp section {
        int local_node = numa_node_of_cpu(sched_getcpu());
        Matrix B_local = {B.nrows, B.ncols, numa_alloc_onnode(sizeof(double)*B.nrows*B.ncols, local_node)};
        memcpy(B_local.data, B.data, sizeof(double)*B.nrows*B.ncols);
        
        double ave_msec = 0;
        for (int i = 0; i < ITERS; ++i) {
            double start = omp_get_wtime();
            SpGEMM(A_upper, B_local, C_upper, 24);
            double end = omp_get_wtime();
            ave_msec += (end - start) * 1000 / ITERS;
        }
        total_ave += ave_msec;
        numa_free(B_local.data, sizeof(double)*B.nrows*B.ncols);
    }
    #pragma omp section {
        int local_node = numa_node_of_cpu(sched_getcpu());
        Matrix B_local = {B.nrows, B.ncols, numa_alloc_onnode(sizeof(double)*B.nrows*B.ncols, local_node)};
        memcpy(B_local.data, B.data, sizeof(double)*B.nrows*B.ncols);
        
        double ave_msec = 0;
        for (int i = 0; i < ITERS; ++i) {
            double start = omp_get_wtime();
            SpGEMM(A_lower, B_local, C_lower, 24);
            double end = omp_get_wtime();
            ave_msec += (end - start) * 1000 / ITERS;
        }
        total_ave += ave_msec;
        numa_free(B_local.data, sizeof(double)*B.nrows*B.ncols);
    }
}

printf("Total average MFLOPS: %.2f\n", (TOTAL_FLOPS * 2) / (total_ave / 1000));

这样线程只创建一次，避免了每次循环的调度开销。

4. 确认数据分区的内存本地性 #

确保A_upper、C_upper是分配在socket 0的本地内存，A_lower、C_lower分配在socket 1的本地内存。如果这些数据是在全局内存分配的，可能被NUMA调度器放到远端节点，导致本地线程访问远端数据。

解决办法： #

用numa_alloc_onnode显式指定数据分配的NUMA节点，比如在主线程里提前分配：

// 分配A_upper到socket 0
A_upper.data = numa_alloc_onnode(sizeof(double)*A_upper.nrows*A_upper.ncols, 0);
// 分配A_lower到socket 1
A_lower.data = numa_alloc_onnode(sizeof(double)*A_lower.nrows*A_lower.ncols, 1);
// 同理分配C_upper和C_lower到对应节点

5. 调整环境变量的亲和性设置 #

可以修改OMP_PLACES为更明确的核心范围，确保外层线程精准绑定到两个socket：

export OMP_PLACES="{0}:24,{1}:24"  # 明确指定socket 0的24核和socket 1的24核
export OMP_PROC_BIND=close,close   # 外层close到socket，内层close到当前socket的核心

按照上面的步骤调整后，应该能把性能提升到接近1400 MFLOPS的预期值。

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