原文: http://numba.pydata.org/numba-doc/latest/roc/kernels.html
HSA 提供类似于 OpenCL 的执行模型。指令由一组硬件线程并行执行。在某种程度上,这类似于 _ 单指令多数据 _(SIMD)模型,但方便的是,细粒度调度对程序员是隐藏的,而不是用 SIMD 向量作为数据结构进行编程。在 HSA 中,您编写的代码将由多个线程同时执行(通常为数百或数千)。您的解决方案将通过定义 grid , workgroup 和 workitem 的线程层次结构来建模。
Numba 的 HSA 支持公开了用于声明和管理这种线程层次结构的工具。
HSA 执行模型类似于 CUDA。 HSA 在 ROC GPU 上采用的存储器模型也类似于 CUDA。 ROC GPU 专用于 GPU 内存,因此根据 CUDA 需要to_device()
和copy_to_host()
等。
以下是 CUDA 术语与 HSA 的快速映射:
workitem
相当于 CUDA 线程。workgroup
相当于 CUDA 线程块。grid
相当于 CUDA 网格。wavefront
相当于 CUDAwarp
。
_ 内核函数 _ 是一个 GPU 函数,可以从 CPU 代码中调用。它赋予它两个基本特征:
- 内核无法显式返回值;所有结果数据必须写入传递给函数的数组(如果计算一个标量,你可能会传递一个单元素数组);
- 内核在调用时显式声明其线程层次结构:即工作组的数量和每个工作组的工作项数量(请注意,虽然内核编译一次,但可以使用不同的工作组大小或网格大小多次调用它)。
乍一看,用 Numba 编写 HSA 内核看起来非常像为 CPU 编写 JIT 函数:
@roc.jit
def increment_by_one(an_array):
"""
Increment all array elements by one.
"""
# code elided here; read further for different implementations
内核通常以以下方式启动:
itempergroup = 32
groupperrange = (an_array.size + (itempergroup - 1)) // itempergroup
increment_by_one[groupperrange, itempergroup](an_array)
我们在这里注意两个步骤:
- 通过指定多个工作组(或“每个网格的工作组”)以及每个工作组的多个工作项来实例化内核。两者的产品将给出推出的工作项目总数。内核实例化是通过编译内核函数(此处为
increment_by_one
)并使用整数元组对其进行索引来完成的。 - 运行内核,通过传递输入数组(以及任何必要的单独输出数组)。默认情况下,运行内核是同步的:当内核完成执行并且数据被同步回来时,函数返回。
在声明内核所需的工作项数时,拥有两级层次结构似乎很奇怪。工作组大小(即每个工作组的工作项数)通常至关重要:
-
在软件方面,工作组大小决定共享共享内存的给定区域的线程数。
-
On the hardware side, the workgroup size must be large enough for full
占领执行单位。
为了帮助处理多维数组,HSA 允许您指定多维工作组和网格。在上面的示例中,您可以使itempergroup
和groupperrange
元组为一个,两个或三个整数。与等效大小的 1D 声明相比,这不会改变生成代码的效率或行为,但可以帮助您以更自然的方式编写算法。
运行内核时,每个线程执行一次内核函数的代码。因此,它必须知道它所在的线程,以便知道它负责哪个数组元素(复杂的算法可能定义更复杂的责任,但基本原理是相同的)。
一种方法是让线程确定它在网格和工作组中的位置,并手动计算相应的数组位置:
@roc.jit
def increment_by_one(an_array):
# workitem id in a 1D workgroup
tx = roc.get_local_id(0)
# workgroup id in a 1D grid
ty = roc.get_group_id(0)
# workgroup size, i.e. number of workitem per workgroup
bw = roc.get_local_size(0)
# Compute flattened index inside the array
pos = tx + ty * bw
# The above is equivalent to pos = roc.get_global_id(0)
if pos < an_array.size: # Check array boundaries
an_array[pos] += 1
注意
除非您确定工作组大小和网格大小是数组大小的除数,否则必须检查边界,如上所示。
get_local_id()
, get_local_size()
, get_group_id()
和 get_global_id()
是 HSA 后端为鞋底提供的特殊功能了解线程层次结构的几何以及当前工作项在该几何中的位置的目的。
numba.roc.get_local_id(dim)
获取要查询的维度的索引
返回给定维度的当前工作组中的本地工作项 ID。对于 1D 工作组,索引是一个整数,范围从 0 到包括 numba.roc.get_local_size()
不包括。
numba.roc.get_local_size(dim)
获取要查询的维度的索引
返回给定维度的工作组的大小。在实例化内核时声明该值。对于给定内核中的所有工作项,此值都是相同的,即使它们属于不同的工作组(即每个工作组都是“完整”)。
numba.roc.get_group_id(dim)
获取要查询的维度的索引
返回启动内核的工作组网格中的工作组 ID。
numba.roc.get_global_id(dim)
获取要查询的维度的索引
返回给定维度的全局工作项 ID。与 <cite>numba.roc .get_local_id()</cite>不同,此数字对于网格中的所有工作项都是唯一的。