用户定义的核函数#
CuPy 提供了定义三种类型 CUDA 核函数的简便方法：逐元素核函数 (elementwise kernels)、归约核函数 (reduction kernels) 和原生核函数 (raw kernels)。本文档中，我们将介绍如何定义和调用每种核函数。
逐元素核函数基础#
逐元素核函数可以通过 ElementwiseKernel 类定义。该类的实例定义了一个 CUDA 核函数，可以通过此实例的 __call__ 方法调用。
逐元素核函数的定义包含四个部分：输入参数列表、输出参数列表、循环体代码和核函数名称。例如，一个计算平方差 \(f(x, y) = (x - y)^2\) 的核函数定义如下：
>>> squared_diff = cp.ElementwiseKernel(
... 'float32 x, float32 y',
... 'float32 z',
... 'z = (x - y) * (x - y)',
... 'squared_diff')
参数列表由逗号分隔的参数定义组成。每个参数定义包含一个 类型说明符 和一个 参数名称。NumPy 数据类型的名称可以用作类型说明符。
注意
n、i 和以下划线 _ 开头的名称保留供内部使用。
上述核函数可以在标量或支持广播的数组上调用：
>>> x = cp.arange(10, dtype=np.float32).reshape(2, 5)
>>> y = cp.arange(5, dtype=np.float32)
>>> squared_diff(x, y)
array([[ 0., 0., 0., 0., 0.],
[25., 25., 25., 25., 25.]], dtype=float32)
>>> squared_diff(x, 5)
array([[25., 16., 9., 4., 1.],
[ 0., 1., 4., 9., 16.]], dtype=float32)
输出参数可以显式指定（紧跟在输入参数之后）：
>>> z = cp.empty((2, 5), dtype=np.float32)
>>> squared_diff(x, y, z)
array([[ 0., 0., 0., 0., 0.],
[25., 25., 25., 25., 25.]], dtype=float32)
类型通用核函数#
如果类型说明符是一个字符，则将其视为 类型占位符。它可用于定义类型通用的核函数。例如，上述 squared_diff 核函数可以通用化类型，如下所示：
>>> squared_diff_generic = cp.ElementwiseKernel(
... 'T x, T y',
... 'T z',
... 'z = (x - y) * (x - y)',
... 'squared_diff_generic')
核函数定义中相同字符的类型占位符表示同一类型。这些占位符的实际类型由实际参数类型确定。ElementwiseKernel 类首先检查输出参数，然后检查输入参数来确定实际类型。如果在核函数调用中未给出输出参数，则仅使用输入参数来确定类型。
类型占位符可以在循环体代码中使用：
>>> squared_diff_generic = cp.ElementwiseKernel(
... 'T x, T y',
... 'T z',
... '''
... T diff = x - y;
... z = diff * diff;
... ''',
... 'squared_diff_generic')
核函数定义中可以使用多个类型占位符。例如，上述核函数可以进一步通用化以处理多个参数：
>>> squared_diff_super_generic = cp.ElementwiseKernel(
... 'X x, Y y',
... 'Z z',
... 'z = (x - y) * (x - y)',
... 'squared_diff_super_generic')
请注意，此核函数需要显式指定输出参数，因为类型 Z 无法从输入参数自动确定。
原生参数说明符 (Raw argument specifiers)#
ElementwiseKernel 类会自动进行广播索引，这对于定义大多数逐元素计算非常有用。另一方面，有时我们希望对某些参数编写带有手动索引的核函数。通过在类型说明符前添加 raw 关键字，可以告知 ElementwiseKernel 类使用手动索引。
我们可以使用特殊变量 i 和方法 _ind.size() 进行手动索引。i 表示循环内的索引。_ind.size() 表示要应用逐元素操作的总元素数。请注意，它表示广播后的大小。
例如，一个反转其中一个向量后相加的核函数可以编写如下：
>>> add_reverse = cp.ElementwiseKernel(
... 'T x, raw T y', 'T z',
... 'z = x + y[_ind.size() - i - 1]',
... 'add_reverse')
(注意，这是一个人工示例，你可以直接使用 z = x + y[::-1] 而无需定义新的核函数)。原生参数可以像数组一样使用。索引操作符 y[_ind.size() - i - 1] 涉及对 y 进行索引计算，因此 y 可以具有任意形状和步长。
请注意，原生参数不参与广播。如果你想将所有参数标记为 raw，则必须在调用时指定 size 参数，该参数定义了 _ind.size() 的值。
纹理内存 (Texture memory)#
纹理对象 (TextureObject) 可以传递给 ElementwiseKernel，其类型使用与同一核函数中使用的任何其他类型不同的唯一类型占位符标记，因为其实际数据类型在填充纹理内存时确定。纹理坐标可以在核函数中通过每线程循环索引 i 计算。
归约核函数 (Reduction kernels)#
归约核函数可以通过 ReductionKernel 类定义。我们可以通过定义核函数的四个部分来使用它：
初始值 (Identity value)：此值用作归约的初始值。
映射表达式 (Mapping expression)：用于对每个要归约的元素进行预处理。
归约表达式 (Reduction expression)：它是用于归约多个映射值的操作符。特殊变量
a和b用作其操作数。后映射表达式 (Post mapping expression)：用于转换归约结果值。特殊变量
a用作其输入。输出应该写入输出参数。
ReductionKernel 类会自动插入高效灵活归约实现所需的其他代码片段。
例如，沿指定轴的 L2 范数可以编写如下：
>>> l2norm_kernel = cp.ReductionKernel(
... 'T x', # input params
... 'T y', # output params
... 'x * x', # map
... 'a + b', # reduce
... 'y = sqrt(a)', # post-reduction map
... '0', # identity value
... 'l2norm' # kernel name
... )
>>> x = cp.arange(10, dtype=np.float32).reshape(2, 5)
>>> l2norm_kernel(x, axis=1)
array([ 5.477226 , 15.9687195], dtype=float32)
注意
raw 说明符限制用于要归约的轴位于形状头部的情况。这意味着，如果你想对至少一个参数使用 raw 说明符，axis 参数必须为 0 或从 0 开始的连续递增整数序列，如 (0, 1)、(0, 1, 2) 等。
注意
ReductionKernel 尚未支持纹理内存。
原生核函数 (Raw kernels)#
原生核函数可以通过 RawKernel 类定义。使用原生核函数，你可以从原始 CUDA 源代码定义核函数。
RawKernel 对象允许你使用 CUDA 的 cuLaunchKernel 接口调用核函数。换句话说，你可以控制网格大小、块大小、共享内存大小和流。
>>> add_kernel = cp.RawKernel(r'''
... extern "C" __global__
... void my_add(const float* x1, const float* x2, float* y) {
... int tid = blockDim.x * blockIdx.x + threadIdx.x;
... y[tid] = x1[tid] + x2[tid];
... }
... ''', 'my_add')
>>> x1 = cp.arange(25, dtype=cp.float32).reshape(5, 5)
>>> x2 = cp.arange(25, dtype=cp.float32).reshape(5, 5)
>>> y = cp.zeros((5, 5), dtype=cp.float32)
>>> add_kernel((5,), (5,), (x1, x2, y)) # grid, block and arguments
>>> y
array([[ 0., 2., 4., 6., 8.],
[10., 12., 14., 16., 18.],
[20., 22., 24., 26., 28.],
[30., 32., 34., 36., 38.],
[40., 42., 44., 46., 48.]], dtype=float32)
也可以创建对复数数组进行操作的原生核函数：
>>> complex_kernel = cp.RawKernel(r'''
... #include <cupy/complex.cuh>
... extern "C" __global__
... void my_func(const complex<float>* x1, const complex<float>* x2,
... complex<float>* y, float a) {
... int tid = blockDim.x * blockIdx.x + threadIdx.x;
... y[tid] = x1[tid] + a * x2[tid];
... }
... ''', 'my_func')
>>> x1 = cupy.arange(25, dtype=cupy.complex64).reshape(5, 5)
>>> x2 = 1j*cupy.arange(25, dtype=cupy.complex64).reshape(5, 5)
>>> y = cupy.zeros((5, 5), dtype=cupy.complex64)
>>> complex_kernel((5,), (5,), (x1, x2, y, cupy.float32(2.0))) # grid, block and arguments
>>> y
array([[ 0. +0.j, 1. +2.j, 2. +4.j, 3. +6.j, 4. +8.j],
[ 5.+10.j, 6.+12.j, 7.+14.j, 8.+16.j, 9.+18.j],
[10.+20.j, 11.+22.j, 12.+24.j, 13.+26.j, 14.+28.j],
[15.+30.j, 16.+32.j, 17.+34.j, 18.+36.j, 19.+38.j],
[20.+40.j, 21.+42.j, 22.+44.j, 23.+46.j, 24.+48.j]],
dtype=complex64)
请注意，虽然我们鼓励对复数使用 complex<T> 类型（通过包含 <cupy/complex.cuh> 获得，如上所示），但对于已经使用 cuComplex.h 中函数编写的 CUDA 代码，无需自己进行转换：只需在创建 RawKernel 实例时设置选项 translate_cucomplex=True。
CUDA 核函数属性可以通过访问 attributes 字典或直接访问 RawKernel 对象的属性来检索；后者也可用于设置某些属性：
>>> add_kernel = cp.RawKernel(r'''
... extern "C" __global__
... void my_add(const float* x1, const float* x2, float* y) {
... int tid = blockDim.x * blockIdx.x + threadIdx.x;
... y[tid] = x1[tid] + x2[tid];
... }
... ''', 'my_add')
>>> add_kernel.attributes
{'max_threads_per_block': 1024, 'shared_size_bytes': 0, 'const_size_bytes': 0, 'local_size_bytes': 0, 'num_regs': 10, 'ptx_version': 70, 'binary_version': 70, 'cache_mode_ca': 0, 'max_dynamic_shared_size_bytes': 49152, 'preferred_shared_memory_carveout': -1}
>>> add_kernel.max_dynamic_shared_size_bytes
49152
>>> add_kernel.max_dynamic_shared_size_bytes = 50000 # set a new value for the attribute
>>> add_kernel.max_dynamic_shared_size_bytes
50000
RawKernel 支持动态并行。你只需要将链接标志（如 -dc）提供给 RawKernel 的 options 参数。CuPy 会自动发现静态 CUDA 设备运行时库 (cudadevrt)。有关详细信息，请参阅 CUDA 工具包文档。
在 RawKernel 中访问纹理（表面）内存通过 CUDA 运行时纹理（表面）对象 API 支持，请参阅 TextureObject (SurfaceObject) 以及 CUDA C 编程指南的文档。对于使用纹理引用 API（自 CUDA 工具包 10.1 起已被标记为已弃用），请参阅下方 RawModule 的介绍。
如果你的核函数依赖于 C++ std 库头文件，例如 <type_traits>，很可能会遇到编译错误。在这种情况下，尝试通过在创建 RawKernel 实例时设置 jitify=True 来启用 CuPy 的 Jitify 支持。它提供基本的 C++ std 支持以弥补常见错误。
注意
核函数没有返回值。你需要将输入数组和输出数组都作为参数传递。
注意
在 CUDA 核函数中使用 printf() 时，可能需要同步流才能看到输出。如果使用默认流，可以使用 cupy.cuda.Stream.null.synchronize()。
注意
在上述所有示例中，我们都将核函数声明在 extern "C" 块中，表示使用 C 链接。这是为了确保核函数名称不被修饰，以便可以通过名称检索它们。
核函数参数#
Python 基本类型和 NumPy 标量以值传递给核函数。数组参数（指针参数）必须作为 CuPy ndarray 传递。CuPy 对传递给核函数的参数不执行验证，包括类型和参数数量。
特别需要注意的是，传递 CuPy ndarray 时，其 dtype 应与 CUDA 源代码函数签名中声明的参数类型匹配（除非你有意进行数组类型转换）。
例如，cupy.float32 和 cupy.uint64 数组必须分别传递给类型为 float* 和 unsigned long long* 的参数。CuPy 不直接支持非基本类型（如 float3）的数组，但在核函数中将 float* 或 void* 转换为 float3* 没有任何问题。
Python 基本类型 int, float, complex 和 bool 分别映射到 long long, double, cuDoubleComplex 和 bool。
NumPy 标量 (numpy.generic) 和大小为一的 NumPy 数组 (numpy.ndarray) 以值传递给核函数。这意味着你可以按值传递任何基础 NumPy 类型，例如 numpy.int8 或 numpy.float64，前提是核函数参数的大小匹配。你可以参考下表匹配 CuPy/NumPy dtype 和 CUDA 类型：
CuPy/NumPy 类型 |
对应的核函数类型 |
itemsize (字节) |
|---|---|---|
bool |
bool |
1 |
int8 |
char, signed char |
1 |
int16 |
short, signed short |
2 |
int32 |
int, signed int |
4 |
int64 |
long long, signed long long |
8 |
uint8 |
unsigned char |
1 |
uint16 |
unsigned short |
2 |
uint32 |
unsigned int |
4 |
uint64 |
unsigned long long |
8 |
float16 |
half |
2 |
float32 |
float |
4 |
float64 |
double |
8 |
complex64 |
float2, cuFloatComplex, complex<float> |
8 |
complex128 |
double2, cuDoubleComplex, complex<double> |
16 |
CUDA 标准保证了基本类型在主机和设备上的大小始终匹配。size_t、ptrdiff_t、intptr_t、uintptr_t、long、signed long 和 unsigned long 的 itemsize 取决于平台。要将任何 CUDA 向量内置类型（如 float3）或任何其他用户定义结构作为核函数参数传递（前提是它与设备端核函数参数类型匹配），请参见下面的自定义用户类型。
自定义用户类型#
可以通过定义自定义 NumPy dtype 来使用自定义类型（复合类型，例如结构体和结构体的结构体）作为核函数参数。执行此操作时，你有责任匹配主机和设备结构体的内存布局。CUDA 标准保证了基本类型在主机和设备上的大小始终匹配。然而，它可能对复合类型施加设备对齐要求。这意味着对于复合类型，结构体成员偏移量可能与你预期的不同。
当按值传递核函数参数时，CUDA 驱动程序将从 NumPy 对象数据指针的起始位置精确复制 sizeof(param_type) 字节，其中 param_type 是核函数中的参数类型。你必须通过定义相应的 NumPy dtype 来匹配 param_type 的内存布局（例如：大小、对齐和结构体填充/打包）。
对于内置 CUDA 向量类型（如 int2 和 double4）以及其他具有命名成员的打包结构体，你可以直接定义此类 NumPy dtypes，如下所示：
>>> import numpy as np
>>> names = ['x', 'y', 'z']
>>> types = [np.float32]*3
>>> float3 = np.dtype({'names': names, 'formats': types})
>>> arg = np.random.rand(3).astype(np.float32).view(float3)
>>> print(arg)
[(0.9940819, 0.62873816, 0.8953669)]
>>> arg['x'] = 42.0
>>> print(arg)
[(42., 0.62873816, 0.8953669)]
此处 arg 可以直接用作核函数参数。如果不需要命名字段，你可以优先使用此语法来定义打包结构体，例如向量或矩阵：
>>> import numpy as np
>>> float5x5 = np.dtype({'names': ['dummy'], 'formats': [(np.float32,(5,5))]})
>>> arg = np.random.rand(25).astype(np.float32).view(float5x5)
>>> print(arg.itemsize)
100
此处 arg 表示一个 100 字节的标量（即大小为 1 的 NumPy 数组），可以按值传递给任何核函数。核函数参数通过专用的 4KB 内存库（带有自己的广播缓存）按值传递。因此，总核函数参数大小的上限为 4KB（参见此链接）。需要注意的是，这个专用内存库不与设备 __constant__ 内存空间共享。
目前，CuPy 不提供用于创建用户定义复合类型的辅助例程。但是，可以使用 NumPy dtype 的 offsets 和 itemsize 功能递归构建此类复合类型，请参阅 cupy/examples/custum_struct 获取高级用法示例。
警告
你不能使用 type arg[N] 语法直接将静态数组作为核函数参数传递，其中 N 是编译时常量。__global__ void kernel(float arg[5]) 的签名被编译器视为 __global__ void kernel(float* arg)。如果你想将五个浮点数按值传递给核函数，你需要定义一个自定义结构体 struct float5 { float val[5]; }; 并修改核函数签名以使其变为 __global__ void kernel(float5 arg)。
原生模块 (Raw modules)#
对于处理大型原始 CUDA 源代码或加载现有 CUDA 二进制文件，RawModule 类可能更方便。它可以通过 CUDA 源代码或 CUDA 二进制文件的路径进行初始化。它接受与 RawKernel 中相同的大部分参数。然后可以通过调用 get_function() 方法检索所需的核函数，该方法返回一个 RawKernel 实例，可以按照上述方式调用。
>>> loaded_from_source = r'''
... extern "C"{
...
... __global__ void test_sum(const float* x1, const float* x2, float* y, \
... unsigned int N)
... {
... unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
... if (tid < N)
... {
... y[tid] = x1[tid] + x2[tid];
... }
... }
...
... __global__ void test_multiply(const float* x1, const float* x2, float* y, \
... unsigned int N)
... {
... unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
... if (tid < N)
... {
... y[tid] = x1[tid] * x2[tid];
... }
... }
...
... }'''
>>> module = cp.RawModule(code=loaded_from_source)
>>> ker_sum = module.get_function('test_sum')
>>> ker_times = module.get_function('test_multiply')
>>> N = 10
>>> x1 = cp.arange(N**2, dtype=cp.float32).reshape(N, N)
>>> x2 = cp.ones((N, N), dtype=cp.float32)
>>> y = cp.zeros((N, N), dtype=cp.float32)
>>> ker_sum((N,), (N,), (x1, x2, y, N**2)) # y = x1 + x2
>>> assert cp.allclose(y, x1 + x2)
>>> ker_times((N,), (N,), (x1, x2, y, N**2)) # y = x1 * x2
>>> assert cp.allclose(y, x1 * x2)
上面关于在 RawKernel 中使用复数的说明也适用于 RawModule。
对于需要访问全局符号（例如常量内存）的 CUDA 核函数，可以使用 get_global() 方法，详细信息请参阅其文档。
请注意，自 CuPy vX.X 起，已弃用的 API cupy.RawModule.get_texref() 由于 CUDA 移除了纹理引用支持而被移除。
为了支持 C++ 模板核函数，RawModule 额外提供了一个 name_expressions 参数。应该提供一个模板特化列表，以便可以按类型生成和检索相应的核函数：
>>> code = r'''
... template<typename T>
... __global__ void fx3(T* arr, int N) {
... unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
... if (tid < N) {
... arr[tid] = arr[tid] * 3;
... }
... }
... '''
>>>
>>> name_exp = ['fx3<float>', 'fx3<double>']
>>> mod = cp.RawModule(code=code, options=('-std=c++11',),
... name_expressions=name_exp)
>>> ker_float = mod.get_function(name_exp[0]) # compilation happens here
>>> N=10
>>> a = cp.arange(N, dtype=cp.float32)
>>> ker_float((1,), (N,), (a, N))
>>> a
array([ 0., 3., 6., 9., 12., 15., 18., 21., 24., 27.], dtype=float32)
>>> ker_double = mod.get_function(name_exp[1])
>>> a = cp.arange(N, dtype=cp.float64)
>>> ker_double((1,), (N,), (a, N))
>>> a
array([ 0., 3., 6., 9., 12., 15., 18., 21., 24., 27.])
注意
用于初始化 RawModule 实例和检索核函数的名字表达式是原始的（未修饰的）核函数名称，并明确指定了所有模板参数。名称修饰和解修饰在内部处理，因此用户无需担心。
核函数融合 (Kernel fusion)#
cupy.fuse() 是一个用于融合函数的装饰器。使用此装饰器可以比 ElementwiseKernel 或 ReductionKernel 更轻松地定义逐元素核函数或归约核函数。
使用此装饰器，我们可以按如下方式定义 squared_diff 核函数：
>>> @cp.fuse()
... def squared_diff(x, y):
... return (x - y) * (x - y)
上述核函数可以在标量、NumPy 数组或 CuPy 数组上调用，就像原始函数一样。
>>> x_cp = cp.arange(10)
>>> y_cp = cp.arange(10)[::-1]
>>> squared_diff(x_cp, y_cp)
array([81, 49, 25, 9, 1, 1, 9, 25, 49, 81])
>>> x_np = np.arange(10)
>>> y_np = np.arange(10)[::-1]
>>> squared_diff(x_np, y_np)
array([81, 49, 25, 9, 1, 1, 9, 25, 49, 81])
在第一次函数调用时，融合函数根据参数的抽象信息（例如它们的 dtype 和 ndim）分析原始函数，并创建和缓存实际的 CUDA 核函数。从第二次使用相同输入类型的函数调用开始，融合函数调用之前缓存的核函数，因此强烈建议重用相同的被装饰函数，而不是多次定义局部函数并对其进行装饰。
cupy.fuse() 也支持简单的归约核函数。
>>> @cp.fuse()
... def sum_of_products(x, y):
... return cp.sum(x * y, axis = -1)
你可以通过使用 kernel_name 关键字参数来指定核函数名称，如下所示：
>>> @cp.fuse(kernel_name='squared_diff')
... def squared_diff(x, y):
... return (x - y) * (x - y)
注意
目前，cupy.fuse() 只能融合简单的逐元素和归约操作。大多数其他例程（例如 cupy.matmul()、cupy.reshape()）尚不支持。
JIT 核函数定义#
cupyx.jit.rawkernel 装饰器可以从 Python 函数创建原生 CUDA 核函数。
本节中，用此装饰器包装的 Python 函数称为 目标函数。
目标函数由基本的标量操作组成，用户必须自行管理如何并行化它们。CuPy 的自动并行化数组操作（例如 add()、sum()）不受支持。如果需要基于此类数组函数的自定义核函数，请参阅核函数融合一节。
基本用法#
以下是一个简短示例，展示如何编写 cupyx.jit.rawkernel 以使用网格跨步循环将 x 的值复制到 y：
>>> from cupyx import jit
>>>
>>> @jit.rawkernel()
... def elementwise_copy(x, y, size):
... tid = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x
... ntid = jit.gridDim.x * jit.blockDim.x
... for i in range(tid, size, ntid):
... y[i] = x[i]
>>> size = cupy.uint32(2 ** 22)
>>> x = cupy.random.normal(size=(size,), dtype=cupy.float32)
>>> y = cupy.empty((size,), dtype=cupy.float32)
>>> elementwise_copy((128,), (1024,), (x, y, size)) # RawKernel style
>>> assert (x == y).all()
>>> elementwise_copy[128, 1024](x, y, size) # Numba style
>>> assert (x == y).all()
如上所示，支持两种启动核函数的方式。前两个条目分别是网格大小和块大小。grid (RawKernel 样式 (128,) 或 Numba 样式 [128]) 是网格的大小，即每个维度中的块数量；block ((1024,) 或 [1024]) 是每个线程块的维度，详细信息请参阅 cupyx.jit._interface._JitRawKernel。在 GPU 上以预定的网格/块大小启动 CUDA 核函数需要对 CUDA 编程模型有基本了解。
编译将在第一次函数调用时延迟进行。CuPy 的 JIT 编译器在调用时推断参数的类型，并缓存编译好的核函数以加速后续调用。
有关完整的 API 列表，请参阅自定义核函数。
基本设计#
CuPy 的 JIT 编译器通过 Python AST 生成 CUDA 代码。我们决定不使用 Python 字节码来分析目标函数，以避免性能下降。从 Python 字节码生成的 CUDA 源代码不会被 CUDA 编译器有效优化，因为将目标函数转换为字节码时，目标函数的 for 循环和其他控制语句会完全转换为跳转指令。
类型规则#
局部变量的类型在函数中的第一次赋值时推断。第一次赋值必须在函数的顶层进行；换句话说，它不得在 if/else 主体或 for 循环中。
限制#
JIT 在 Python 的交互式解释器 (REPL) 中不起作用，因为编译器需要获取目标函数的源代码。