Intro

如何将写好的kernel优雅地嵌入Python流程中。
Pytorch提供了 torch.utils.cpp_extension 工具箱，让我们能够以极小代价将C++/CUDA代码编译为Python模块。这里会梳理三种主流方式，并解释为什么要自定义算子，如何避免环境坑以及如何让自定义kernel无缝支持 autograd

核心概念：

概念	作用	典型位置
torch.utils.cpp_extension.load_inline	直接编译字符串中的C++/CUDA代码	Jupyter/快速原型
torch.utils.cpp_extension.load	编译指定 .cpp/.cu 文件	脚本中临时编译
CUDAExtension / CppExtension	在 setup.py 中声明扩展模块	正式打包分发
BuildExtension	替换 setuptools 默认构建命令，注入PyTorch编译参数	setup.py 的 cmdclass
PYBIND11_MODULE	将C++函数暴露给Python的宏	.cpp 文件末尾
torch.autograd.Function	自定义算子前向/反向传播的包装类	让算子支持自动微分

本质关系：load 和 load_inline 是JIT便捷工具，背后调用相同地编译器逻辑； setup.py + CUDAExtension是预编译方案，更适合生产关系

方式一： load_inline 即时编译

什么时候用：在notebook中快速验证小kernel或者写一次性实验脚本
C++编译需要ninja，需要提前安装

import torch
from torch.utils.cpp_extension import load_inline

# C++ 头文件声明
cpp_src = """
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
"""

# CUDA 实现（注意：不包含 PYBIND11_MODULE 宏）
cuda_src = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void add_kernel(const float* a, const float* b, float* c, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) c[i] = a[i] + b[i];
}

torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
    auto c = torch::empty_like(a);
    int N = a.numel();
    const int threads = 256;
    const int blocks = (N + threads - 1) / threads;
    add_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), c.data_ptr<float>(), N);
    return c;
}
"""

# 自动编译并加载，add_cuda 将绑定为模块的 add 方法
module = load_inline(
    name="inline_cuda_add",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["add_cuda"],          # 指定要暴露的函数名
    verbose=True
)

# 测试
a = torch.randn(1000, device='cuda')
b = torch.randn(1000, device='cuda')
c = module.add_cuda(a, b)            # 调用暴露的 C++ 函数
print(torch.allclose(c, a + b))      # True

优点：无需进行任何文件操作，代码与结果同屏，调试只管
缺点：每次运行都要重新编译且错误定位较难

方式二：load与预编译扩展

load 是 load_inline的文件版，适合脚本中编译已有的 .cu/.cpp 文件

from torch.utils.cpp_extension import load

vector_add = load(
    name="vector_add_ext",
    sources=["add_kernel.cu"],
    verbose=True,
    extra_cuda_cflags=["-O3"]
)

x = torch.ones(10, device='cuda')
y = torch.ones(10, device='cuda')
print(vector_add.add(x, y))

对应的 add_kernel.cu 与后卫的 vector_add.cu 类似。load会在 ~/.cache/torch_extension 下生成预编译产物。第二次运行将跳过编译直接加载

常见的坑：
- CUDA版本不匹配，load 会检测nvcc版本，若与torch.version.cuda 不一致则报错。
- 重复编译耗时：代码量较大时建议使用setup.py预编译

load 和 load_inline参数解释

这两者都是PyTorch提供的JIT编译工具，核心参数基本一致
- verbose(bool): 控制是否打印详细的编译日志(False 为静默编译，True输出完整的编译指令)
- name(str): 生成对应的python模块名，也是编译中间文件的目录名
- sources：需要指定编译的源文件
- functions(仅load_inline)：指定需要暴露给Python的C++函数名列表。绑定的Python 方法名与C++函数名相同
- extra_cuda_cflags/ extra_cxxflags: 向nvcc或者C++传递额外的编译选项

1 2	extra_cuda_cflags = ["-O3", "-arch=sm_80"] extra_cxxflags = ["-O3"]

常用选项：-O3 (最高优化), -g (调试符号), -arch=sm_xx(指定GPU计算能力)

方式三：setup.py 与 CUDAExtension

这是正式项目的首选，只需要写好setup.py，运行pip install . 即可将苦熬站永久安装到当前Python环境

1. 项目结构

custom_ops/
├── setup.py
├── cuda_add/
│   └── vector_add.cu

2. 算子编写注意事项

#include <torch/extension.h>
#include <cuda_runtime.h>

// 高性能向量加法 kernel (float4 向量化)
__global__ void add_kernel_vec4(const float* __restrict__ a,
                                const float* __restrict__ b,
                                float* __restrict__ c,
                                int N) {
    const float4* a4 = reinterpret_cast<const float4*>(a);
    const float4* b4 = reinterpret_cast<const float4*>(b);
    float4* c4 = reinterpret_cast<float4*>(c);

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    int N4 = N / 4;

    for (int i = idx; i < N4; i += stride) {
        float4 av = a4[i];
        float4 bv = b4[i];
        float4 cv;
        cv.x = av.x + bv.x;
        cv.y = av.y + bv.y;
        cv.z = av.z + bv.z;
        cv.w = av.w + bv.w;
        c4[i] = cv;
    }

    int remainder_start = N4 * 4;
    for (int i = remainder_start + idx; i < N; i += stride) {
        c[i] = a[i] + b[i];
    }
}

torch::Tensor add_forward(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.device().is_cuda(), "a must be CUDA tensor");
    TORCH_CHECK(b.device().is_cuda(), "b must be CUDA tensor");
    auto c = torch::empty_like(a);
    int N = a.numel();
    const int threads = 256;
    int blocks = ((N / 4) + threads - 1) / threads;
    if (blocks == 0) blocks = 1;
    add_kernel_vec4<<<blocks, threads>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        c.data_ptr<float>(),
        N
    );
    return c;
}

// 自动微分包装
class AddFunction : public torch::autograd::Function<AddFunction> {
public:
    static torch::Tensor forward(
        torch::autograd::AutogradContext* ctx,
        torch::Tensor a,
        torch::Tensor b) {
        ctx->save_for_backward({a, b});
        return add_forward(a, b);
    }

    static torch::autograd::variable_list backward(
        torch::autograd::AutogradContext* ctx,
        torch::autograd::variable_list grad_output) {
        auto grad = grad_output[0];
        return {grad, grad};
    }
};

torch::Tensor add_autograd(torch::Tensor a, torch::Tensor b) {
    return AddFunction::apply(a, b);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("add", &add_autograd, "Vector addition with autograd");
    m.def("add_forward", &add_forward, "Raw forward (no autograd)");
}

3. 编写setup.py

from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

setup(
    name="cuda_add",
    version="0.1.0",
    ext_modules=[
        CUDAExtension(
            name="cuda_add",
            sources=["vector_add.cu"],
            extra_compile_args={"cxx": ["-O3"], "nvcc": ["-O3"]}
        )
    ],
    cmdclass={"build_ext": BuildExtension},
    install_requires=["torch"],
)

安装与使用

1	python setup.py install # 或 pip install .

import torch
import cuda_add
a = torch.randn(10_000_000, device='cuda')
b = torch.randn(10_000_000, device='cuda')
c = cuda_add.add(a, b)          # 带 autograd
c_no_grad = cuda_add.add_forward(a, b)  # 纯前向，稍快

为什么推荐预编译？
- 一次编译，随处import：省去JIT等待时间
- 依赖管理清晰：可指定pytorch版本，打包上传PyPI
- 错误提示友好：编译失败会生成完整日志

自动微分

案例：RGB转灰度与三维归一化

1. 算子定义

#include <torch/extension.h>
#include <cuda_runtime.h>

// 定义 uchar3 结构体（与 OpenCV 对应）
struct uchar3 {
    unsigned char x, y, z;
};

__global__ void rgbToGrayKernel(const uchar3* img, unsigned char* gray,
                                int width, int height) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x < width && y < height) {
        int idx = y * width + x;
        uchar3 pixel = img[idx];
        gray[idx] = (unsigned char)(0.299f * pixel.x + 0.587f * pixel.y + 0.114f * pixel.z);
    }
}

torch::Tensor rgb_to_gray_cuda(torch::Tensor img) {
    // 假设输入是 uint8 [H, W, 3] 且已在 GPU 上
    TORCH_CHECK(img.dim() == 3 && img.size(2) == 3, "Input must be HxWx3");
    TORCH_CHECK(img.dtype() == torch::kUInt8, "Input must be uint8");
    int height = img.size(0);
    int width = img.size(1);
    auto gray = torch::empty({height, width}, img.options().dtype(torch::kUInt8));

    dim3 block(16, 16);
    dim3 grid((width + 15) / 16, (height + 15) / 16);
    rgbToGrayKernel<<<grid, block>>>(
        reinterpret_cast<const uchar3*>(img.data_ptr<unsigned char>()),
        gray.data_ptr<unsigned char>(),
        width, height
    );
    return gray;
}

__global__ void normalizeVolumeKernel(const unsigned short* in, float* out,
                                      int dimX, int dimY, int dimZ, float maxVal) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int z = blockIdx.z * blockDim.z + threadIdx.z;
    if (x < dimX && y < dimY && z < dimZ) {
        int idx = z * dimY * dimX + y * dimX + x;
        out[idx] = (float)in[idx] / maxVal;
    }
}

torch::Tensor normalize_volume_cuda(torch::Tensor volume, float maxVal) {
    TORCH_CHECK(volume.dim() == 3, "Volume must be 3D");
    TORCH_CHECK(volume.dtype() == torch::kUInt16, "Input must be uint16");
    int dimX = volume.size(2);  // 注意 PyTorch 维度顺序 DxHxW，这里假设 Z,Y,X
    int dimY = volume.size(1);
    int dimZ = volume.size(0);
    auto out = torch::empty_like(volume, volume.options().dtype(torch::kFloat32));

    dim3 block(8, 8, 4);
    dim3 grid((dimX + 7) / 8, (dimY + 7) / 8, (dimZ + 3) / 4);
    normalizeVolumeKernel<<<grid, block>>>(
        volume.data_ptr<unsigned short>(),
        out.data_ptr<float>(),
        dimX, dimY, dimZ, maxVal
    );
    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("rgb_to_gray", &rgb_to_gray_cuda, "RGB to grayscale (CUDA)");
    m.def("normalize_volume", &normalize_volume_cuda, "3D volume normalization (CUDA)");
}

Python 测试代码

import torch
import cv2
import numpy as np
import custom_image_ops  # 编译后的模块名

# RGB 转灰度测试
img = cv2.imread("input.jpg")
img_tensor = torch.from_numpy(img).cuda()  # HxWx3
gray = custom_image_ops.rgb_to_gray(img_tensor)
cv2.imwrite("gray.png", gray.cpu().numpy())

# 三维归一化测试
vol = torch.randint(0, 4096, (128, 256, 256), dtype=torch.uint16, device='cuda')
norm = custom_image_ops.normalize_volume(vol, 4095.0)
print(norm.min(), norm.max())  # 0.0 ~ 1.0