从 Python 和环境变量调用 OpenMP C 库

Calling an OpenMP C library from Python and environment variables

我想从 Python 调用一些 OpenMP 程序,更改线程数 (OMP_NUM_THREADS) 及其绑定 (OMP_PLACESOMP_PROC_BIND)。我写了这个程序:

import os
from ctypes import cdll

lib = cdll.LoadLibrary("/home/fayard/Desktop/libf.so")

nb_socket = 2
nb_core_per_socket = 14
nb_thread_per_core = 2

n = nb_socket * nb_core_per_socket * nb_thread_per_core

for nb_thread in range(1, n + 1):
    os.environ['OMP_NUM_THREADS'] = str(nb_thread)
    print("nb_thread: {}, omp_get_num_threads: {}".format(
        nb_thread, lib.num_threads()))

OpenMP 库如下:

#include <omp.h>

extern "C" {

int num_threads() {
  int ans;
#pragma omp parallel
  {
#pragma omp single
    ans = omp_get_num_threads();
  }

  return ans;
}

}

并编译为:

g++ -c -fPIC -fopenmp f.cpp -o f.o
g++ -shared -fopenmp -Wl,soname,libf.so -o libf.so f.o

当我 运行 python program.py 时,我得到:

nb_thread: 1, omp_get_num_threads: 56
...
nb_thread: 56, omp_get_num_threads: 56

这不是我想要的!我还意识到,当使用具有完全相同参数的英特尔编译器编译时,我得到:

nb_thread: 1, omp_get_num_threads: 1
...
nb_thread: 56, omp_get_num_threads: 1

有什么问题吗?

环境变量只控制内部控制变量的初始设置。

[OpenMP 4.5] 4. Environment variables

Modifications to the environment variables after the program has started, even if modified by the program itself, are ignored by the OpenMP implementation. However, the settings of some of the ICVs can be modified during the execution of the OpenMP program by the use of the appropriate directive clauses or OpenMP API routines.

您可以围绕 omp_set_num_threads 编写一个小包装器,但是您不能动态更改绑定。

不幸的是,unload shared librareies in ctypes 没有干净的解决方案。另一种方法是 运行 使用 subprocess 的实际程序,而不是加载库,但是你有一个不同的界面。

如果您必须使用共享库并控制动态绑定,您可以在 python.

调用的共享库中手动使用 sched_setaffinity 做一些魔术

gcc 和 intel 运行 次表现不同的原因,可能是因为您在加载库后设置了环境变量,并且它们的初始化完成方式不同。