Использование mkl_set_num_threads с numpy

Ответ 1

Офион привел меня правильно. Несмотря на документацию, нужно передать параметр 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-функций записываются в капитал-case, которые ожидают параметры по значению:

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

Ответ 2

Короче говоря, используйте MKL_Set_Num_Threads и его друзей CamelCased при вызове MKL из Python. То же самое относится к C, если вы не #include <mkl.h>.

Документация 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 и стрелы, 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_? Это версия Fortran MKL_Set_Num_Threads! Как мы в итоге позвонили в версию Fortran? (Имейте в виду, что соглашение о вызове Fortran требует, чтобы аргументы передавались как указатели, а не по значению.)

Оказывается, документация была полным фасадом. Если вы действительно просматриваете файлы заголовков для последних версий MKL, вы найдете это довольно маленькое определение:

void    MKL_Set_Num_Threads(int nth);
#define mkl_set_num_threads         MKL_Set_Num_Threads

... и теперь все имеет смысл! Правильная функция call (для кода 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_
…

Почему Intel поставила четыре разных варианта одной функции, несмотря на то, что в документации есть только варианты C и Fortran? Я не знаю точно, но я подозреваю, что он совместим с разными компиляторами Fortran. Вы видите, что соглашение о назначении Фортрана не стандартизировано. Различные компиляторы будут поменять имена по-разному:

• некоторые используют верхний регистр,
• некоторые используют нижний регистр с завершающим подчеркиванием и
• некоторые используют нижний регистр без каких-либо украшений.

Могут быть и другие способы, о которых я не знаю. Этот трюк позволяет библиотеке MKL использоваться с большинством компиляторов Fortran без каких-либо изменений, недостатком является то, что функции C должны быть "искалечены", чтобы освободить место для трех вариантов соглашения о назначении Fortran.

Ответ 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())

Код также доступен в этой сути.

Ответ 4

Для тех, кто ищет кроссплатформенное и пакетное решение, обратите внимание, что мы недавно выпустили threadpoolctl, модуль для ограничения количества потоков, используемых в пулах потоков уровня C, называемых python (OpenBLAS, OpenMP и MKL). Смотрите этот ответ для получения дополнительной информации.