关于 Numba 的线程实现说明

由 Numba parallel 目标呈现的工作的执行是由 Numba 线程层承担的。实际上，“线程层”是一个 Numba 内置库，能够执行所需的并发执行。在撰写本文时，有三个线程层可用，每个都是通过不同的底层本地线程库实现的。有关线程层的更多信息以及为特定应用程序/系统选择合适的线程层的适当选择，请参阅 线程层文档。

以下部分需要注意的相关信息是，在 threading 库中执行并行执行的函数是 parallel_for 函数。该函数的工作是协调和执行并行任务。

本文档中引用的相关源文件是

  • numba/np/ufunc/tbbpool.cpp

  • numba/np/ufunc/omppool.cpp

  • numba/np/ufunc/workqueue.c

    这些文件分别包含了 TBB、OpenMP 和工作队列线程池的实现。每个文件都包括了 set_num_threads()、get_num_threads() 和 get_thread_id() 函数，以及各自调度器中相关的线程掩码逻辑。请注意，基本的线程本地变量逻辑在每个文件中都有重复，而不是在它们之间共享。

  • numba/np/ufunc/parallel.py

    此文件包含适用于 Python 和 JIT 的 set_num_threads()、get_num_threads() 和 get_thread_id() 的包装器，以及将上述库加载到 Python 并启动线程池的代码。

  • numba/parfors/parfor_lowering.py

    此文件包含为并行后端生成代码的主要逻辑。在线程掩码在生成调度器代码的代码中访问，并传递给相关的后端调度器函数（见下文）。

线程掩码

作为其设计的一部分，Numba在首次运行并行执行时，不会启动超出``numba.np.ufunc.parallel._launch_threads()``初始启动的线程。这是由于在实现线程掩码之前，Numba中已经实现了线程的方式。为了保持设计的简单性，这一限制被保留下来，尽管未来可能会移除。因此，可以通过编程方式设置线程数量，但仅限于已启动线程总数或更少。这是通过“掩码”未使用的线程来实现的，使它们不进行任何工作。例如，在一台16核的机器上，如果用户调用``set_num_threads(4)``，Numba将始终保持16个线程，但其中12个线程在并行计算中将处于空闲状态。进一步调用``set_num_threads(16)``将使这些相同的线程在后续计算中工作。

线程掩码 的添加使得用户可以通过编程方式改变线程层中执行工作的线程数量。线程掩码的实现被证明具有挑战性，因为它需要开发一个适合用户的编程模型，易于理解，并且可以在各种线程层中安全地实现，行为一致。

编程模型

所选择的编程模型类似于OpenMP中的模型。选择这个模型的原因是它对许多用户来说很熟悉，范围受限且简单。使用的线程数通过调用``set_num_threads``来指定，并且可以通过调用``get_num_threads``来查询使用的线程数。这两个函数与其OpenMP的对应函数是同义的（上述限制是掩码必须小于或等于启动的线程数）。执行语义也类似于OpenMP，即一旦启动并行区域，更改线程掩码不会影响当前执行的区域，但会影响随后执行的并行区域。

实现

为了不对用户代码施加除线程层库中已存在的限制之外的任何进一步限制，需要仔细考虑线程掩码的设计。”线程掩码”不能存储在全局值中，因为并发使用线程层可能会导致对该值本身的经典形式的竞争条件。讨论了涉及此类全局值的各种类型的互斥锁的多种设计，所有这些设计最终都仅通过思想实验被打破。最终发现，在一些OpenMP实现之后，”线程掩码”最好实现为``线程局部``。这意味着执行Numba并行函数的每个线程都将有一个线程局部存储（TLS）槽，其中包含在``parallel_for``函数中调度线程时要使用的线程掩码值。

上述关于TLS用于线程掩码的概念相对容易实现，get_num_threads 和 set_num_threads 只需在给定的线程层中处理TLS槽。这也意味着并行区域的执行计划可以从运行时调用 get_num_threads 中派生出来。这是通过一个众所周知的、相对容易实现的 C 库函数注册模式，并在Numba内部实现中对其进行包装来实现的。

除了满足最初的线程掩码前置要求外，还需要考虑以下几种更复杂的场景。

嵌套并行性

在所有线程层中，一个“主线程”将调用 parallel_for 函数，然后在并行区域内，根据线程层的情况，一些额外的线程将协助完成实际工作。如果工作包含对另一个并行函数的调用（即嵌套并行性），则调用线程需要知道“主线程”的“线程掩码”是什么，以便在执行嵌套并行函数时将其传播到它所做的 parallel_for 调用中。这种行为的实现是特定于线程层的，但一般原则是“主线程”始终将其TLS槽中的线程掩码值“发送”给并行区域中所有活跃的线程层中的线程。这些活跃线程在执行任何工作之前，会用这个值更新它们的TLS槽。这种实现细节的最终结果是：

  • 线程掩码正确地传播到嵌套函数中

  • 在并行区域中，每个线程仍然可以安全地拥有不同的掩码来调用嵌套函数，如果没有显式设置，则使用从“主线程”继承的掩码。

  • 具有动态调度功能的线程层，在 parallel_for 执行期间，线程可能加入和离开活动池，这些情况都能成功处理。

  • 任何“主线程”线程掩码与活动线程池中线程的线程掩码的动态特性完全解耦。

