在CUDA+Numba中的行上循环

Question

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

我用numba的cuda.jit来加速这个过程。我的方法是让一个1D线程块等于我的矩阵的宽度（w）（其中矩阵是灰度图像，即尺寸为hxw的2D img）因此，每个线程负责一个列

如何正确循环每一行？在下面的代码中，由于每个线程处理一个列，因此我在行中取得了长足的进步

@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循环被提升出内核，代码现在在内核中有for循环：

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()才能进入下一个循环。这会强制按顺序执行每一行。好东西