协同组

CUDA 内置目标弃用通知

Numba 内置的 CUDA 目标已弃用，后续开发已移至 NVIDIA numba-cuda 包。请参阅内置 CUDA 目标弃用和维护状态。

支持的功能

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 – 块维度，一维块为标量，二维或三维块为元组。
dynsmemsize – 动态共享内存大小（以字节为单位）。

返回

网格中的最大块数。

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

应用和示例

网格组同步可用于实现跨网格中所有线程的全局屏障——这方面的应用包括全局归约到单个值，或使用整个网格并行操作列元素来顺序遍历大型矩阵的行。

在以下示例中，行由网格顺序写入。网格中的每个线程从其“对立”线程写入的上一行读取一个值。需要网格同步以确保网格中的线程不会领先于其他块中的线程，或无法看到其对立线程的更新。

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

来自 numba/cuda/tests/doc_example/test_cg.py 的 test_grid_sync

from numba import cuda, int32
import numpy as np

sig = (int32[:,::1],)

@cuda.jit(sig)
def sequential_rows(M):
    col = cuda.grid(1)
    g = cuda.cg.this_grid()

    rows = M.shape[0]
    cols = M.shape[1]

    for row in range(1, rows):
        opposite = cols - col - 1
        # Each row's elements are one greater than the previous row
        M[row, col] = M[row - 1, opposite] + 1
        # Wait until all threads have written their column element,
        # and that the write is visible to all other threads
        g.sync()

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

来自 numba/cuda/tests/doc_example/test_cg.py 的 test_grid_sync

# Empty input data
A = np.zeros((1024, 1024), dtype=np.int32)
# A somewhat arbitrary choice (one warp), but generally smaller block sizes
# allow more blocks to be launched (noting that other limitations on
# occupancy apply such as shared memory size)
blockdim = 32
griddim = A.shape[1] // blockdim

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

来自 numba/cuda/tests/doc_example/test_cg.py 的 test_grid_sync

# Kernel launch - this is implicitly a cooperative launch
sequential_rows[griddim, blockdim](A)

# What do the results look like?
# print(A)
#
# [[   0    0    0 ...    0    0    0]
#  [   1    1    1 ...    1    1    1]
#  [   2    2    2 ...    2    2    2]
#  ...
#  [1021 1021 1021 ... 1021 1021 1021]
#  [1022 1022 1022 ... 1022 1022 1022]
#  [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)