Numba CUDA缩减为数组

问题描述

我分别有两个长度分别为N和M的相当大的数组。对于N个元素中的每个元素，我需要对M个元素中的每个元素进行计算，然后减少这些结果，以得出长度为N的另一个数组。这听起来像是非常适合GPU加速和I因此，我想使用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],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操作，然后在GPU上减少这些结果之后是cpu？

解决方法

有没有办法通过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],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维度。

cuda cuda gpu gpu numba python