Numba 线程实现说明

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

以下部分需要注意的相关信息是，线程库中执行并行操作的函数是 parallel_for 函数。该函数的作用是编排和执行并行任务。

本文中引用的相关源文件包括：

  • numba/np/ufunc/tbbpool.cpp

  • numba/np/ufunc/omppool.cpp

  • numba/np/ufunc/workqueue.c

    这些文件分别包含 TBB、OpenMP 和 workqueue 线程池的实现。每个文件都包含 set_num_threads()、get_num_threads() 和 get_thread_id() 函数，以及在其各自调度器中进行线程掩码的相关逻辑。请注意，基本的线程局部变量逻辑在这些文件中是重复的，并且不共享。

  • numba/np/ufunc/parallel.py

    此文件包含 set_num_threads()、get_num_threads() 和 get_thread_id() 的 Python 和 JIT 兼容封装器，以及将上述库加载到 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，则表明上述流程中存在 bug。

操作系统 fork() 调用

TLS 的使用也部分受到 Linux（迄今为止 Numba 使用最流行的平台）的驱动，因为它具有 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 的防御性检查并非严格必要，但它可以防止在线程掩码逻辑存在 bug 时出现意外的错误行为。

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

并行块大小详情

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