Skip to content

3.12。 GPU 减少

原文： http://numba.pydata.org/numba-doc/latest/cuda/reduction.html

为 CUDA GPU 编写缩减算法可能很棘手。 Numba 提供了一个@reduce装饰器，用于将简单的二进制操作转换为简化内核。

3.12.1。 @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)

3.12.2。减少

reduce装饰器创建Reduce类的实例。 （目前，reduce是Reduce的别名，但不保证这种行为。）

class numba.cuda.Reduce(functor)

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

完全减少。

| 参数： |

  • arr - 主机或设备阵列。如果给出了设备数组，则会在原地执行缩减，并覆盖数组中的值。如果给出了主机阵列，则会自动将其复制到设备。
  • size - 可选整数，指定要减少的arr中的元素数。如果未指定此参数，则会减少整个数组。
  • res - 将减少结果写入的可选设备数组。结果写入此数组的第一个元素。如果指定了此参数，则不会从设备到主机进行减少输出的通信。
  • init - 还原的可选初始值，其类型必须与arr.dtype匹配。
  • 流 - 执行还原的可选 CUDA 流。如果未指定流，则使用默认流 0。

| | --- | --- | | 返回： | 如果指定了res，则返回None。否则，返回减少的结果。 | | --- | --- |

__init__(functor)

创建一个使用给定二进制函数减少值的缩减对象。二进制函数编译一次并缓存在此对象中。保持此对象存活将阻止重新编译。

参数： binop - 要编译为 CUDA 设备函数的函数，该函数将用作 CUDA 设备上的二进制运算。在内部，它使用cuda.jit(device=True)编译。
回到顶部