线程层

本节介绍 Numba 线程层，它是内部用于通过 CPU 的 parallel 目标执行并行计算的库，具体如下：

  • 在 @jit 和 @njit 中使用 parallel=True kwarg。

  • 在 @vectorize 和 @guvectorize 中使用 target='parallel' kwarg。

注意

如果代码库不使用 threading 或 multiprocessing 模块（或任何其他形式的并行性），Numba 附带的线程层默认设置将运行良好，无需进一步操作！

有哪些线程层可用？

有三种线程层可用，它们命名如下：

  • tbb - 由 Intel TBB 支持的线程层。

  • omp - 由 OpenMP 支持的线程层。

  • workqueue - 一个简单的内置工作共享任务调度器。

实际上，唯一保证存在的线程层是 workqueue。omp 层需要存在合适的 OpenMP 运行时库。tbb 层需要存在 Intel 的 TBB 库，可以通过 conda 命令获取：

$ conda install tbb

如果您使用 pip 安装了 Numba，可以通过运行以下命令启用 TBB：

$ pip install tbb

注意

Numba 搜索和加载线程层的默认方式能够容忍缺少库、不兼容的运行时等情况。

设置线程层

线程层通过环境变量 NUMBA_THREADING_LAYER 或通过赋值给 numba.config.THREADING_LAYER 来设置。如果使用编程方式设置线程层，则必须在任何基于 Numba 的并行目标编译发生之前在逻辑上进行。选择线程层有两种方法：第一种是选择在各种形式的并行执行下都安全的线程层；第二种是通过线程层名称（例如 tbb）进行显式选择。

设置线程层选择优先级

默认情况下，线程层按 'tbb'、'omp'，然后 'workqueue' 的顺序搜索。为了在保持基于可用性选择线程层的情况下更改此搜索顺序，可以使用环境变量 NUMBA_THREADING_LAYER_PRIORITY。

请注意，它也可以通过 numba.config.THREADING_LAYER_PRIORITY 设置。与 numba.config.THREADING_LAYER 类似，它必须在任何基于 Numba 的并行目标编译发生之前在逻辑上进行。

例如，要指示 Numba 在可用时首先选择 omp，然后是 tbb 等等，请将环境变量设置为 NUMBA_THREADING_LAYER_PRIORITY="omp tbb workqueue"。或者通过编程方式，numba.config.THREADING_LAYER_PRIORITY = ["omp", "tbb", "workqueue"]。

为安全的并行执行选择线程层

并行执行主要通过核心 Python 库以四种形式派生（前三种也适用于通过其他方式使用并行执行的代码！）：

  • 来自 threading 模块的 threads。

  • 通过 spawn 从 multiprocessing 模块 spawn 进程（Windows 上的默认设置，仅在 Unix 上的 Python 3.4+ 中可用）

  • 通过 fork 从 multiprocessing 模块 fork 进程（Unix 上的默认设置）。

  • 通过使用 forkserver 从 multiprocessing 模块 fork 进程（仅在 Unix 上的 Python 3 中可用）。本质上是先启动一个新进程，然后根据请求从该新进程中派生（fork）子进程。

任何与这些并行形式一起使用的库都必须在给定范式下表现出安全行为。因此，线程层选择方法旨在以一种简单、跨平台和环境容错的方式，为给定范式选择安全的线程层库。可以提供给设置机制的选项如下：

  • default 不提供特定的安全保证，是默认值。

  • safe 既是 fork 安全的也是线程安全的，这需要安装 tbb 包（Intel TBB 库）。

  • forksafe 提供一个 fork 安全的库。

  • threadsafe 提供一个线程安全的库。

要查找已选择的线程层，可以在并行执行后调用函数 numba.threading_layer()。例如，在没有安装 TBB 的 Linux 机器上：

from numba import config, njit, threading_layer
import numpy as np

# set the threading layer before any parallel target compilation
config.THREADING_LAYER = 'threadsafe'

@njit(parallel=True)
def foo(a, b):
    return a + b

x = np.arange(10.)
y = x.copy()

# this will force the compilation of the function, select a threading layer
# and then execute in parallel
foo(x, y)

# demonstrate the threading layer chosen
print("Threading layer chosen: %s" % threading_layer())

这将产生

Threading layer chosen: omp

这是有道理的，因为 Linux 上存在的 GNU OpenMP 是线程安全的。

选择一个命名的线程层

高级用户可能希望为其用例选择特定的线程层，这通过直接将线程层名称提供给设置机制来完成。选项和要求如下：

线程层名称

平台

要求

tbb

所有

tbb 包 ($ conda install tbb)

omp

Linux

Windows

OSX

GNU OpenMP 库（很可能已经存在）

MS OpenMP 库（很可能已经存在）

intel-openmp 包或 llvm-openmp 包（按名称 conda install）。

workqueue

所有

None

如果线程层未能正确加载，Numba 将检测到此问题并提供解决提示。另请注意，Numba 诊断命令 numba -s 有一个 __Threading Layer Information__ 部分，报告当前环境中线程层的可用性。

额外注意事项

线程层与 CPython 内部和系统级库之间存在相当复杂的交互，还需要注意一些事项：

  • 安装 Intel 的 TBB 库将大大扩展线程层选择过程中可用的选项。

  • 在 Linux 上，omp 线程层不是 fork-safe 的，因为 GNU OpenMP 运行时库 (libgomp) 不支持 fork-safe。如果在使用 omp 线程层的程序中发生 fork，则存在一个检测机制，它将尝试优雅地终止分叉的子进程并将错误消息打印到 STDERR。

  • 在支持 fork(2) 系统调用的系统上，如果正在使用 TBB 支持的线程层，并且 fork 调用是由启动 TBB 的线程（通常是主线程）以外的线程发出的，则会导致未定义行为，并且会在 STDERR 上显示警告。由于 spawn 本质上是 fork 后跟 exec，因此从非主线程 spawn 是安全的，但由于这无法与单纯的 fork 调用区分开来，因此警告消息仍会显示。

  • 在 OSX 上，需要 intel-openmp 包才能启用基于 OpenMP 的线程层。

