循环遍历 CUDA + Numba 中的行

Question

我正在做一个动态编程问题，其中每一行都依赖于前一行，即ith 行依赖于i-1th 行，从第 1 行开始（0 索引）。

我正在使用numba 的cuda.jit 来加速这个过程。我的方法是让一维线程块等于我的矩阵的宽度（w）（其中矩阵是灰度图像，即尺寸为h x w 的二维图像）。 因此每个线程只负责一个列。

如何正确循环每一行？在下面的代码中，我沿着行大步前进，因为每个线程处理一列。

@cuda.jit
def forward_energy(im, energy, m):
  row, col = cuda.grid(2)
  xstride, _ = cuda.gridsize(2)

  height, width = im.shape[0], im.shape[1]

  if row >= height or col >= width: return


  for i in range(row, im.shape[0], xstride):
      # example code below is dependent on the previous row of `im`
      energy[i, col] = min(im[i - 1, col], im[i - 1, col - 1])
      im[i, col] = # update current row

这是对的吗？据我所知，每个线程都是异步执行的，因此线程 1（处理第 1 列）可能位于第 5 行，而线程 2（位于第 2 列）可能仍位于第 3 行。如何确保每一行在继续之前完成，如果用 cuda 可能的话？

Answer 1

为了完整起见，这是我最后的做法：

1. 我从上面的代码中删除了跨步。
2. 添加了一个调用新内核的for 循环：

for row in range(1, image.shape[0]):
  forward_energy[(1,),(1, image.shape[1])](row, imd, ed, md)
  cuda.synchronize()

并更新了函数参数以接受该行：

def foward_energy(row, im, energy, m):
    # algo stuff

这篇thread 对异步内核调用及其行为方式很有帮助。 我不完全确定在运行下一个循环之前需要cuda.synchronize()，但我宁愿谨慎行事。

更新（2020 年 12 月 23 日）：

虽然上述解决方案运行良好，但速度非常缓慢。值得庆幸的是，Numba 即将在0.53 发布合作组（根据他们的 github 里程碑）。如果您在 0.53 发布后阅读此内容，请直接查看下面的代码；如果没有，您可以将最新的 numba 与此 setup guide 一起使用。这将使您能够访问 numba 的 cooperative groups、cuda.cg，从而允许您执行网格范围的同步。

该解决方案基本上消除了多次内核启动的需要。之前，for-loop 是从内核中提升出来的，现在代码在内核中有 for-loop：

threads_per_block = 128
blocks_per_grid = math.ceil(image.shape[1] / threads_per_block)

循环内核中的行：

@cuda.jit
def forward_energy(im, *args):
    col = cuda.grid(1)
    g = cuda.cg.this_grid()
    for row in range(1, im.shape[0]):
         # do stuff
         g.sync()

调用内核一次：

forward_energy[blocks_per_grid, threads_per_block](im, *args)

就是这样。比以前管理顺序行的方式快得多。据我了解，g.sync() 所做的是每个线程都需要调用g.sync() 才能进入下一个循环。这具有强制每一行按顺序执行的效果。好东西。