使用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)

但我不断收到分段错误（segmentation fault）：

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

获取线程数量并不是问题：

print mkl_rt.mkl_get_max_threads()

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

- Daniel

---

当我使用Enthought Python发行版时，我可以将mkl作为模块导入。我想知道背后发生了什么。 - user2379410

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

- Daniel

Answer 2

长话短说，当从Python调用MKL时，请使用MKL_Set_Num_Threads及其驼峰式的伙伴。如果您不使用#include <mkl.h>，则同样适用于C语言。

---

MKL文档似乎表明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编译后，砰，再次出现分段错误。所以问题似乎不仅限于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。检查符号表发现实际上定义了四种不同的变体：

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调用约定并没有标准化。不同的编译器会以不同的方式改变函数名称：一些使用大写字母，一些使用小写字母加下划线，还有一些则完全不添加修饰符。甚至可能还有其他我不知道的方式。这个技巧允许MKL库在大多数Fortran编译器上使用而无需进行任何修改，缺点是需要对C函数进行"改名"以容纳三种Fortran调用约定的变体。

- Rufflewind

Answer 3

对于寻找跨平台和打包解决方案的人，请注意我们最近发布了threadpoolctl，这是一个模块，用于限制Python调用的C级线程池（OpenBLAS、OpenMP和MKL）中使用的线程数。有关更多信息，请参见此答案。

- Thomas Moreau

Answer 4

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

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中找到。

- Alex Maystrenko