协作组

支持的功能

Numba 的协作组支持目前提供了网格组和网格同步，以及协作内核启动。

协作组在 Linux 和 Windows 上受支持，适用于处于 TCC 模式 的设备。

使用网格组

要获取当前的网格组，请使用 cg.this_grid() 函数：

g = cuda.cg.this_grid()

同步网格是通过网格组的 sync() 方法完成的：

g.sync()

合作启动

与 CUDA C/C++ API 不同，协作启动使用与普通内核启动相同的语法调用 - Numba 会根据内核中是否同步了网格组来自动确定是否需要协作启动。

协作启动的网格大小限制比正常启动更为严格 - 网格的大小不能超过设备上启动时的最大活动块数。要获取具有给定块大小和动态共享内存需求的内核协作启动的最大网格大小，请使用内核重载的 max_cooperative_grid_blocks() 方法：

_Kernel.max_cooperative_grid_blocks(blockdim, dynsmemsize=0)[源代码]

计算在当前上下文中，对于给定的块和动态共享内存大小，可以在协作网格中启动的此内核的最大块数。

参数:

  • blockdim – 块的维度，对于1D块可以是一个标量，对于2D或3D块可以是一个元组。

  • dynsmemsize – 动态共享内存大小，以字节为单位。

返回:

网格中块的最大数量。

这可以用来确保内核以不超过最大块数的数量启动。超过协同启动的最大块数将导致 CUDA_ERROR_COOPERATIVE_LAUNCH_TOO_LARGE 错误。

应用程序和示例

网格组同步可以用来在网格中的所有线程之间实现一个全局屏障——这包括将全局缩减为一个单一值，或者使用整个网格按顺序循环遍历大型矩阵的行，同时并行操作列元素。

在下面的例子中，行是由网格顺序写入的。网格中的每个线程从前一行中读取其*相反*线程写入的值。需要进行网格同步以确保网格中的线程不会超过其他块中的线程，或者未能看到其相反线程的更新。

首先，我们将定义我们的内核：

来自 numba/cuda/tests/doc_example/test_cg.py 中的 test_grid_sync
 1from numba import cuda, int32
 2import numpy as np
 3
 4sig = (int32[:,::1],)
 5
 6@cuda.jit(sig)
 7def sequential_rows(M):
 8    col = cuda.grid(1)
 9    g = cuda.cg.this_grid()
10
11    rows = M.shape[0]
12    cols = M.shape[1]
13
14    for row in range(1, rows):
15        opposite = cols - col - 1
16        # Each row's elements are one greater than the previous row
17        M[row, col] = M[row - 1, opposite] + 1
18        # Wait until all threads have written their column element,
19        # and that the write is visible to all other threads
20        g.sync()

然后创建一些空输入数据并确定网格和块的大小：

来自 numba/cuda/tests/doc_example/test_cg.py 中的 test_grid_sync
1# Empty input data
2A = np.zeros((1024, 1024), dtype=np.int32)
3# A somewhat arbitrary choice (one warp), but generally smaller block sizes
4# allow more blocks to be launched (noting that other limitations on
5# occupancy apply such as shared memory size)
6blockdim = 32
7griddim = A.shape[1] // blockdim

最后我们启动内核并打印结果：

来自 numba/cuda/tests/doc_example/test_cg.py 中的 test_grid_sync
 1# Kernel launch - this is implicitly a cooperative launch
 2sequential_rows[griddim, blockdim](A)
 3
 4# What do the results look like?
 5# print(A)
 6#
 7# [[   0    0    0 ...    0    0    0]
 8#  [   1    1    1 ...    1    1    1]
 9#  [   2    2    2 ...    2    2    2]
10#  ...
11#  [1021 1021 1021 ... 1021 1021 1021]
12#  [1022 1022 1022 ... 1022 1022 1022]
13#  [1023 1023 1023 ... 1023 1023 1023]]

sequential_rows 的最大网格尺寸可以通过以下方式查询：

overload = sequential_rows.overloads[(int32[:,::1],)
max_blocks = overload.max_cooperative_grid_blocks(blockdim)
print(max_blocks)
# 1152 (e.g. on Quadro RTX 8000 with Numba 0.52.1 and CUDA 11.0)