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 块级别执行经典的并行归约，每个块将负责一个结果元素。因此，每个块必须处理N元素，并且我们必须启动的块总数等于M。对于块级处理，我们将实现一个块跨步循环，其概念上类似于网格跨步循环。

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

$ 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维度。