用户定义的内核#

CuPy 提供了定义三种 CUDA 内核的简便方法：元素内核、归约内核和原始内核。在本文档中，我们描述了如何定义和调用每种内核。

元素级内核的基础#

逐元素内核可以通过 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 无法从输入参数自动确定。

原始参数说明符#

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() 的值。

纹理内存#

纹理对象（TextureObject）可以传递给 ElementwiseKernel，其类型由一个独特的类型占位符标记，该占位符与同一内核中使用的任何其他类型不同，因为其实际数据类型是在填充纹理内存时确定的。纹理坐标可以通过每个线程的循环索引 i 在内核中计算。

简化内核#

可以通过 ReductionKernel 类定义缩减内核。我们可以通过定义内核代码的四个部分来使用它：

  1. 标识值：此值用于归约的初始值。

  2. 映射表达式：用于对每个要减少的元素进行预处理。

  3. 归约表达式：它是一个用于归约多个映射值的运算符。特殊变量 a 和 b 用于其操作数。

  4. 后映射表达式：用于转换得到的归约值。特殊变量 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 中尚未支持。

原始内核#

原始内核可以通过 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。

可以通过访问 attributes 字典或直接访问 RawKernel 对象的属性来获取CUDA内核属性；后者也可以用于设置某些属性：

>>> 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 支持。您只需向 RawKernel 的 options 参数提供链接标志（如 -dc）。静态 CUDA 设备运行时库（cudadevrt）由 CuPy 自动发现。更多详情，请参阅 CUDA Toolkit 的文档。

在 RawKernel 中访问纹理（表面）内存是通过 CUDA Runtime 的纹理（表面）对象 API 支持的，请参阅 TextureObject (SurfaceObject) 的文档以及 CUDA C 编程指南。对于使用纹理引用 API，自 CUDA Toolkit 10.1 起已被标记为弃用，请参阅下面的 RawModule 介绍。

如果你的内核依赖于 C++ 标准库头文件，如 <type_traits>，你很可能会遇到编译错误。在这种情况下，尝试在创建 RawKernel 实例时通过设置 jitify=True 来启用 CuPy 的 Jitify 支持。它提供了基本的 C++ 标准库支持，以修复常见的错误。

备注

内核没有返回值。你需要将输入数组和输出数组都作为参数传递。

备注

在你的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 数据类型和 CUDA 类型：

CuPy/NumPy 类型

对应的内核类型

itemsize (字节)

布尔

布尔

toctree 是一个 reStructuredText 指令 ，这是一个非常多功能的标记。指令可以有参数、选项和内容。

int8

char, signed char

toctree 是一个 reStructuredText 指令 ，这是一个非常多功能的标记。指令可以有参数、选项和内容。

int16

短整型，有符号短整型

2

int32

int, 有符号 int

4

int64

长长，有符号长长

8

uint8

无符号字符

toctree 是一个 reStructuredText 指令 ，这是一个非常多功能的标记。指令可以有参数、选项和内容。

uint16

无符号短整型

2

uint32

无符号整数

4

uint64

无符号长长整型

8

float16

一半

2

float32

浮动

4

float64

双

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数据类型，如下所示：

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

原始模块#

对于处理大量的原始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() 方法，详见其文档。

请注意，由于 CUDA 移除了对纹理引用支持，自 CuPy vX.X 起已移除已弃用的 API cupy.RawModule.get_texref()。

为了支持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 实例和检索内核的名称表达式是原始的（未混淆的）内核名称，其中所有模板参数都明确指定。名称的混淆和去混淆在幕后处理，因此用户无需担心。

内核融合#

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])

在第一次函数调用时，融合函数会根据参数的抽象信息（例如它们的dtypes和ndims）分析原始函数，并创建并缓存一个实际的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）中不起作用，因为编译器需要获取目标函数的源代码。