mxnet.rtc 运行时 CUDA 编译指南：在 MXNet 中直接编写并启动自定义 CUDA Kernel

• 深度学习
• 人工智能
• 机器学习
• 分布式训练

导读

mxnet.rtc是 Apache MXNet 面向 GPU 用户的「运行时 CUDA 编译（Runtime CUDA Compilation, RTC）」模块：它允许你在 Python 中直接书写 CUDA C/C++ 内核源码，在运行时通过 NVRTC 编译成 PTX，再与 MXNet 的 NDArray 和自动调度引擎无缝衔接并启动执行，全程无需重编译 MXNet、无需编写 C++ 插件或自定义算子。读完本文，你将掌握CudaModule/CudaKernel的完整用法、extern "C"与exports两种导出方式的取舍、签名解析与参数类型规则、编译选项与 SASS 直编的底层原理，并能把自定义 CUDA Kernel 直接作用于mx.nd.NDArray，复用 MXNet 的引擎调度与依赖管理。本文以官方 API 文档 docs/python_docs/python/api/rtc/index.rst（其内容由 python/mxnet/rtc.py 的 docstring 承载）为骨架，结合源码与测试展开。

1. rtc 是什么：面向 MXNet 的运行时 CUDA 编译接口

mxnet.rtc在 MXNet Python 包中通过 python/mxnet/init.py 的from . import rtc导出，其模块 docstring 自述为 "Interface to runtime cuda kernel compile module"（python/mxnet/rtc.py）。它解决的核心问题很直接：当标准算子无法覆盖某些自定义 GPU 计算时，传统做法是编译进自定义算子（涉及 C++/CUDA 代码、MSHADOW_XINLINE宏与整个 MXNet 重编译），而 rtc 让你把 CUDA 源码作为字符串传入，在运行期完成「编译 → 加载 → 启动」全流程：

1. 编译：借助 NVIDIA NVRTC（libnvrtc）在运行时把 CUDA C 源码编译为 PTX（或特定架构的 SASS/cubin）。
2. 加载：通过 CUDA Driver API（cuModuleLoadDataEx）把编译产物加载为 CUmodule。
3. 启动：通过cuLaunchKernel按用户指定的 grid/block 维度启动内核，并同步等待完成。

这一整套逻辑的 C++ 实现位于 src/common/rtc.cc 与 include/mxnet/rtc.h，全部在MXNET_USE_CUDA宏保护下，即该功能仅在启用 CUDA 的构建中可用；从 include/mxnet/rtc.h 的CHECK_EQ(ctx.dev_mask(), Context::kGPU)也可确认内核只能在 NVIDIA GPU 上启动。Python 层通过 include/mxnet/c_api.h 暴露的MXRtcCudaModuleCreate/MXRtcCudaKernelCreate/MXRtcCudaKernelCall等 C API 与底层对接。

从源码结构看，rtc模块对外只暴露两个类：CudaModule（编译并持有 CUDA 源码对应的模块）与CudaKernel（由CudaModule.get_kernel产生、负责启动的 kernel 句柄），本文其余章节将围绕这两个类的完整使用展开。

2. 快速上手：第一个 rtc 内核（axpy 示例）

CudaModule的 docstring（python/mxnet/rtc.py）给出了一个完整的端到端示例。在 CUDA 7.5 时代（也兼容所有后续版本）最稳妥的写法是用extern "C"修饰内核，以避免 C++ 名称修饰（name mangling）导致无法按原名查找：

import mxnet as mx source = r''' extern "C" __global__ void axpy(const float *x, float *y, float alpha) { int i = threadIdx.x + blockIdx.x * blockDim.x; y[i] += alpha * x[i]; } ''' module = mx.rtc.CudaModule(source) func = module.get_kernel("axpy", "const float *x, float *y, float alpha") x = mx.nd.ones((10,), ctx=mx.gpu(0)) y = mx.nd.zeros((10,), ctx=mx.gpu(0)) func.launch([x, y, 3.0], mx.gpu(0), (1, 1, 1), (10, 1, 1)) print(y)

运行后y的每个元素都为3.0。这个例子覆盖了 rtc 的全部核心流程：

2.1 参数命名与作用

CudaModule构造函数的三个参数（python/mxnet/rtc.py）：

• source : str— 完整的 CUDA 源码字符串，是唯一必填参数。
• options : tuple of str— 传给 NVRTC 的编译选项，例如"-I/usr/local/cuda/include"用于向 include path 追加 CUDA 头文件目录；也可直接传单个字符串（构造时若检测到string_types会自动包装成元组）。多个选项可组合使用。
• exports : tuple of str— 需要按名称导出的内核名（仅 CUDA 8.0+ 支持，详见第 3 节）；同样支持字符串自动包装。

一个典型的实际用法是配合mxnet.util.get_rtc_compile_opts获取针对当前 GPU 架构的编译选项（详见第 5 节），例如官方 GPU 测试 tests/python/gpu/test_operator_gpu.py 中module = mx.rtc.CudaModule(source, options=compile_opts)的写法。

3. 两种内核导出方式：extern "C" 与 exports

C++ 会对函数名做名称修饰（name mangling），导致内核的实际符号名与源码中的函数名不一致。mxnet.rtc提供了两条解决路径，这也是官方文档（python/mxnet/rtc.py）重点区分的内容。

3.1 方式一：extern "C"（CUDA 7.5 及所有版本通用）

在 CUDA 7.5 中（以及任何版本下不想用 exports 时），内核定义需以extern "C"开头以避免名称修饰：

extern "C" __global__ void axpy(const float *x, float *y, float alpha) { ... }

此时get_kernel("axpy", ...)直接使用"axpy"作为查找名。C++ 端在 src/common/rtc.cc 的GetKernel中，若名字未命中 exports 集合，mangled_name就保持用户传入的原名；随后 src/common/rtc.cc 通过cuModuleGetFunction从已加载模块中查找该符号，若找不到会输出明确提示：请为内核定义加extern "C"，或在创建CudaModule时把名字加入exports。

3.2 方式二：exports（CUDA 8.0+，支持模板）

从 CUDA 8.0 开始，可以通过exports按名称导出函数，这同时解锁了模板内核的使用，例如同一份模板源码分别实例化float与double两个版本：

source = r''' template<typename DType> __global__ void axpy(const DType *x, DType *y, DType alpha) { int i = threadIdx.x + blockIdx.x * blockDim.x; y[i] += alpha * x[i]; } ''' module = mx.rtc.CudaModule(source, exports=['axpy<float>', 'axpy<double>']) func32 = module.get_kernel("axpy<float>", "const float *x, float *y, float alpha") x = mx.nd.ones((10,), dtype='float32', ctx=mx.gpu(0)) y = mx.nd.zeros((10,), dtype='float32', ctx=mx.gpu(0)) func32.launch([x, y, 3.0], mx.gpu(0), (1, 1, 1), (10, 1, 1)) print(y) func64 = module.get_kernel("axpy<double>", "const double *x, double *y, double alpha") x = mx.nd.ones((10,), dtype='float64', ctx=mx.gpu(0)) y = mx.nd.zeros((10,), dtype='float64', ctx=mx.gpu(0)) func64.launch([x, y, 3.0], mx.gpu(0), (1, 1, 1), (10, 1, 1)) print(y)

注意上面示例中两处launch都打印的是func32的结果，这是官方 docstring 的既有写法；实际使用中第二次应调用func64.launch(...)才会得到 float64 的输出。此外，模板实例名（如axpy<float>）含特殊字符，不能用extern "C"修饰，因此这类场景必须走exports路径。

底层实现上，src/common/rtc.cc 在 CUDA 8.0+ 时对每个导出名调用nvrtcAddNameExpression登记表达式，编译完成后在GetKernel（src/common/rtc.cc）中通过nvrtcGetLoweredName查询真实修饰名；而在低于 CUDA 8.0 的构建中，src/common/rtc.cc 会直接CHECK_EQ(exports.size(), 0)并给出错误信息：导出功能仅在 CUDA 8.0 及以上可用，低版本请改用extern "C"。

4. get_kernel 与签名解析：参数类型如何被解读

4.1 签名语法

CudaModule.get_kernel(name, signature)（python/mxnet/rtc.py）的signature参数描述内核的函数签名。例如内核声明为

extern "C" __global__ void axpy(const float *x, double *y, int alpha)

则签名可写作完整的

const float *x, double *y, int alpha

也可以省略参数名、只保留类型与指针标记：

const float *, double *, int

其中两条规则决定了参数如何映射到 Python 侧：

• 签名中的*标记该参数为数组（NDArray）；
• 签名中的const标记该参数为常量（只读输入）数组。

4.2 解析器实现与支持的类型

签名在 Python 侧由 python/mxnet/rtc.py 中的正则表达式解析：

pattern = re.compile(r"""^(const)?\s?([\w_]+)\s?(\*)?\s?([\w_]+)?$""")

每个参数被拆解为(const) 类型 (*) (参数名)四段：第一组捕获const（记入is_const）、第二组捕获类型名、第三组捕获*（记入is_ndarray）。若格式不合法（如组内出现const位置错误），会抛出ValueError，提示签名必须符合"(const) type (*) (name)"形式。支持的类型映射定义在模块级字典_DTYPE_CPP_TO_NP（python/mxnet/rtc.py）：

不在表内的类型会在get_kernel阶段直接抛出TypeError，并列出全部受支持类型，因此类型检查发生在取内核时，而不是启动时。

4.3 C++ 侧的一致性校验

即使 Python 侧通过了类型检查，C++ 端在启动时仍会做严格校验（src/common/rtc.cc）：对每个被标记为 NDArray 的参数，会CHECK_EQ(array.dtype(), arg_types[i].dtype)确认传入 NDArray 的 dtype 与签名声明一致，不一致时给出包含期望类型与实际类型的错误信息。同时内核被提交到 MXNet 引擎时，is_const标记决定了依赖关系方向：只读数组进入read_vars、输出数组进入write_vars（src/common/rtc.cc），从而让 MXNet 的依赖调度器（dep scheduler）自动保证内核与前序/后续算子之间的数据依赖，这正是 rtc 能无缝融入mx.nd计算图的关键设计。

5. CudaKernel.launch：启动参数与调度细节

CudaKernel.launch(args, ctx, grid_dims, block_dims, shared_mem=0)（python/mxnet/rtc.py）的参数约定如下：

Python 侧在 python/mxnet/rtc.py 做前置断言：ctx.device_type必须为'gpu'，grid_dims/block_dims必须是长度 3 的元组，参数个数必须与签名类型数一致（否则报 "CudaKernel({name}) expects {n} arguments but got {m}"）。随后将 NDArray 参数直接取其handle，数值参数则按声明 dtype 用np.array(arg, dtype=dtype)打包并取内存地址，最终调用 C APIMXRtcCudaKernelCall启动。

C++ 端 src/common/rtc.cc 的Kernel::Launch展示了完整的调度链路：

1. 按设备缓存CUfunction（func_[ctx.dev_id]），避免重复查找；
2. 收集只读/可写 NDArray 的Engine::VarHandle作为依赖变量；
3. Engine::Get()->PushSync(...)把启动动作作为 MXNet 引擎任务提交——引擎会等待所有read_vars依赖的算子完成、并保证后续算子等待本内核的write_vars；
4. 任务内部对 NDArray 参数取dptr_指针、数值参数做MSHADOW_TYPE_SWITCH分派，随后cuLaunchKernel在 MXNet 的 GPU 流上启动，并cudaStreamSynchronize同步等待内核完成。

5.1 实战：测试用例中的共享内存用法

仓库自带的 GPU 测试 tests/python/gpu/test_rtc.py 演示了动态共享内存与expf数学函数的组合用法：

import mxnet as mx import numpy as np from numpy.testing import assert_allclose x = mx.nd.zeros((10,), ctx=mx.gpu(0)) x[:] = 1 y = mx.nd.zeros((10,), ctx=mx.gpu(0)) y[:] = 2 rtc = mx.rtc('abc', [('x', x)], [('y', y)], """ __shared__ float s_rec[10]; s_rec[threadIdx.x] = x[threadIdx.x]; y[threadIdx.x] = expf(s_rec[threadIdx.x]*5.0);""") rtc.push([x], [y], (1, 1, 1), (10, 1, 1)) assert_allclose(y.asnumpy(), np.exp(x.asnumpy()*5.0))

该测试在python/mxnet/rtc.py之外、通过mx.rtc模块名直接以旧式函数接口（rtc(name, in_args, out_args, source)+push）运行，并验证了结果与np.exp(x * 5.0)一致。测试文件中虽然用了__shared__静态共享内存，但launch的shared_mem参数同样支持为extern __shared__数组动态分配共享内存（默认 0 字节）。

6. 编译选项：架构目标、include 路径与 SASS 直编

CudaModule(source, options=...)的options会原样透传给 NVRTC。常见用途包括：

• "-I/path/to/cuda/include"：追加 CUDA 头文件搜索路径，用于在源码中使用自定义头文件；
• "--gpu-architecture=sm_80"（或compute_80等）：指定编译目标架构。

6.1 自动选择架构编译选项

MXNet 在 python/mxnet/util.py 提供get_rtc_compile_opts(device)工具函数，可针对运行设备自动生成合适的--gpu-architecture选项：

def get_rtc_compile_opts(device): device_cc = get_cuda_compute_capability(device) # 当前设备算力 max_supported_cc = get_max_supported_compute_capability() # NVRTC 支持的最大算力 can_compile_to_SASS = max_supported_cc >= 86 # CUDA 11.1+ 支持 sm_86 直编 SASS should_compile_to_SASS = can_compile_to_SASS and device_cc <= max_supported_cc device_cc_as_used = min(device_cc, max_supported_cc) arch_opt = "--gpu-architecture={}_{}".format("sm" if should_compile_to_SASS else "compute", device_cc_as_used) return [arch_opt]

其决策逻辑值得展开：若 NVRTC 支持的最高算力 ≥ 86（即 CUDA 11.1 及以后）且设备算力不超过该上限，则直接编译为sm_XX目标（SASS）；否则回退为compute_XX目标（PTX，运行时再 JIT）。官方 GPU 测试 tests/python/gpu/test_operator_gpu.py 在构建CudaModule时正是传入options=get_rtc_compile_opts(ctx)。

6.2 底层如何决定 PTX 还是 CUBIN

C++ 端 src/common/rtc.cc 会根据编译选项决定产物形式：遍历options，若发现包含"sm_"的选项（即显式指定了具体 SM 架构），则use_ptx = false，转而通过nvrtcGetCUBINSize/nvrtcGetCUBIN获取 CUBIN（SASS）二进制——但该路径要求 CUDA 11.1+，否则LOG(FATAL)提示"请改用 compute_XX 目标或升级到 CUDA 11.1+"；默认情况下（无sm_选项）则走nvrtcGetPTXSize/nvrtcGetPTX生成 PTX。无论哪种产物，最终都在 src/common/rtc.cc 通过cuModuleLoadDataEx按设备加载，并对每个设备（ctx.dev_id）缓存一个CUmodule。

7. 与 MXNet 生态的协作方式

7.1 在 NDArray 计算流中使用 rtc

rtc 内核的输入输出就是mx.nd.NDArray，因此天然可以出现在任意 MXNet 计算序列中。由于内核通过引擎的PushSync提交并登记了read_vars/write_vars（src/common/rtc.cc），它可以与前序算子（生产者）和后序算子（消费者）自动建立依赖，无需手动同步——这是 rtc 相比"裸"cudaLaunchKernel的最大优势。

7.2 C API 层面对接

Python 层封装通过 include/mxnet/c_api.h 的四个接口完成：MXRtcCudaModuleCreate（编译源码并返回模块句柄）、MXRtcCudaModuleFree（释放模块）、MXRtcCudaKernelCreate（按名称/签名创建内核）、MXRtcCudaKernelCall（启动内核），均通过check_call检查返回码。此外，同文件还保留了一套更早期的MXRtcCreate/MXRtcPush/MXRtcFree接口（include/mxnet/c_api.h），对应 tests/python/gpu/test_rtc.py 中旧式mx.rtc(...)+push(...)的用法。若你有更复杂的自定义算子需求，也可以参考仓库中 example/extensions/lib_custom_op 的算子注册流程，但 rtc 的优势在于无需任何编译期集成。

8. 使用限制与注意事项

综合文档与源码，使用mxnet.rtc时需注意以下边界：

• 仅限 NVIDIA GPU：从 Python 的assert ctx.device_type == 'gpu'（python/mxnet/rtc.py）到 C++ 的CHECK_EQ(ctx.dev_mask(), Context::kGPU)（include/mxnet/rtc.h、src/common/rtc.cc），全链路仅支持 GPU 上下文；CPU 环境无法使用。
• 仅限 CUDA 构建：rtc.h整体包裹在#if MXNET_USE_CUDA中，未启用 CUDA 的 MXNet 构建不包含该功能。
• 名称查找规则：不用extern "C"也不加入exports的内核将无法被get_kernel找到，运行时cuModuleGetFunction返回CUDA_ERROR_NOT_FOUND并给出修正提示（src/common/rtc.cc）。
• 类型限制：内核参数类型仅限第 4.2 节表格中的类型；数值参数按签名 dtype 打包，NDArray 参数 dtype 必须与签名严格一致，否则在启动时抛错。
• CUDA 版本相关行为：exports需要 CUDA 8.0+；sm_XX目标直编 SASS 需要 CUDA 11.1+；低版本只能输出 PTX。
• 同步语义：每次launch都会cudaStreamSynchronize等待内核完成（src/common/rtc.cc），即当前实现是同步启动，适合自定义 kernel 的调试与确定性执行。

9. 小结

mxnet.rtc为 MXNet 用户提供了一条零重编译的自定义 CUDA 内核路径：Python 中书写源码 → NVRTC 运行时编译 → Driver API 加载 → 引擎调度启动，输入输出全部复用mx.nd.NDArray，并自动纳入 MXNet 的依赖调度。官方文档的核心内容——CudaModule的source/options/exports参数、extern "C"与 CUDA 8.0+exports两种导出方式、模板内核实例化、get_kernel签名语法（*表示数组、const表示只读）、launch的 grid/block/shared_mem 约定、受支持的类型表——在本文中均有完整继承与展开；编译产物选择（PTX vs CUBIN）、按设备缓存、引擎依赖登记、get_rtc_compile_opts的架构决策等实现细节则补充自 src/common/rtc.cc、include/mxnet/rtc.h、python/mxnet/util.py 与 tests/python/gpu/test_rtc.py 等仓库证据。相关 API 的正式索引可见 docs/python_docs/python/api/rtc/index.rst。

• 深度学习
• 人工智能
• 机器学习
• 分布式训练