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

Question

我想从 Python 调用一些 OpenMP 程序，更改线程数 (OMP_NUM_THREADS) 及其绑定 (OMP_PLACES, OMP_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

有什么想法吗？

Answer 1

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

[OpenMP 4.5] 4.环境变量

程序启动后对环境变量的修改，即使由程序本身修改，也会被 OpenMP 实现忽略。但是，在 OpenMP 程序执行期间，可以通过使用适当的指令子句或 OpenMP API 例程来修改某些 ICV 的设置。

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

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

如果您必须使用共享库并控制动态绑定，您可以在由 python 调用的共享库中手动使用 sched_setaffinity 做一些魔术。

gcc 和 intel 运行时行为不同的原因，可能是因为你在加载库后设置了环境变量，并且它们的初始化方式不同。