将 mkl_set_num_threads 与 numpy 一起使用

Question

我正在尝试使用 mkl_set_num_threads 设置 numpy 计算的线程数，就像这样

import numpy
import ctypes
mkl_rt = ctypes.CDLL('libmkl_rt.so')
mkl_rt.mkl_set_num_threads(4)

但我一直收到分段错误：

Program received signal SIGSEGV, Segmentation fault.
0x00002aaab34d7561 in mkl_set_num_threads__ () from /../libmkl_intel_lp64.so

获取线程数没问题：

print mkl_rt.mkl_get_max_threads()

如何让我的代码正常工作？或者还有其他方法可以在运行时设置线程数吗？

Answer 1

Ophion 带领我走上了正确的道路。尽管有文档，但必须通过引用传递 mkl_set_num_thread 的参数。

现在我已经定义了函数，用于获取和设置线程

import numpy
import ctypes
mkl_rt = ctypes.CDLL('libmkl_rt.so')
mkl_get_max_threads = mkl_rt.mkl_get_max_threads
def mkl_set_num_threads(cores):
    mkl_rt.mkl_set_num_threads(ctypes.byref(ctypes.c_int(cores)))

mkl_set_num_threads(4)
print mkl_get_max_threads() # says 4

它们按预期工作。

编辑：根据 Rufflewind 的说法，C 函数的名称以大写形式书写，按值期望参数：

import ctypes

mkl_rt = ctypes.CDLL('libmkl_rt.so')
mkl_set_num_threads = mkl_rt.MKL_Set_Num_Threads
mkl_get_max_threads = mkl_rt.MKL_Get_Max_Threads

Answer 2

长话短说，从 Python 调用 MKL 时使用 MKL_Set_Num_Threads 及其 CamelCased 朋友。如果你不这样做，这同样适用于 C #include <mkl.h>.

MKL documentation 似乎表明 C 中正确的类型签名是：

void mkl_set_num_threads(int nt);

好，那我们来试试小程序：

void mkl_set_num_threads(int);
int main(void) {
    mkl_set_num_threads(1);
    return 0;
}

用 GCC 和 boom、Segmentation fault 再次编译它。所以看来问题并不局限于 Python.

运行它通过调试器 (GDB) 显示：

Program received signal SIGSEGV, Segmentation fault.
0x0000… in mkl_set_num_threads_ ()
   from /…/mkl/lib/intel64/libmkl_intel_lp64.so

等一下，mkl_set_num_threads_??那是 mkl_set_num_threads 的 Fortran 版本！我们是如何最终调用 Fortran 版本的？（请记住，Fortran 的调用约定要求参数作为指针而不是按值传递。）

原来文档是一个完整的门面。如果您实际检查最新版本的 MKL 的头文件，您会发现这个可爱的小定义：

void    MKL_Set_Num_Threads(int nth);
#define mkl_set_num_threads         MKL_Set_Num_Threads

……现在一切都说得通了！正确的调用函数（对于 C 代码）是 MKL_Set_Num_Threads，而不是 mkl_set_num_threads。检查符号 table 表明实际上定义了 四种不同的变体 ：

nm -D /…/mkl/lib/intel64/libmkl_rt.so | grep -i mkl_set_num_threads
00000000000e3060 T MKL_SET_NUM_THREADS
…
00000000000e30b0 T MKL_Set_Num_Threads
…
00000000000e3060 T mkl_set_num_threads
00000000000e3060 T mkl_set_num_threads_
…

尽管文档中只有 C 和 Fortran 变体，但英特尔为什么要为一个函数添加四种不同的变体？我不确定，但我怀疑这是为了与不同的 Fortran 编译器兼容。你看，Fortran 调用约定不是标准化的。不同的编译器会 mangle the names 不同的功能：

有些使用大写，
有些使用小写字母并带有尾随下划线，并且
有些使用小写字母，完全没有修饰。

甚至可能还有其他我不知道的方式。这个技巧允许 MKL 库与 大多数 Fortran 编译器 一起使用而无需任何修改，缺点是 C 函数需要 "mangled" 才能为 3 个变体腾出空间Fortran 调用约定。

Answer 3

对于寻找完整解决方案的人，您可以使用上下文管理器：

import ctypes


class MKLThreads(object):
    _mkl_rt = None

    @classmethod
    def _mkl(cls):
        if cls._mkl_rt is None:
            try:
                cls._mkl_rt = ctypes.CDLL('libmkl_rt.so')
            except OSError:
                cls._mkl_rt = ctypes.CDLL('mkl_rt.dll')
        return cls._mkl_rt

    @classmethod
    def get_max_threads(cls):
        return cls._mkl().mkl_get_max_threads()

    @classmethod
    def set_num_threads(cls, n):
        assert type(n) == int
        cls._mkl().mkl_set_num_threads(ctypes.byref(ctypes.c_int(n)))

    def __init__(self, num_threads):
        self._n = num_threads
        self._saved_n = self.get_max_threads()

    def __enter__(self):
        self.set_num_threads(self._n)
        return self

    def __exit__(self, type, value, traceback):
        self.set_num_threads(self._saved_n)

然后像这样使用它：

with MKLThreads(2):
    # do some stuff on two cores
    pass

或者只是通过调用以下函数来操作配置：

# Example
MKLThreads.set_num_threads(3)
print(MKLThreads.get_max_threads())

此 gist 中也提供了代码。

Answer 4

对于寻找跨平台和打包解决方案的人，请注意我们最近发布了 threadpoolctl, a module to limit the number of threads used in C-level threadpools called by python (OpenBLAS, OpenMP and MKL). See this 以获取更多信息。

将 mkl_set_num_threads 与 numpy 一起使用

Using mkl_set_num_threads with numpy

python

numpy

intel-mkl