在 CUDA + Numba 中循环遍历行

Question

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

我正在使用 numba 的 cuda.jit 来加速它。我的方法是让一个 1D 线程块等于我的矩阵的宽度 (w)（其中矩阵是一个灰度图像，即尺寸 h x w 的 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

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

我从上面的代码中删除了步幅。
添加了调用新内核的 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() 在下一个循环的运行之前是必要的，但我宁愿谨慎行事。

更新 (23/12/2020):

虽然上面的解决方案工作正常，但速度非常慢。值得庆幸的是，Numba 即将在 0.53 发布合作组（根据他们的 github 里程碑）。如果您在 0.53 发布后阅读本文，请直接阅读下面的代码；如果没有，您可以将最新的 numba 与此 setup guide. That would give you access to numba's 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() 才能进入下一个循环。这具有强制按顺序执行每一行的效果。好东西。

在 CUDA + Numba 中循环遍历行

Looping over rows in CUDA + Numba

python

cuda

numba