CUDA Ufuncs 和广义 Ufuncs

本页描述了 CUDA ufunc-like 对象。

为了支持CUDA程序的编程模式，CUDA Vectorize和GUVectorize不能生成常规的ufunc。相反，会返回一个类似ufunc的对象。这个对象是一个接近的模拟，但与常规的NumPy ufunc不完全兼容。CUDA ufunc增加了对传递设备内数组（已经在GPU设备上）的支持，以减少通过PCI-express总线的流量。它还接受一个`stream`关键字，用于在异步模式下启动。

示例：基本示例

import math
from numba import vectorize, cuda
import numpy as np

@vectorize(['float32(float32, float32, float32)',
            'float64(float64, float64, float64)'],
           target='cuda')
def cu_discriminant(a, b, c):
    return math.sqrt(b ** 2 - 4 * a * c)

N = 10000
dtype = np.float32

# prepare the input
A = np.array(np.random.sample(N), dtype=dtype)
B = np.array(np.random.sample(N) + 10, dtype=dtype)
C = np.array(np.random.sample(N), dtype=dtype)

D = cu_discriminant(A, B, C)

print(D)  # print result

示例：调用设备功能

所有CUDA ufunc内核都有能力调用其他CUDA设备函数:

from numba import vectorize, cuda

# define a device function
@cuda.jit('float32(float32, float32, float32)', device=True, inline=True)
def cu_device_fn(x, y, z):
    return x ** y / z

# define a ufunc that calls our device function
@vectorize(['float32(float32, float32, float32)'], target='cuda')
def cu_ufunc(x, y, z):
    return cu_device_fn(x, y, z)

广义 CUDA ufuncs

广义ufuncs可以通过CUDA在GPU上执行，类似于CUDA ufunc功能。这可以通过以下方式实现:

from numba import guvectorize

@guvectorize(['void(float32[:,:], float32[:,:], float32[:,:])'],
             '(m,n),(n,p)->(m,p)', target='cuda')
def matmulcore(A, B, C):
    ...