如何验证Libtorch CPU版是否使用MKL及多线程加速矩阵乘法

1. Libtorch底层是否使用MKL？ #

官网提供的CPU-only版本Libtorch默认链接MKL库，可通过你给出的nm命令输出进一步确认：

• MKL的BLAS核心函数（如复数矩阵乘法的zgemm、向量乘法的zgemv）均以带下划线的后缀形式存在（例如zgemm_）。
• 如果nm -a --demangle libtorch_cpu.so | grep 'zgemm'的输出包含U zgemm_（表示动态链接MKL库的未定义符号），或是直接显示zgemm_的定义，即可证明Libtorch依赖MKL执行底层矩阵运算。

2. Libtorch矩阵乘法是否使用多线程？之前的线程数检查能否保证？ #

多线程有效性判断 #

Libtorch的torch::matmul处理std::complex<double>矩阵时，底层会调用MKL的zgemm，其线程调度受以下设置控制：

• 代码中mkl_set_num_threads(64)直接指定MKL运算线程数，mkl_get_max_threads()返回64说明该设置已生效。
• omp_set_num_threads(64)会影响MKL（若MKL以OMP作为线程后端），omp_get_max_threads()返回64也验证了OMP线程设置有效。
• torch::get_num_threads()返回的是PyTorch内部ATen线程池的大小，CPU运算时PyTorch会通过该线程池调度任务，配合MKL的线程设置，理论上会启用多线程运算。

速度慢的潜在排查点 #

尽管线程数检查结果正常，仍有几个因素可能导致运算速度未达预期：

• 矩阵尺寸过小：MKL对小矩阵（如几十阶）会跳过多线程优化，因为线程调度开销超过运算收益，建议测试≥500×500的大尺寸矩阵验证速度。
• 环境变量覆盖：若系统中设置了MKL_NUM_THREADS、OMP_NUM_THREADS等环境变量，其优先级可能高于代码中的mkl_set_num_threads/omp_set_num_threads调用，运行时可通过echo $MKL_NUM_THREADS检查。
• 嵌套线程竞争：PyTorch的ATen线程池和MKL线程可能出现嵌套（如PyTorch开64个线程，每个线程内MKL又开64个线程），导致CPU核心过度竞争。可临时设置torch::set_num_threads(1)，仅保留MKL的64线程设置，观察速度变化。

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