Python 线程独立调用并行函数

线程层启动序列受到严格保护，以确保启动既线程安全又进程安全，并且每个进程只运行一次。在一个系统中，如果许多使用Python threading 模块的线程都使用Numba，第一个通过启动序列的线程将适当地设置其线程掩码，但其他线程无法运行启动序列。这意味着其他线程需要通过其他方式设置其初始线程掩码。当调用 get_num_threads 且没有线程掩码时，这可以通过将线程掩码设置为默认值来实现。在实现中，“没有线程掩码”由值 -1 表示，而“默认线程掩码”（未设置）由值 0 表示。实现还会在执行此操作后立即调用 set_num_threads(NUMBA_NUM_THREADS)，因此，如果 get_num_threads() 的结果为 -1 或 0，则表明上述过程中存在错误。

操作系统 fork() 调用

TLS 的使用也在一定程度上受到 Linux（到目前为止，Numba 使用最流行的平台）的影响，Linux 有一个 fork(2, 3P) 调用，该调用会将 TLS 传播到子进程中，参见 clone(2) 的 CLONE_SETTLS。

线程 ID

在每个线程后端中添加了一个私有的 get_thread_id() 函数，该函数返回每个线程的唯一ID。可以通过 numba.np.ufunc.parallel._get_thread_id() 从Python中访问（也可以在JIT编译的函数内部使用）。线程ID函数对于测试线程掩码行为是否正确很有用，但不应在测试之外使用。例如，可以调用 set_num_threads(4) 然后在并行区域中收集所有唯一的 _get_thread_id() 来验证是否只有4个线程在运行。

注意事项

测试线程掩码时需要注意的一些事项：

  • TBB 后端可能会选择调度少于给定掩码数量的线程。因此，如上所述的测试可能会返回少于 4 个的唯一线程。

  • workqueue 后端不是线程安全的，因此尝试使用它进行多线程嵌套并行可能会导致死锁或其他未定义行为。如果检测到嵌套并行，workqueue 后端将引发 SIGABRT 信号。

  • 某些后端可能会重用主线程进行计算，但这种行为不应被依赖（例如，在传播异常时）。

在代码生成中的使用

在代码生成中使用 get_num_threads 的一般模式是

from llvmlite import ir as llvmir

get_num_threads = cgutils.get_or_insert_function(builder.module
    llvmir.FunctionType(llvmir.IntType(types.intp.bitwidth), []),
    name="get_num_threads")

num_threads = builder.call(get_num_threads, [])

with cgutils.if_unlikely(builder, builder.icmp_signed('<=', num_threads,
                                              num_threads.type(0))):
    cgutils.printf(builder, "num_threads: %d\n", num_threads)
    context.call_conv.return_user_exc(builder, RuntimeError,
                                              ("Invalid number of threads. "
                                               "This likely indicates a bug in Numba.",))

# Pass num_threads through to the appropriate backend function here

查看 numba/parfors/parfor_lowering.py 中的代码。

对 num_threads 小于等于 0 的检查并非绝对必要，但它可以在线程掩码逻辑中存在错误时，防止意外的错误行为。

num_threads 变量应传递给适当的后端函数，如 do_scheduling 或 parallel_for。如果它以其他方式使用，而不是传递给后端函数，则应考虑上述因素，以确保 num_threads 变量的使用是安全的。最好将此类逻辑保留在线程后端中，而不是尝试在代码生成中实现。

并行块大小详情

在某些情况下，实际的并行工作块大小可能与通过 numba.set_parallel_chunksize() 请求的块大小不同。首先，如果基于指定块大小所需的块数少于配置的线程数，那么 Numba 将使用所有配置的线程来执行并行区域。在这种情况下，实际块大小将小于请求的块大小。其次，由于截断，在迭代次数略小于块大小的倍数的情况下（例如，14 次迭代和指定的块大小为 5），实际块大小将大于指定的块大小。如给定示例中，块数将为 2，实际块大小将为 7（即 14 / 2）。最后，由于 Numba 将 N 维迭代空间划分为 N 维（超）矩形块，可能存在没有 N 个整数因子其乘积等于块大小的情况。在这种情况下，一些块的面积/体积将大于块大小，而其他块将小于指定的块大小。