我们知道深度学习里，虽然接口语言是python，但是实际上还是调用的cuda代码。那么我们自己写的cuda代码该怎么提供给python调用？

最普遍的一种方法就是把C++/cuda代码导出成库，然后用python的ctypes包去调用

不过本文将专门使用PyTorch提供的转换方法

cuda代码

首先我们随便写一个cuda的函数，以及启动函数

my_cuda_op_kernel.cu

// my_cuda_op_kernel.cu
#include <cuda_runtime.h>

__global__ void multiply_by_two_kernel(float* data, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) data[idx] *= 2.0f;
}

void multiply_by_two(float* data, int size, cudaStream_t stream) {
    int block = 256;
    int grid = (size + block - 1) / block;
    multiply_by_two_kernel<<<grid, block, 0, stream>>>(data, size);
}

C++代码

然后按照PyTorch的规则，写C++函数接口

my_cuda_op.cpp

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>   // ✅ 新API头文件
#include <vector>
#include <cuda_runtime.h>

void multiply_by_two(float* data, int size, cudaStream_t stream); // CUDA kernel

torch::Tensor multiply_by_two_forward(torch::Tensor input) {
    auto output = input.clone();
    int size = output.numel();
    float* d_ptr = output.data_ptr<float>();

    // 获取当前设备
    int device = output.device().index();

    // ✅ 新的API：c10命名空间，并取出实际cudaStream_t
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream(device).stream();

    multiply_by_two(d_ptr, size, stream);
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("multiply_by_two_forward", &multiply_by_two_forward, "Multiply tensor by 2 (CUDA)");
}

setup.py

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

setup(
    name='my_cuda_op',
    ext_modules=[
        CUDAExtension(
            'my_cuda_op',
            ['my_cuda_op.cpp', 'my_cuda_op_kernel.cu'],
        )
    ],
    cmdclass={'build_ext': BuildExtension}
)

这是python编译我们c++接口的代码，输入python命令

python setup.py build_ext --inplace

会生成一个.so文件，比如我当前生成的：my_cuda_op.cpython-312-x86_64-linux-gnu.so

如果要直接使用，要显式写明白这个so文件的路径，python好像是通过修改 sys.path来实现的（如果python文件和so文件同一路径自然不用写）。

如果你想要直接当成本地环境包安装到python包列表里，可以用下面的命令

python -m pip install .

这样你使用该python环境，任何地方都可以import这个模块。

Python代码使用

import torch
import my_cuda_op

x = torch.arange(10, dtype=torch.float32, device='cuda')
print("Before:", x)

y = my_cuda_op.multiply_by_two_forward(x)
print("After:", y)

然后就可以在PyTorch里使用自己实现的cuda算子了。

输出：

Before: tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], device='cuda:0')
After: tensor([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], device='cuda:0')

C++接口代码的规则

一个 PyTorch CUDA 扩展，一般由 3 个部分组成：

文件	作用
xxx.cpp	C++ 层接口，负责 PyTorch Tensor 与 CUDA kernel 的桥接（定义 Python 调用函数）
xxx_kernel.cu	CUDA 层实现，编写具体的 GPU 算子逻辑
setup.py	构建配置文件，告诉 Python 如何编译 .cpp 和 .cu

当然为了工程可拓展性，还有头文件用来声明函数。

1️⃣ 包含头文件

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>   // 获取当前 CUDA 流
#include <cuda_runtime.h>
#include <vector>

• torch/extension.h：暴露给 Python 的扩展接口。

• c10/cuda/CUDAStream.h：新版 API，用于获取 PyTorch 的 CUDA 流。

• cuda_runtime.h：CUDA runtime 函数声明。

2️⃣ 声明 CUDA kernel 函数

C++ 文件只需要声明（不是定义）CUDA 实现，具体实现是放在cu代码里的。（如果有头文件，这个部分就放在头文件，C++文件include即可）

