如何在OpenMP并行区域内将BLAS gemm()调用设为单线程？

问题背景 #

在C++代码中通过#omp parallel for并行化循环，每次迭代调用BLAS的gemm()。由于使用Cray编译器并链接了OpenMP版本的cray-libsci（因需支持并行区外的多线程库shtns，必须保留-fopenmp编译标志），导致每个工作线程调用gemm()时会启动多线程，总线程数远超核心数，引发资源竞争。

解决方案 #

1. 代码内动态控制线程数 #

在并行循环的每个迭代中，调用gemm()前临时设置当前线程的OpenMP线程数为1，避免gemm()启动多线程：

#include <omp.h>
#include <cblas.h>

// ... 其他业务代码 ...

#pragma omp parallel for
for (int i = 0; i < iter_count; ++i) {
    // 保存原线程数（可选，若后续有需要多线程的操作）
    int prev_threads = omp_get_max_threads();
    
    // 强制后续并行操作仅用1线程
    omp_set_num_threads(1);
    
    // 单线程执行gemm
    cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans,
                M, N, K, 1.0, A_ptr, M, B_ptr, K, 0.0, C_ptr, M);
    
    // 恢复原线程数配置（可选）
    omp_set_num_threads(prev_threads);
}

2. 环境变量强制gemm单线程 #

利用Cray-libsci专属环境变量SCI_OMP_THREADS控制BLAS函数的线程数，无需修改代码：

• 设置SCI_OMP_THREADS=1，强制gemm()仅用1线程
• 保留OMP_NUM_THREADS为外层并行所需的线程数，确保shtns正常工作

示例运行命令：

export OMP_NUM_THREADS=16  # 外层OpenMP并行线程数
export SCI_OMP_THREADS=1    # gemm单线程限制
./your_executable

3. 关闭OpenMP嵌套并行 #

若环境中意外开启了OpenMP嵌套并行，可通过以下方式关闭，避免内层线程膨胀：

• 环境变量：export OMP_NESTED=false
• 代码内设置：在程序初始化时添加omp_set_nested(false);

注意：此方法仅当gemm()的并行依赖OpenMP嵌套机制时有效，若cray-libsci使用独立线程池，优先选择前两种方案。

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