C++/CUDA 自定义算子开发
当 Python 实现的性能不足时，PyTorch 的 C++ Extension 机制允许用 C++/CUDA 编写高性能自定义算子。
开发流程
cpp
// custom_kernel.cpp
#include <torch/extension.h>
torch::Tensor custom_add_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor custom_add(torch::Tensor a, torch::Tensor b) {
if (a.is_cuda()) return custom_add_cuda(a, b);
return a + b;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("custom_add", &custom_add, "Custom Add");
}python
# setup.py
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(ext_modules=[CUDAExtension(
name='custom_kernel',
sources=['custom_kernel.cpp', 'custom_kernel.cu']
)], cmdclass={'build_ext': BuildExtension})调试难度
CUDA Kernel 的调试比 Python 复杂得多。建议先用 Python 原型验证正确性，再用 C++/CUDA 优化性能。