设置线程数

Numba 使用的线程数基于可用 CPU 核数（参见 numba.config.NUMBA_DEFAULT_NUM_THREADS），但可以通过 NUMBA_NUM_THREADS 环境变量覆盖。

Numba 启动的总线程数在变量 numba.config.NUMBA_NUM_THREADS 中。

对于某些用例，可能希望将线程数设置为较低的值，以便 Numba 可以与更高级别的并行性一起使用。

线程数可以在运行时使用 numba.set_num_threads() 动态设置。请注意，set_num_threads() 仅允许将线程数设置为小于 NUMBA_NUM_THREADS 的值。Numba 总是启动 numba.config.NUMBA_NUM_THREADS 个线程，但 set_num_threads() 会使其屏蔽未使用的线程，从而这些线程不参与计算。

Numba 当前使用的线程数可以通过 numba.get_num_threads() 访问。这两个函数都可以在 JIT 编译的函数内部工作。

限制线程数示例

在此示例中，假设我们运行的机器有 8 个核心（因此 numba.config.NUMBA_NUM_THREADS 将是 8）。假设我们想使用 @njit(parallel=True) 运行一些代码，但我们也想在 4 个不同的进程中并发运行我们的代码。在默认线程数下，每个 Python 进程将运行 8 个线程，总计 4*8 = 32 个线程，这对于我们的 8 个核心来说是超额订阅。我们应该将每个进程限制为 2 个线程，这样总数将是 4*2 = 8，这与我们的物理核心数相匹配。

有两种方法可以做到这一点。一种是将 NUMBA_NUM_THREADS 环境变量设置为 2。

$ NUMBA_NUM_THREADS=2 python ourcode.py

然而，这种方法有两个缺点：

  1. NUMBA_NUM_THREADS 必须在导入 Numba 之前设置，理想情况下是在启动 Python 之前。一旦导入 Numba，就会读取环境变量，并且该线程数将作为 Numba 启动的线程数被锁定。

  2. 如果我们想稍后增加进程使用的线程数，则无法实现。NUMBA_NUM_THREADS 设置了为进程启动的最大线程数。调用 set_num_threads() 时提供的值大于 numba.config.NUMBA_NUM_THREADS 将导致错误。

这种方法的优点是，我们可以在不更改代码的情况下从进程外部完成它。

另一种方法是在我们的代码中使用 numba.set_num_threads() 函数。

from numba import njit, set_num_threads

@njit(parallel=True)
def func():
    ...

set_num_threads(2)
func()

如果我们在执行并行代码之前调用 set_num_threads(2)，它与使用 NUMBA_NUM_THREADS=2 调用进程具有相同的效果，即并行代码将只在 2 个线程上执行。然而，我们稍后可以调用 set_num_threads(8) 将线程数增加回默认大小。我们也不必担心在 Numba 导入之前进行设置。它只需要在并行函数运行之前调用即可。

获取线程 ID

在某些情况下，访问当前在 Numba 中执行并行区域的线程的唯一标识符可能很有益。为此，Numba 提供了 numba.get_thread_id() 函数。此函数是 OpenMP 的 omp_get_thread_num 函数的推论，并返回一个介于 0（包含）和上述配置线程数（不包含）之间的整数。

API 参考

numba.config.NUMBA_NUM_THREADS

Numba 启动的总（最大）线程数。

默认为 numba.config.NUMBA_DEFAULT_NUM_THREADS，但可以通过 NUMBA_NUM_THREADS 环境变量覆盖。

numba.config.NUMBA_DEFAULT_NUM_THREADS

系统上可用的 CPU 核心数（由 len(os.sched_getaffinity(0)) 确定，如果操作系统支持，否则由 multiprocessing.cpu_count() 确定）。这是 numba.config.NUMBA_NUM_THREADS 的默认值，除非设置了 NUMBA_NUM_THREADS 环境变量。

numba.set_num_threads(n)

设置用于并行执行的线程数。

默认情况下，使用所有 numba.config.NUMBA_NUM_THREADS 线程。

此功能通过屏蔽未使用的线程来工作。因此，线程数 n 必须小于或等于 NUMBA_NUM_THREADS，即启动的总线程数。详情请参阅其文档。

此函数可以在 JIT 编译的函数内部使用。

参数
n: 线程数。必须介于 1 和 NUMBA_NUM_THREADS 之间。

另请参阅

get_num_threads, numba.config.NUMBA_NUM_THREADS
numba.config.NUMBA_DEFAULT_NUM_THREADS, NUMBA_NUM_THREADS
numba.get_num_threads()

获取用于并行执行的线程数。

默认情况下（如果从未调用 set_num_threads()），将使用所有 numba.config.NUMBA_NUM_THREADS 线程。

此数字小于或等于启动的总线程数 numba.config.NUMBA_NUM_THREADS。

此函数可以在 JIT 编译的函数内部使用。

返回
线程数。

另请参阅

set_num_threads, numba.config.NUMBA_NUM_THREADS
numba.config.NUMBA_DEFAULT_NUM_THREADS, NUMBA_NUM_THREADS
numba.get_thread_id()

返回每个线程的唯一 ID，范围从 0（包含）到 get_num_threads()（不包含）。