PyTorch 自定义 CUDA 内核教程

本仓库包含一个用于为 PyTorch 编写自定义 CUDA 函数的教程代码。该代码基于 PyTorch 的 C 扩展示例。

免责声明

• 2019年1月2日：我撰写了 另一份更现代的教程，介绍如何使用 Makefile 构建 PyTorch C++/CUDA 扩展。相关 Git 页面位于 https://github.com/chrischoy/MakePytorchPlusPlus。
• 2018年12月9日：PyTorch CFFI 现已弃用，取而代之的是从 PyTorch 1.0 开始提供的 C++ 扩展。

本教程是在 PyTorch 尚不支持广播求和时编写的。现在 PyTorch 已经支持广播求和，因此你可能不再需要自己实现广播求和函数，但仍然可以按照本教程来构建带有自定义 CUDA 内核的自定义层。

在本仓库中，我们将构建一个简单的基于 CUDA 的广播求和函数。当前版本的 PyTorch 不支持 广播求和，因此我们必须手动扩展张量，例如使用 expand_as 方法，这会创建一个新的张量并占用额外的内存和计算资源。

例如：

a = torch.randn(3, 5)
b = torch.randn(3, 1)
# 下面这行代码会报错
# a += b

# 将 b 扩展成与 a 相同的维度
b_like_a = b.expand_as(a)
a += b_like_a

在本文中，我们将构建一个可以直接对 a 和 b 进行广播求和而无需显式扩展 b 的函数。

mathutil.broadcast_sum(a, b, *map(int, a.size()))

编写 CUDA 内核

首先，我们编写一个 CUDA 内核，它可以在不复制张量 b 的情况下将 b 加到 a 上。

__global__ void broadcast_sum_kernel(float *a, float *b, int x, int y, int size)
{
    int i = (blockIdx.x + blockIdx.y * gridDim.x) * blockDim.x + threadIdx.x;
    if(i >= size) return;
    int j = i % x; i = i / x;
    int k = i % y;
    a[IDX2D(j, k, y)] += b[k];
}

编写 C 语言封装器

编写好 CUDA 内核后，我们需要用 C 语言对其进行封装。不过，此时我们尚未使用 PyTorch 后端。请注意，输入参数已经是设备指针。

void broadcast_sum_cuda(float *a, float *b, int x, int y, cudaStream_t stream)
{
    int size = x * y;
    cudaError_t err;

    broadcast_sum_kernel<<<cuda_gridsize(size), BLOCK, 0, stream>>>(a, b, x, y, size);

    err = cudaGetLastError();
    if (cudaSuccess != err)
    {
        fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
        exit(-1);
    }
}

将 PyTorch 后端与 C 语言封装器连接

接下来，我们需要将 PyTorch 后端与我们的 C 语言封装器连接起来。我们可以使用 THCudaTensor_data 函数来暴露设备指针。指针 a 和 b 都是设备指针（位于 GPU 上）。

extern THCState *state;

int broadcast_sum(THCudaTensor *a_tensor, THCudaTensor *b_tensor, int x, int y)
{
    float *a = THCudaTensor_data(state, a_tensor);
    float *b = THCudaTensor_data(state, b_tensor);
    cudaStream_t stream = THCState_getCurrentStream(state);

    broadcast_sum_cuda(a, b, x, y, stream);

    return 1;
}

编写 Python 封装器

现在我们已经构建了 CUDA 函数和 PyTorch 函数，接下来需要将这些函数暴露给 Python，以便在 Python 中使用它们。

首先，我们将使用 nvcc 编译生成共享库。

nvcc ... -o build/mathutil_cuda_kernel.so src/mathutil_cuda_kernel.cu

然后，我们将使用 PyTorch 的 torch.utils.ffi.create_extension 函数，它会自动添加适当的头文件并构建一个可加载的 Python 共享库。

from torch.utils.ffi import create_extension

...

ffi = create_extension(
    'mathutils',
    headers=[...],
    sources=[...],
    ...
)

ffi.build()

测试！

最后，我们可以通过构建项目来测试我们的函数。在 README 中，我省略了许多细节，但你可以看到一个可运行的示例。

git clone https://github.com/chrischoy/pytorch-cffi-tutorial
cd pytorch-cffi-tutorial
make

注意事项

该函数仅接受 THCudaTensor 类型的数据，即 Python 中的 torch.FloatTensor().cuda()。

PyTorch 自定义 CUDA Kernel 快速上手指南

本指南基于 pytorch-custom-cuda-tutorial 项目，演示如何为 PyTorch 编写自定义的 CUDA 广播求和（Broadcasting Sum）函数。

重要提示：本教程基于已废弃的 PyTorch CFFI 接口（适用于 PyTorch v1.0 之前）。自 PyTorch v1.0 起，官方推荐使用 C++ Extension (torch.utils.cpp_extension)。若需生产环境使用，建议参考作者更新的教程 MakePytorchPlusPlus。本指南仅用于理解底层原理及旧代码维护。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (Ubuntu/CentOS 等)
• GPU: NVIDIA GPU 且已安装对应的驱动程序
• 编译器: gcc / g++ 和 nvcc (CUDA Toolkit)
• Python: Python 3.x
• PyTorch: 本教程代码依赖旧版 API，建议在与原项目兼容的 PyTorch 版本（< 1.0）中运行，或仅作为逻辑参考。
• 依赖库:

  pip install torch cffi numpy
  

安装步骤

1. 克隆项目仓库 获取源代码到本地：

   git clone https://github.com/chrischoy/pytorch-cffi-tutorial.git
   cd pytorch-cffi-tutorial
   

2. 编译扩展模块 进入项目目录后，直接运行 make 命令。该命令将调用 nvcc 编译 CUDA 内核，并使用 cffi 构建 Python 可调用的共享库。

   make
   

   注：如果编译成功，将在 build/ 目录下生成 mathutil_cuda_kernel.so 文件。

基本使用

编译完成后，您可以在 Python 中导入该模块并测试自定义的广播求和功能。该功能允许在不显式扩展张量内存的情况下执行 a += b 操作。

使用示例：

import torch
import mathutil  # 编译生成的模块

# 创建测试数据
# a: shape (3, 5)
a = torch.randn(3, 5).cuda()
# b: shape (3, 1)，需要广播到 (3, 5) 进行相加
b = torch.randn(3, 1).cuda()

# 记录原始值以便验证
a_original = a.clone()

# 调用自定义 CUDA 函数
# 参数说明：tensor_a, tensor_b, dim_x, dim_y
# 此处对应 a 的形状 (3, 5)
mathutil.broadcast_sum(a, b, 3, 5)

# 验证结果：自定义函数结果应等同于 expand_as 后的结果
b_expanded = b.expand_as(a_original)
expected = a_original + b_expanded

# 检查差异
diff = torch.abs(a - expected).max()
print(f"Max difference: {diff.item()}")

if diff.item() < 1e-5:
    print("Success: Custom CUDA kernel works correctly!")
else:
    print("Error: Results mismatch.")

核心优势： 传统方法需要使用 b.expand_as(a) 创建新张量，消耗额外显存和计算资源。本教程实现的 broadcast_sum 直接在 GPU 内核中处理索引映射，实现了零拷贝的广播加法。