CUDA功能和通用功能

CUDA功能和通用功能

本文描述了类似于CUDA ufunc的对象。

为了支持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功能

类似于CUDA ufunc功能，可以使用CUDA在GPU上执行通用ufunc。这可以通过以下方式完成：

from numba import guvectorize

 

@guvectorize(['void(float32[:,:], float32[:,:], float32[:,:])'],

             '(m,n),(n,p)->(m,p)', target='cuda')

def matmulcore(A, B, C):

    ...

有时gufunc内核会使用过多的GPU资源，这可能会导致内核启动失败。用户可以通过在已编译的gufunc对象上设置max_blocksize属性来显式控制线程块的最大size。

from numba import guvectorize

 

@guvectorize(..., target='cuda')

def very_complex_kernel(A, B, C):

    ...

 

very_complex_kernel.max_blocksize = 32  # limits to 32 threads per block