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

有没有办法用 Numba 实现这一点

是的。您想要做的是分段（或矢量化）变换归约操作。您调用的 processArr 实际上是您的转换操作。

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

对于您选择的N ~= 100，我建议每次减少使用一个 CUDA 块。我们将在 CUDA 块级别执行classical parallel reduction，每个块将负责一个结果元素。因此每个块必须处理N 元素，并且我们必须启动等于M 的块总数。对于块级处理，我们将实现一个块步长循环，它在概念上类似于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.]
$

我并不是说上述代码没有缺陷或适用于任何特定目的。我的意图是概述一种可能的方法。请注意，上面的代码应该能够处理或多或少的任意M 和N 维度。