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朋友。如果您不#include <mkl.h>，则同样适用于C.

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。检查符号表显示实际上定义了四种不同的变体：

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函数需要被“修复”以为Fortran调用约定的3个变体腾出空间。 / p>

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，该模块可限制python（OpenBLAS在C级线程池中使用的线程数， OpenMP和MKL）。有关更多信息，请参见此answer。

将mkl_set_num_threads与numpy一起使用

4 个答案: