企业官网建设流程全解析

CuPy 自定义结构体内核实战：基于 RawKernel 与 RawModule 的三种递进式写法

【免费下载链接】cupyNumPy & SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy

本指南以examples/custom_struct/目录下的三个示例脚本为核心，系统讲解如何在 CuPy 的RawKernel与RawModule中传递自定义结构体：从 CUDA 内置向量float4/double3的标量与数组传参，到模板化打包结构体（packed struct）的实例化，再到递归构造与设备端内存布局严格一致的 NumPy dtype。读完本文，你将掌握在 CuPy 中自定义任意设备结构体、并在 Host 侧精确复刻其内存布局的完整方案。

三个示例按复杂度递增排列，均可作为普通 Python 脚本直接运行（python3.x example_name.py），对应的官方说明文档位于 docs/source/reference/kernel.rst。

一、示例总览：从最简单到最复杂的三个台阶

examples/custom_struct/目录为自定义结构体在 RawKernel 中的使用提供了三个渐进式示例：

三者均以「Host 侧构造数据结构 → 传入内核 → 设备端计算 → 与 NumPy/CuPy 参考结果比对」的闭环演示验证正确性。下面逐一深入。

二、第一步：CUDA 内置向量double3的两种传参方式

builtin_vectors.py 演示了 CUDA 内置向量类型在 RawKernel 中最基础也最常用的两种用法：数组参数（指针）与标量参数（按值传递）。

2.1 内核代码：自定义运算符的 CUDA 源码

脚本先定义了double3的加法运算符重载与内核sum_kernel：

__device__ double3 operator+(const double3& lhs, const double3& rhs) { return make_double3(lhs.x + rhs.x, lhs.y + rhs.y, lhs.z + rhs.z); } extern "C" __global__ void sum_kernel(const double3* lhs, double3 rhs, double3* out) { int i = threadIdx.x; out[i] = lhs[i] + rhs; }

注意两个细节：

• lhs与out是const double3*指针，对应设备端数组；rhs是double3值类型，对应从 Host 按值传入的标量参数。
• 内核声明了extern "C"以避免 C++ 名字修饰（name mangling），这是 RawKernel 定位内核函数名的前提。

2.2 Host 侧：用 NumPy dtype 描述double3

设备端是struct { double x, y, z; }，Host 侧用带字段的 NumPy dtype 精确复刻：

double3 = numpy.dtype( { 'names': ['x', 'y', 'z'], 'formats': [numpy.float64]*3 } )

• names字段名与设备结构体成员一一对应；
• formats字段类型numpy.float64与 CUDA 的double一致；
• 三个连续double无对齐间隙，因此不需要显式指定offsets与itemsize。

2.3 组装参数并启动内核

lhs = cupy.random.rand(3*N, dtype=numpy.float64).reshape(N, 3) rhs = numpy.random.rand(3).astype(numpy.float64) out = cupy.empty_like(lhs) kernel = cupy.RawKernel(code, 'sum_kernel') args = (lhs, rhs.view(double3), out) kernel((1,), (N,), args)

关键点在于rhs.view(double3)：把形状为(3,)的普通float64数组重新解释（reinterpret）为形状为(1,)、类型为double3的数组，从而按值传给内核的double3 rhs形参。lhs与out由于形状为(N, 3)，直接以数组指针身份传入。

最终验证：

expected = lhs + cupy.asarray(rhs[None, :]) cupy.testing.assert_array_equal(expected, out)

2.4 底层依据：RawKernel 是什么

从 raw.pyx 的类定义看，RawKernel接收用户提供的原始 CUDA 源码字符串并在首次调用时编译，编译结果按设备缓存，同时落盘缓存到$HOME/.cupy/kernel_cache/目录，可供其他进程复用。其核心构造参数包括：

• code：CUDA 源码字符串；
• name：内核函数名；
• options：传给后端（NVRTC 或 NVCC）的编译选项；
• backend：'nvrtc'（默认）或'nvcc'；
• translate_cucomplex、enable_cooperative_groups、jitify等可选特性开关。

从 CuPy v14.0.0 起，RawKernel 默认以 C++17 标准（-std=c++17）编译。另外需要特别留意的是：与ElementwiseKernel不同，RawKernel不会自动处理视图（view），传入matrix.T这类转置视图会被当作原始矩阵处理，跨步（stride）需要自行管理。

三、第二步：RawModule 与模板化打包结构体

packed_matrix.py 将复杂度提升一个台阶：结构体本身是 C++ 模板，内核也以模板函数形式定义，需要借助RawModule与name_expressions完成显式实例化。

3.1 模板化结构体与运算符

设备端定义了一个 4×4 矩阵模板结构体，并重载了+与*运算符：

template<typename T> struct Matrix { T value[4][4]; __device__ T& operator() (int i, int j) { return this->value[i][j]; } __device__ const T& operator() (int i, int j) const { return this->value[i][j]; } }; template<typename T> __device__ Matrix<T> operator+ (const Matrix<T>& lhs, const Matrix<T>& rhs) { Matrix<T> res; for (int i = 0; i<4; i++) { for (int j = 0; j<4; j++) { res(i,j) = lhs(i,j) + rhs(i,j); } } return res; } template<typename T> __device__ Matrix<T> operator* (const Matrix<T>& lhs, const Matrix<T>& rhs) { Matrix<T> res; for (int i = 0; i<4; i++) { for (int j = 0; j<4; j++) { res(i,j) = T(0); for (int k = 0; k<4; k++) { res(i,j) += lhs(i,k) * rhs(k,j); } } } return res; } template<typename T> __global__ void kernel(const Matrix<T>* A, const Matrix<T>* B, const Matrix<T> C, Matrix<T>* out) { int i = threadIdx.x; out[i] = A[i] * B[i] + C; }

内核语义为out = A*B+C，其中A、B是Matrix<T>数组，C是Matrix<T>标量参数。

3.2 用 RawModule 实例化模板内核

模板内核不能直接交给RawKernel（因为它不知道要实例化成哪个具体类型），必须交给RawModule并用name_expressions声明需要的特化版本：

module = cupy.RawModule(code=code, name_expressions=( 'kernel<float>', 'kernel<double>'))

随后在循环中对float（对应numpy.float32）与double（对应numpy.float64）两个类型分别构造参数、取函数、启动内核：

for (ctype, dtype) in zip(('float', 'double'), (numpy.float32, numpy.float64)): A = cupy.random.rand(16*N, dtype=dtype).reshape(N, 4, 4) B = cupy.random.rand(16*N, dtype=dtype).reshape(N, 4, 4) C = numpy.random.rand(16).astype(dtype).reshape(4, 4) out = cupy.empty_like(A) Matrix = numpy.dtype( { 'names': ['value'], 'formats': [(dtype, (4, 4))] } ) kernel = module.get_function('kernel<{}>'.format(ctype)) args = (A, B, C.ravel().view(Matrix), out) kernel((1,), (N,), args) expected = cupy.matmul(A, B) + cupy.asarray(C[None, :, :]) cupy.testing.assert_array_almost_equal(expected, out)

Host 侧用嵌套 dtype 描述设备端结构体：Matrix只有一个成员value，类型为(dtype, (4, 4))，即 4×4 的标量数组——与设备端T value[4][4]的内存排布一致。C.ravel().view(Matrix)再次运用「重新解释视图」技巧，把扁平数组包装成结构体数组后按值传递。

3.3 底层依据：name_expressions 与 get_function 的原理

从 raw.pyx 的文档可知，name_expressions是「指向 C++ 全局/模板内核名称的字符串序列」，例如('func1<int>', 'func1<double>', 'func2')，其中模板内核必须通过get_function逐一取出对应特化版本。

RawModule的约束条件（见 raw.pyx）：

• code与path（指向预编译.cubin/.ptx的路径）互斥，必须且只能提供一个；
• 指定name_expressions时必须有源码且后端必须是nvrtc——因为 C++ 名字修饰（mangled name）解析依赖 NVRTC 的nvrtcAddNameExpression/nvrtcGetLoweredName机制；
• get_function内部先通过module.mapping把用户提供的 C++ 表达式映射为修饰后的内核名，再包装成RawKernel返回（见 raw.pyx）。

从 CuPy v8.0.0 起，RawModule 的编译延迟到首次取对象（内核或指针）时进行，而非初始化时。对于模板特化的行为，仓库测试 tests/cupy_tests/core_tests/test_raw.py 中的test_template_specialization给出了可复现的验证路径。

四、第三步：递归构造与设备布局完全一致的复杂 dtype

complex_struct.py 是三个示例中信息密度最高、最贴近真实工程需求的：目标结构体包含多种异构成员（向量、标量、数组），存在对齐与填充，无法靠「恰好连续」的假设蒙混过关，必须显式查询设备端布局。

4.1 目标结构体

struct complex_struct { int4 a; // 16 字节对齐体 char b; // 1 字节 double c[2]; // 2 个 double short1 d; // 1 个 short unsigned long long int e[3]; // 3 个 64 位整数 };

混合了int4（含x/y/z/w四个成员）、char、double[2]、short1、unsigned long long[3]，由于对齐规则会引入填充字节，Host 侧无法凭直觉手写 layout。

4.2 第一步：用探针内核查询布局

脚本先编译一个「探针内核」get_struct_layout，直接在设备端用sizeof与成员地址差求出结构体的整体大小（itemsize）、各成员大小（sizes）与各成员偏移（offsets）：

extern "C" __global__ void get_struct_layout( unsigned long long *itemsize, unsigned long long *sizes, unsigned long long *offsets) { const complex_struct* ptr = nullptr; itemsize[0] = sizeof(complex_struct); sizes[0] = sizeof(ptr->a); sizes[1] = sizeof(ptr->b); sizes[2] = sizeof(ptr->c); sizes[3] = sizeof(ptr->d); sizes[4] = sizeof(ptr->e); offsets[0] = (unsigned long long)&ptr->a; offsets[1] = (unsigned long long)&ptr->b; offsets[2] = (unsigned long long)&ptr->c; offsets[3] = (unsigned long long)&ptr->d; offsets[4] = (unsigned long long)&ptr->e; }

Host 侧分配三个uint64数组并启动内核，随后拷回主机并打印：

itemsize = cupy.ndarray(shape=(1,), dtype=numpy.uint64) sizes = cupy.ndarray(shape=(5,), dtype=numpy.uint64) offsets = cupy.ndarray(shape=(5,), dtype=numpy.uint64) kernel = cupy.RawKernel(struct_layout_code, 'get_struct_layout',) kernel((1,), (1,), (itemsize, sizes, offsets)) (itemsize, sizes, offsets) = map(cupy.asnumpy, (itemsize, sizes, offsets)) print("Overall structure itemsize: {} bytes".format(itemsize.item())) print("Structure members itemsize: {}".format(sizes)) print("Structure members offsets: {}".format(offsets))

cupy.asnumpy将设备数组转回 NumPy 数组，输出的字节级信息即为构造 Host dtype 的全部依据。这一「设备端自报布局」的思路，避免了手工对齐计算在异构平台上的脆弱性。

4.3 第二步：用 make_packed 构造各成员的 dtype

def make_packed(basetype, N, itemsize): # A small utility function to make packed structs # Can represent simple packed vectors such as float4 or double[3]. assert 0 < N <= 4, N names = list('xyzw')[:N] formats = [basetype]*N return numpy.dtype(dict(names=names, formats=formats, itemsize=itemsize))

该工具函数用names/formats/itemsize三元组构造一个「打包」的 dtype：字段名取xyzw的前 N 个，字段类型统一为basetype，并显式指定itemsize（取自设备端查到的各成员实际大小），从而把int4、short1这类内置向量与double[2]这类数组统一表示为可用的 dtype。

4.4 第三步：递归组合出完整结构体 dtype

有了每个成员的类型与设备端实测偏移，就能组装出complex_struct的完整 dtype：

atype = make_packed(numpy.int32, 4, sizes[0]) btype = make_packed(numpy.int8, 1, sizes[1]) ctype = make_packed(numpy.float64, 2, sizes[2]) dtype = make_packed(numpy.int16, 1, sizes[3]) etype = make_packed(numpy.uint64, 3, sizes[4]) names = list('abcde') formats = [atype, btype, ctype, dtype, etype] complex_struct = numpy.dtype(dict(names=names, formats=formats, offsets=offsets, itemsize=itemsize.item()))

注意offsets与itemsize均来自探针内核的实测值，因此最终 dtype 的内存布局与设备端complex_struct严格一致。这正是示例标题中「递归构建」的含义：先构造成员 dtype，再以成员 dtype 为「格式」构造外层结构体 dtype。

4.5 填充数据、启动内核并验证

s = numpy.empty(shape=(1,), dtype=complex_struct) s['a'] = numpy.arange(0, 4).astype(numpy.int32).view(atype) s['b'] = numpy.arange(4, 5).astype(numpy.int8).view(btype) s['c'] = numpy.arange(5, 7).astype(numpy.float64).view(ctype) s['d'] = numpy.arange(7, 8).astype(numpy.int16).view(dtype) s['e'] = numpy.arange(8, 11).astype(numpy.uint64).view(etype)

每个成员都用view(...)包装成对应的打包 dtype 后写入。设备端内核test_kernel按值接收整个结构体并求和：

extern "C" __global__ void test_kernel(const complex_struct s, double* out) { int i = threadIdx.x; double sum = 0.0; sum += s.a.x + s.a.y + s.a.z + s.a.w; sum += s.b; sum += s.c[0] + s.c[1]; sum += s.d.x; sum += s.e[0] + s.e[1] + s.e[2]; out[i] = i * sum; }

填充值0..10之和恰为 55.0，因此验证断言为：

expected = cupy.arange(N) * 55.0 cupy.testing.assert_array_almost_equal(expected, out)

4.6 适用范围与边界

make_packed中的assert 0 < N <= 4表明该工具适用于简单打包向量（如float4、double[3]）。若设备结构体包含更复杂的嵌套或对齐规则，核心方法论不变：始终以探针内核实测的sizeof/偏移为准，再用names/formats/offsets/itemsize完整刻画 dtype。

五、运行方式与工程建议

三个脚本均为自包含的可执行文件，无额外依赖（仅需numpy与已安装的cupy），直接运行即可：

python3 builtin_vectors.py python3 packed_matrix.py python3 complex_struct.py

每个脚本成功时会打印形如Kernel output matches expected value.的确认信息，complex_struct.py还会打印设备端查到的布局细节（结构体总大小、各成员大小与偏移）。

综合三个示例，在 CuPy 中传递自定义结构体的工程建议可归纳为：

1. 优先考虑内置向量：若数据可用float4、double3等 CUDA 内置向量表达，直接配合「names+formats」的简单 dtype 即可，无需处理对齐；
2. 模板化结构体交给 RawModule：模板内核必须用name_expressions显式声明特化名单，再经get_function('kernel<T>')取出，底层依赖 NVRTC 的 lowered-name 机制（见 raw.pyx）；
3. 复杂结构体以设备实测布局为准：用探针内核读取sizeof与成员偏移，再递归构造 dtype，避免手工对齐在异构平台上的误差；
4. 牢记 RawKernel 的视图语义：它不处理视图与跨步，传参前应确认数组的内存连续性。

仓库中 tests/cupy_tests/core_tests/test_raw.py 提供了 RawKernel/RawModule 在模板特化、错误处理等多维度上的测试覆盖，可作为进一步阅读与验证的入口。

【免费下载链接】cupyNumPy & SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy