如何启用numpy.linalg.svd及矩阵乘法的多线程计算

问题现象 #

• 脚本大量调用NumPy及numpy.linalg模块接口，核心计算逻辑依赖np.linalg.svd，按照特性这类线性代数运算应自动启用多线程加速，但通过htop监控发现脚本运行时始终仅占用1个CPU线程，多线程加速未生效。
• 当前环境numpy.show_config()输出如下，可见NumPy底层绑定的是64位OpenBLAS后端：

openblas64__info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
openblas64__lapack_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2
    not found = AVX512F,AVX512CD,AVX512_KNL,AVX512_KNM,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL

• 问题最小复现代码：

import numpy as np 
import tensorly as ty 

tensor = np.random.rand(32,32,32,32)
unfolding = ty.unfold(tensor,0)
unfolding = unfolding @ unfolding.transpose()
U,S,_ = np.linalg.svd(unfolding)

根因说明 #

NumPy自身不实现多线程调度逻辑，numpy.linalg下的SVD、矩阵乘法等运算的多线程能力完全由底层链接的BLAS/LAPACK后端提供。当前环境绑定的是自行编译的ILP64版本OpenBLAS，这类版本出现单线程锁死的常见原因包括：

• 编译OpenBLAS时未开启多线程编译选项，默认编译为单线程版本
• 系统或用户环境变量全局将OPENBLAS_NUM_THREADS、OMP_NUM_THREADS设置为1
• 编译OpenBLAS时线程模型配置错误，与系统pthread库适配异常导致多线程调度失效

解决方案 #

方案1：替换为MKL后端NumPy（已验证可解决问题） #

Intel MKL后端对线性代数运算的多线程调度适配成熟，SVD、矩阵乘等常用接口的优化完善，不会出现OpenBLAS编译配置导致的线程锁死问题，是生产环境最稳定的选择。
操作步骤：

1. 卸载当前绑定OpenBLAS的NumPy版本：执行命令pip uninstall -y numpy
2. 安装MKL绑定的NumPy版本，通过conda安装是最简便的方式，conda源分发的NumPy默认预编译链接MKL后端，无需额外配置：执行命令conda install numpy
3. 安装完成后执行numpy.show_config()，输出中出现mkl相关配置项即说明替换成功，运行脚本时会自动根据CPU核心数调度多线程计算。

方案2：保留OpenBLAS后端排查修复 #

如果需要保留现有OpenBLAS后端，可按以下顺序排查：

• 在脚本最开头、导入NumPy之前手动指定OpenBLAS线程数，代码如下：

import os
# 数值替换为当前机器的CPU逻辑核心数
os.environ["OPENBLAS_NUM_THREADS"] = "8"
os.environ["OMP_NUM_THREADS"] = "8"

• 检查系统全局profile、用户bashrc/zshrc配置文件中是否存在export OPENBLAS_NUM_THREADS=1这类配置，若存在则删除对应配置后重新登录终端生效。
• 若上述操作无效，说明当前安装的OpenBLAS本身是单线程编译版本，需要重新编译OpenBLAS，编译时添加USE_THREAD=1参数指定启用pthread多线程支持，编译完成后重新编译NumPy链接到新的多线程版OpenBLAS即可。

验证结果 #

按照方案1替换为MKL后端的NumPy后，多线程未生效问题已解决，线性代数运算可正常占用多个CPU核心加速。

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