CUDA Add 算子工程

一个精简的 Python 调用 CUDA ADD 算子工程。

目录结构

cuda_add_example/
├── CMakeLists.txt        # CMake 构建配置
├── setup.py              # Python 安装脚本 (推荐使用)
├── test_add.py           # 测试脚本
├── include/
│   └── add.h             # 头文件
└── src/
    ├── add_cuda.cu       # CUDA kernel 实现
    └── add_pybind.cpp    # pybind11 封装

依赖

1. PyTorch (CUDA 版本)

conda create -n torch-env python=3.12
conda activate torch-env
which python  # 确认当前是 torch-env 环境

# CUDA 13.1
pip install torch==2.10.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

2. 验证安装

python -c "import torch; print(torch.__version__)"  # 显示 PyTorch 版本
python -c "import torch; print(torch.cuda.is_available())"  # True 表示 CUDA 可用
python -c "import torch; print(torch.version.cuda)"         # 显示 CUDA 版本

3. 其他依赖

pip install setuptools  # 通常已有
pip install cmake        # 如使用 CMake 构建

构建

方式1: 使用 setup.py (推荐)

# 将 CUDA 共享库加入系统搜索路径中
echo 'export LD_LIBRARY_PATH=/root/miniconda3/envs/torch-env/lib/python3.12/site-packages/torch/lib:$LD_LIBRARY_PATH' >> ~/.bashrc
pip install -e . --no-build-isolation

pip install -e . 会找到目录下的 setup.py 并按其配置编译 C/C++ 扩展。

方式2: 使用 CMake

cd cuda_add_example
mkdir build && cd build
cmake ..
make

运行测试

python test_add.py

工作原理

test_add.py
    │
    ▼
import torch_add_cuda        # Python 加载编译好的扩展
    │
    ▼
add_pybind.cpp               # pybind11 封装层
    │
    ├──► add_cpu()           # CPU 实现 (fallback)
    │
    └──► add_cuda()          # CUDA 实现
            │
            ▼
        add_cuda.cu          # CUDA Kernel
        kernel: add_kernel    # 实际在 GPU 上执行

核心代码说明

add_cuda.cu

__global__ void add_kernel(const float* a, const float* b, float* c, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];  // 逐元素加法
    }
}

• 每个线程处理一个元素
• 使用 block/thread 索引计算全局位置
• 简单的数据并行模式

add_pybind.cpp

PYBIND11_MODULE(torch_add_cuda, m) {
    m.def("add", &add, "Element-wise add of two tensors");
}

• 暴露 add 函数给 Python
• 自动处理 Tensor 类型转换
• 同时支持 CPU 和 CUDA

扩展为真正的 PyTorch 算子

如果要集成到 PyTorch 算子注册表中，需要：

1. 实现 torch::Tensor 类型的算子
2. 使用 TORCH_LIBRARY 注册算子
3. 使用 m.def("add", &add) 导出