GPU 归约

CUDA 内置目标弃用通知

Numba 内置的 CUDA 目标已弃用，后续开发已移至 NVIDIA numba-cuda 包。请参阅 内置 CUDA 目标弃用和维护状态。

为 CUDA GPU 编写归约算法可能很棘手。Numba 提供了一个 @reduce 装饰器，用于将简单的二元操作转换为归约核函数。示例如下

import numpy
from numba import cuda

@cuda.reduce
def sum_reduce(a, b):
    return a + b

A = (numpy.arange(1234, dtype=numpy.float64)) + 1
expect = A.sum()      # NumPy sum reduction
got = sum_reduce(A)   # cuda sum reduction
assert expect == got

此处也可以使用 Lambda 函数

sum_reduce = cuda.reduce(lambda a, b: a + b)

Reduce 类

reduce 装饰器会创建一个 Reduce 类的实例。目前，reduce 是 Reduce 的别名，但此行为不保证。

class numba.cuda.Reduce(functor)

创建一个归约对象，该对象使用给定的二元函数对值进行归约。二元函数编译一次并在此对象内部缓存。保持此对象存活将防止重新编译。

__init__(functor)

参数

functor – 一个实现归约二元操作的函数。它将使用 cuda.jit(device=True) 编译为 CUDA 设备函数。

__call__(arr, size=None, res=None, init=0, stream=0)

执行完全归约。

参数

• arr – 主机或设备数组。

• size – 可选整数，指定 arr 中要归约的元素数量。如果未指定此参数，则归约整个数组。

• res – 可选的设备数组，用于写入归约结果。结果将写入此数组的第一个元素。如果指定此参数，则归约输出不会从设备传输到主机。

• init – 可选的归约初始值，其类型必须与 arr.dtype 匹配。

• stream – 可选的 CUDA 流，用于执行归约。如果未指定流，则使用默认流 0。

返回

如果指定 res，则返回 None。否则，返回归约结果。