3️⃣ 写 PyTorch 接口函数

torch::Tensor multiply_by_two_forward(torch::Tensor input) {
    // ✅ 确保 Tensor 在 GPU 上且类型为 float32
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "Only float32 supported");

    // 克隆一个输出副本
    auto output = input.clone();
    int size = output.numel();
    float* d_ptr = output.data_ptr<float>();

    // 获取当前设备编号
    int device = output.device().index();

    // ✅ 从 PyTorch 获取当前 CUDA 流（而不是自己新建）
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream(device).stream();

    // 调用 CUDA 核函数
    multiply_by_two(d_ptr, size, stream);

    return output;
}

• input.data_ptr<float>()：获取 GPU 上的实际数据指针；

• 不做 cudaMemcpy，因为 PyTorch Tensor 已经管理了显存；

• 使用 c10::cuda::getCurrentCUDAStream() 让 kernel 运行在 PyTorch 的流中；

• 返回新的 Tensor 给 Python。

4️⃣ 注册为 Python 模块

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("multiply_by_two_forward", &multiply_by_two_forward, "Multiply tensor by 2 (CUDA)");
}

当你在 Python 中导入 my_cuda_op 时，这个 .so 模块中有一个函数叫 multiply_by_two_forward。

如果想要一个模块有多个接口，就要在这个文件里全部实现并注册接口，比如

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>

// ======================== 声明 CUDA kernel ========================
void multiply_by_two(float* data, int size, cudaStream_t stream);
void add_scalar(float* data, float scalar, int size, cudaStream_t stream);
void relu_forward(float* data, int size, cudaStream_t stream);

// ======================== 封装 PyTorch 接口 ========================
torch::Tensor multiply_by_two_forward(torch::Tensor input) {
    auto output = input.clone();
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
    multiply_by_two(output.data_ptr<float>(), output.numel(), stream);
    return output;
}

torch::Tensor add_scalar_forward(torch::Tensor input, float scalar) {
    auto output = input.clone();
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
    add_scalar(output.data_ptr<float>(), scalar, output.numel(), stream);
    return output;
}

torch::Tensor my_custom_relu(torch::Tensor input) {
    auto output = input.clone();
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
    relu_forward(output.data_ptr<float>(), output.numel(), stream);
    return output;
}

// ======================== 模块注册 ========================
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("multiply_by_two_forward", &multiply_by_two_forward, "Multiply tensor by 2 (CUDA)");
    m.def("add_scalar_forward", &add_scalar_forward, "Add scalar to tensor (CUDA)");
    m.def("my_custom_relu", &my_custom_relu, "ReLU activation (CUDA)");
}

setup.py 写法规则

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

setup(
    name='my_cuda_op',
    ext_modules=[
        CUDAExtension(
            'my_cuda_op',
            ['my_cuda_op.cpp', 'my_cuda_op_kernel.cu'],
            extra_compile_args={
                'cxx': ['-O3'],
                'nvcc': ['-O3', '-arch=sm_86'],  # 改成你的 GPU 架构
            }
        )
    ],
    cmdclass={'build_ext': BuildExtension}
)

我这个例子没写头文件，其实一般的工程，都会把

// ======================== 声明 CUDA kernel ========================
void multiply_by_two(float* data, int size, cudaStream_t stream);
void add_scalar(float* data, float scalar, int size, cudaStream_t stream);
void relu_forward(float* data, int size, cudaStream_t stream);

放到单独的头文件，cu文件做具体实现。然后C++文件仅仅是调用cuda kernel接口的函数。但是在setup编译的时候，还是只需要写cu和cpp文件即可。

如果头文件，目标文件不在同目录，则setup.py可以改成这个结构

CUDAExtension(
    name='my_cuda_ops',
    sources=['src/my_cuda_ops.cpp', 'src/my_cuda_ops_kernel.cu'],
    include_dirs=['include'],  # 👈 告诉编译器去哪里找头文件
)