如何在 numba CUDA 中对行进行切片？

Question

我是 Numba 的初学者。我很难在 GPU 中重新排列数组的行。

例如，在 Numba CPU 中，这可以通过

from numba import njit
import numpy as np

@njit
def numba_cpu(A, B, ind):
    for i, t in enumerate(ind):
        B[i, :] = A[t, :]

ind = np.array([3, 2, 0, 1, 4])
A = np.random.rand(5, 3)
B = np.zeros((5, 3))
numba_cpu(A, B, ind)

但它不适用于 cuda.jit

from numba import cuda
import numpy as np

@cuda.jit
def numba_gpu(A, B, ind):
    for i, t in enumerate(ind):
        B[i, :] = A[t, :]

d_ind = cuda.to_device(np.array([3, 2, 0, 1, 4]))
d_A = cuda.to_device(np.random.rand((5, 3)))
d_B = cuda.to_device(np.zeros((5, 3)))
numba_gpu[16,16](d_A, d_B, d_ind)

程序失败并出现很多异常，并显示“需要 NRT 但未启用”。

当然，我可以使用嵌套循环逐个复制条目，但它看起来很糟糕，因为我知道 a 行在连续的内存中。即使是类似 C 语言的 memcpy 也会更好。但似乎 Numba 不支持memcpy。

Answer 1

Numba 的设备不支持这些切片运算符是正确的。根本问题是切片操作需要一个中间数组构造，而 Numba 编译器目前无法做到这一点。

可能有两种替代方法可以做到这一点：

1. 使用单线程在源和目标之间复制一行数据（numba_gpu1如下所示）
2. 使用单个块在源和目标之间复制一行数据。这可以利用跨步循环设计模式来提高内存合并和缓存一致性，并且应该在非平凡大小下表现更好（numba_gpu2 如下所示，用于行主要有序数据）。

在代码中，这看起来像：

from numba import cuda
import numpy as np

@cuda.jit
def numba_gpu1(A, B, ind):
  idx = cuda.grid(1)
  if idx < len(ind):
    t = ind[idx]
    for k in range(A.shape[1]):
       B[idx, k] = A[t, k]

@cuda.jit
def numba_gpu2(A, B, ind):
  idx = cuda.grid(1)
  if idx < len(ind):
    t = ind[idx]
    k = cuda.threadIdx.y
    while k < A.shape[1]:
       B[idx, k] = A[t, k]
       k += cuda.blockDim.y

def test1():
  d_ind = cuda.to_device(np.array([3, 2, 0, 1, 4], dtype=np.int32))
  d_A = cuda.to_device(np.random.rand(5, 3).astype(np.float32))
  d_B = cuda.to_device(np.zeros((5, 3), dtype=np.float32))

  numba_gpu1[1,32](d_A, d_B, d_ind)

  ind = d_ind.copy_to_host()
  A = d_A.copy_to_host()
  B = d_B.copy_to_host()

  return np.allclose(A[ind],B)

def test2():
  d_ind = cuda.to_device(np.array([3, 2, 0, 1, 4], dtype=np.int32))
  d_A = cuda.to_device(np.random.rand(5, 3).astype(np.float32))
  d_B = cuda.to_device(np.zeros((5, 3), dtype=np.float32))

  numba_gpu2[5,32](d_A, d_B, d_ind)

  ind = d_ind.copy_to_host()
  A = d_A.copy_to_host()
  B = d_B.copy_to_host()

  return np.allclose(A[ind],B)

print("Test 1:", test1())
print("Test 2:", test2())

[免责声明：此代码是在 Google colab 中用两分钟编写的，运行一次。它需要进一步的代码以将其推广到任何大小并进一步测试以确保其正常工作。不提供任何担保，使用风险自负]

您选择哪个版本可能取决于您应用它的问题的大小。我希望numba_gpu2 将优先用于处理比 CUDA 扭曲大小长很多倍的行的问题，而numba_gpu1 可能更适合处理许多短行的问题。一种聪明的方法是根据数据的大小和形状在它们之间进行启发式选择。所有这些推测的验证和实现留给读者作为练习。

Answer 2

我想我自己已经找到了解决方案。我需要的是在 CUDA 设备中操作 Numpy 数组。为此，CuPy 比 Numba 好得多。 CuPy 以一种高效便捷的方式支持许多类似 Numpy 的操作（包括我的问题中的操作）。