Numba CUDA 缩减为数组

Question

我有两个相当大的数组，长度分别为 N 和 M 个元素。对于 N 个元素中的每一个，我需要对 M 个元素中的每一个进行计算，然后减少这些结果以获得另一个长度为 N 的数组。这听起来像是非常适合 GPU 加速的问题类型，我因此想使用 Numba CUDA 来实现它，但我正在努力寻找如何处理这个问题的减少部分。关于缩减的 Numba 文档 https://numba.pydata.org/numba-doc/dev/cuda/reduction.html 仅显示了如何将所有内容缩减为一个数字，但我本质上需要缩减为一个数组。下面是一个超级简化的例子，基本上是我想要实现的目标

from numba import cuda
import numpy as np

@cuda.jit
def processArr(A, B):
    i = cuda.grid(1)
    if i < A.size:
        A[i] = A[i] * B

@cuda.jit
def reduceArr(A, B, C):
    i = cuda.grid(1)
    if i < A.size:
        total = 1
        processArr(A[i], B[i])

        for j in range(A.shape[1]):
            total *= A[i, j]

        C[i] = total

a = np.array([[0, 0], [1, 1], [1, 2]])
b = np.array([1, 2, 3])
c = np.zeros(3)

threadsperblock = 32
blockspergrid = math.ceil(b.shape[0] / threadsperblock)

reduceArr[blockspergrid, threadsperblock](a, b, c)
print(c)

这段代码显然没有运行但希望它能说明我想要实现的目标。

有没有办法用 Numba 实现这一点，或者首先尝试在 GPU 上执行减少步骤是否愚蠢，即是否最好只在 GPU 上执行 NxM 操作然后减少这些结果在 CPU 之后？

Answer 1

Is there a way to achieve this with Numba

是的。您要做的是分段（或向量化）transform-reduce 操作。您调用的 processArr 实际上是您的转换操作。

要执行多个并行分段归约操作，有多种方法，其中“最佳”方法取决于 N 和 M 的具体大小。

对于 N ~= 100 的选择，我建议每次归约使用一个 CUDA 块。我们将执行 classical parallel reduction but at the CUDA block level, and each block will be responsible for one result element. Therefore each block must process N elements, and we must launch a total number of blocks equal to M. For the block level processing, we will implement a block-stride loop which is conceptually similar to a grid-stride loop.

这是一个例子，大致基于您所展示的内容：

$ cat t47.py
from numba import cuda
import numpy as np

# must be power of 2, less than 1025
nTPB = 128

reduce_init_val = 0

@cuda.jit(device=True)
def reduce_op(x,y):
    return x+y

@cuda.jit(device=True)
def transform_op(x,y):
    return x*y

@cuda.jit
def transform_reduce(A, B, C):
    s = cuda.shared.array(nTPB, dtype=A.dtype)
    b = cuda.blockIdx.x
    t = cuda.threadIdx.x
    if (b < A.shape[0]):
        s[t] = reduce_init_val
        #block-stride loop
        for i in range(t, A.shape[1], nTPB):
            s[t] = reduce_op(s[t], transform_op(A[b,i], B[b]))
        #parallel sweep reduction
        l = nTPB//2
        while (l > 0):
            cuda.syncthreads()
            if  (t < l):
                s[t] = reduce_op(s[t], s[t+l])
            l = l//2
        if (t == 0):
            C[b] = s[0]

a = np.array([[0, 0], [1, 1], [1, 2]], dtype=np.float64)
b = np.array([1, 2, 3], dtype=np.float64)
c = np.zeros(3, dtype=np.float64)

threadsperblock = nTPB
blockspergrid = a.shape[0]
transform_reduce[blockspergrid, threadsperblock](a, b, c)
print(c)
$ python t47.py
[0. 4. 9.]
$

我并不是在暗示以上代码 defect-free 或适合任何特定用途。我的目的是概述一种可能的方法。请注意，上面的代码应该能够处理 more-or 较少的任意 M 和 N 维度。

Numba CUDA 缩减为数组

Numba CUDA reduce into array

python

cuda

gpu

numba