如何在 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 但未启用”。

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

Answer 1

您说得对，Numba 设备不支持这些切片运算符。潜在的问题是切片操作需要中间数组构造，而 Numba 编译器目前无法做到这一点。

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

使用单线程在源和目标之间复制一行数据（numba_gpu1如下所示）
使用单个块在源和目标之间复制一行数据。这可以利用跨步循环设计模式来改进内存合并和缓存一致性，并且应该在非平凡的大小下表现更好（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 warp 大小长很多倍的行的问题是首选，而 numba_gpu1 可能更适合许多短行的问题。一种聪明的方法是根据数据的大小和形状对它们之间的 select 进行试探。所有这些推测的验证和实施留作 reader.

的练习。

Answer 2

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

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

How to slice rows in numba CUDA?

arrays

cuda

memcpy

numba