关于Numba的线程实现的说明

关于Numba的线程实现的说明

由Numbaparallel目标执行的工作由Numba线程层执行。实际上，“线程层”是Numba内置库，可以执行所需的并发执行。在撰写本文时，有三个可用的线程层，每个线程层都通过不同的较低级别的host线程库实现。上thread线程和对于给定的应用/系统的thread线程的适当选择的更多信息可在发现 threading layer documentatio。

以下各节要关注的相关信息是，执行并行线程库中的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，则Numba将始终存在16个线程，但是其中12个线程将处于空闲状态以进行并行计算。进一步调用 set_num_threads(16) ，导致这些相同的线程在以后的计算中起作用。

添加了线程掩码，使用户可以以编程方式更改在线程层中执行工作的线程数。事实证明，线程屏蔽难以实现，需要开发一种适合用户，易于推理且可以安全实现的编程模型，并在各个线程层上具有一致的行为。

编程模型

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

实施

除了对线程层库中已经存在的用户代码进行限制以外，对用户代码没有其它限制，需要仔细考虑线程掩码的设计。无法将“线程掩码”存储在全局值中，因为同时使用线程层，可能会导致值本身出现竞争形式。涉及具有这种全局价值的各种互斥量的众多设计，最终仅通过理想实验就打破了所有这些互斥量。最终发现，在某些OpenMP实现之后，“线程掩码”最好以a实现。这意味着每个执行Numba并行函数的线程，都将具有一个线程本地存储（TLS）插槽，其中包含在线程中调度线程thread localparallel_for时，要使用的线程掩码的值。

TLS使用一个线程掩模的上述概念是相对容易实现的， get_num_threads和set_num_threads简单地需要解决的TLS时隙在给定的线程层。这也意味着可以从运行时调用中得出并行区域的执行调度get_num_threads。这是通过众所周知的且相对容易实现的C 库函数注册模式，并将其包装在内部Numba来实现的。

除了满足原始的前期线程屏蔽要求之外，还需要考虑以下一些更复杂的方案。

嵌套并行

在所有线程层中，“主线程”将调用该parallel_for 函数，然后在并行区域中，根据线程层的不同，一些其它线程将有助于完成实际工作。如果工作包含对另一个并行函数的调用（即嵌套并行性），则使调用的线程必须知道主线程的“线程掩码”是什么，以便可以将其传递到主线程中。 parallel_for在执行嵌套并行函数时调用。此行为的实现是特定于线程层的，但一般原则是“主线程”始终将线程掩码的值从其TLS插槽“发送”到线程层中，在并行区域中处于活动状态的所有线程。这些活动线程将在执行任何工作之前，使用此值更新其TLS插槽。该实现细节的最终结果是：

• 线程掩码正确传递到嵌套函数中
• 并行区域中的每个线程仍然可以安全地使用不同的掩码来调用嵌套函数，如果未明确设置，则使用从“主线程”继承的掩码
• parallel_for成功容纳具有动态调度，且线程可能在执行期间加入和离开活动池的线程层
• 任何“主线程”线程掩码都与活动线程池中线程的线程掩码的流入特性完全分离

Python线程独立调用并行函数

严格保护线程层启动顺序，确保启动既是线程安全的，又是进程安全的，并且每个进程运行一次。在具有大量threading都使用Numba的Python模块线程的系统中，启动序列中的第一个线程将正确设置其线程掩码，但是没有其他线程可以运行启动序列。这意味着其它线程将需要以其它方式设置其初始线程掩码。这是在get_num_threads调用，且不存在线程掩码时实现的，在这种情况下，线程掩码将设置为默认值。在该实现中，“不存在线程掩码”由值表示，-1。“默认线程掩码”（未设置）由值表示0。执行set_num_threads(NUMBA_NUM_THREADS)，此算子也会立即调用，因此如果有-1或0，由于遇到此结果get_num_threads()，而表示上述过程中存在错误。

操作系统fork()调用

使用TLS也是在由Linux（用于Numba使用最流行的平台）驱动，将TLS传输到子进程fork(2, 3P)clone(2)CLONE_SETTLS。

线程ID

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

注意事项

测试线程屏蔽时需要注意的一些注意事项：

• TBB后端可以选择调度少于给定掩码数的线程。因此，诸如上述测试，可能返回少于4个唯一线程。
• 工作队列后端不是线程安全的，因此尝试对其执行多线程嵌套并行处理，可能会导致死锁或其他未定义的行为。如果工作队列后端检测到嵌套的并行性，它将发出SIGABRT信号。
• 某些后端可能会重用主线程进行计算，但是不应依赖此行为（例如，如果传递异常）。

在代码生成中使用

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

import llvmlite.llvmpy.core as lc

get_num_threads = builder.module.get_or_insert_function(

lc.Type.function(lc.Type.int(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安全使用变量。最好将这样的逻辑保留在线程后端中，而不是尝试在代码生成中进行。