如何理解Numba CUDA中add2函数的stride（步长）机制

这问题问得好！我刚接触CUDA时也对这个步长机制产生过类似困惑，咱们一步步拆解add2的逻辑，你就能彻底明白啦。

首先先回顾你给出的代码：

from numba import cuda
import numpy as np

@cuda.jit
def add1(x, y, r):
    i = cuda.grid(1)
    if i < len(r):
        r[i] = x[i] + y[i]

@cuda.jit
def add2(x, y, r):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)
    for i in range(start, x.shape[0], stride):
        r[i] = x[i] + y[i]

if __name__ == '__main__':
    n = 50000000
    x = np.arange(n)
    y = 2*x
    out = np.empty_like(x)
    threads_per_block = 128
    blocks_per_grid = n//threads_per_block + 1
    add1[blocks_per_grid, threads_per_block](x, y, out)
    d = abs(out - x - y).sum()
    print(d)
    out = np.empty_like(x)
    add2[blocks_per_grid, threads_per_block](x, y, out)
    d = abs(out - x - y).sum()
    print(d)

关键函数的含义 #

要理解add2，首先得搞清楚两个CUDA API的返回值：

• cuda.grid(1)：返回当前线程在整个GPU网格中的全局索引，而不是线程块内的局部索引。比如你设置threads_per_block=128、blocks_per_grid=4，全局索引范围就是0 ~ 4*128-1 = 511。
• cuda.gridsize(1)：返回整个网格的总线程数，也就是blocks_per_grid * threads_per_block，上面的例子里就是512。

add2的执行逻辑 #

你之前的误解在于对stride的取值判断——stride不是数组长度，而是总线程数！

add2里的循环for i in range(start, x.shape[0], stride)，每个线程会从自己的全局索引start开始，以总线程数为步长，依次处理数组元素。举个小例子直观感受：
假设总线程数stride=4，数组长度n=10：

• 全局索引为0的线程：处理索引0,4,8
• 全局索引为1的线程：处理索引1,5,9
• 全局索引为2的线程：处理索引2,6
• 全局索引为3的线程：处理索引3,7

这样所有数组元素（0~9）都被恰好一个线程处理一次，完全没有重复计算！

对比add1与add2的差异 #

• add1是单线程单元素模式：每个线程只处理一个元素，但如果数组长度大于总线程数，需要额外判断i < len(r)避免越界。
• add2是网格跨步循环模式：每个线程处理一组间隔为总线程数的元素，循环自动终止于数组末尾，无需额外边界判断。这种方式更适合处理超大数组——当数组长度远大于GPU能同时启动的线程数时，每个线程处理多个元素，能更高效地利用GPU计算资源，减少线程调度的开销。

验证方法 #

如果你想亲自验证，可以把n改成小数值（比如10），在add2的循环里加一句print(i)，然后运行代码，就能看到每个索引只被打印一次啦。

内容的提问来源于stack exchange，提问作者Qiang Zhang