【CUDA】Pytorch_Extensions
为什么要开发CUDA扩展？
当我们在PyTorch中实现自定义算子时，通常有两种选择：
- 使用纯Python实现（简单但效率低）
- 使用C++/CUDA扩展（高效但需要编译）
对于计算密集型的操作（如神经网络中的自定义激活函数），使用CUDA扩展可以获得接近硬件极限的性能。本文将以实现一个多项式激活函数x² + x + 1为例，展示完整的开发流程。
完整CUDA扩展代码解析
polynomial_cuda.cu
1. 编写CUDA内核
首先，我们来编写CUDA内核。我们将实现一个简单的多项式激活函数 x^2 + x + 1，该操作在GPU上执行。
template <typename scalar_t>
__global__ void polynomial_activation_kernel(const scalar_t* __restrict__ x, scalar_t* __restrict__ output,size_t size) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < size) {scalar_t val = x[idx];output[idx] = val * val + val + 1; // x^2 + x + 1}
}
在这个内核中，x 是输入张量，output 是输出张量。每个线程处理输入张量中的一个元素并计算结果。
__restrict__ 是一个关键字，用于告诉编译器某个指针是“独占”的，不会与其他指针指向的内存区域重叠。这允许编译器进行更激进的优化，例如：避免重复加载内存；重新排序指令以提高性能。
2. 使用PyTorch进行内核封装
为了使CUDA内核在Python中可用，我们需要通过PyTorch的扩展API来封装它。这一步使用了 AT_DISPATCH_FLOATING_TYPES 宏，它会根据输入张量的类型选择合适的内核类型。
AT_DISPATCH_FLOATING_TYPES(scalar_type, name, code)
scalar_type：这通常是张量类型，如x.scalar_type()，表示输入张量的实际数据类型（例如float32、float64）。name：这是一个字符串，用于表示当前的操作名称，它通常用于调试信息中，帮助识别当前的操作。code：这是一个 lambda 表达式（或其他可以执行的代码块），在这个块内会根据具体的类型选择合适的内核类型并执行。
torch::Tensor polynomial_activation_cuda(torch::Tensor x) {auto output = torch::empty_like(x);int threads = 1024;int blocks = (x.numel() + threads - 1) / threads;AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "polynomial_activation_cuda", ([&] {polynomial_activation_kernel<scalar_t><<<blocks, threads>>>(x.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(), x.numel());}));return output;
}
3. 绑定CUDA内核到Python
接下来，通过 pybind11 来将这个CUDA内核暴露给Python。PYBIND11_MODULE 是 pybind11 提供的一个宏，用于创建一个 Python 扩展模块。它的语法如下：
PYBIND11_MODULE(name, m) {// 绑定的内容
}
name：模块的名字。这里使用了TORCH_EXTENSION_NAME宏，它是一个由 PyTorch 在编译时定义的名字，通常是你扩展的 C++ 文件的名字。PyTorch 会自动为扩展生成这个名字。使用TORCH_EXTENSION_NAME可以确保模块名与扩展文件名一致。m：这是一个模块对象，是由pybind11自动传递的，通过它我们可以将函数、类等绑定到 Python 模块中。你可以将其看作是 Python 模块的“句柄”，所有要暴露给 Python 的 C++ 函数都需要通过m来注册。
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("polynomial_activation", &polynomial_activation_cuda, "Polynomial activation (CUDA)");
}
m.def("polynomial_activation", &polynomial_activation_cuda, "Polynomial activation (CUDA)")
这行代码将 C++ 中的函数 polynomial_activation_cuda 绑定到 Python 模块，并指定其在 Python 中的名称。
m.def()：这是pybind11用于绑定函数的接口。它的作用是将 C++ 函数暴露给 Python，使得在 Python 中可以直接调用。def函数接受三个参数：- 第一个参数是 Python 中的函数名
"polynomial_activation"。这个名字是在 Python 中用来调用这个函数的。 - 第二个参数是 C++ 函数的指针
&polynomial_activation_cuda。&polynomial_activation_cuda是指向 C++ 函数的指针，它告诉pybind11这个 Python 函数应当调用哪个 C++ 函数。 - 第三个参数是函数的文档字符串
"Polynomial activation (CUDA)"，它用来描述函数的功能。在 Python 中使用help()查看这个函数时会显示这个文档。
- 第一个参数是 Python 中的函数名
这个模块将使得我们可以像使用任何其他Python函数一样，直接在Python中调用 polynomial_activation。
这三个部分构成完整的polynomial_cuda.cu:
#include <cuda.h>
#include <cuda_runtime.h>
#include <torch/extension.h>template <typename scalar_t>
__global__ void polynomial_activation_kernel(const scalar_t* __restrict__ x, scalar_t* __restrict__ output,size_t size) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < size) {scalar_t val = x[idx];output[idx] = val * val + val + 1; // x^2 + x + 1}
}torch::Tensor polynomial_activation_cuda(torch::Tensor x) {auto output = torch::empty_like(x);int threads = 1024;int blocks = (x.numel() + threads - 1) / threads;AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "polynomial_activation_cuda", ([&] {polynomial_activation_kernel<scalar_t><<<blocks, threads>>>(x.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(), x.numel());}));return output;
}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("polynomial_activation", &polynomial_activation_cuda, "Polynomial activation (CUDA)");
}
setup.py
然后通过使用PyTorch提供的 setuptools 工具来构建和安装扩展
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension,CUDAExtension
# setuptools.setup是 Python 的标准构建工具，用于定义和构建 Python 包。
# torch.utils.cpp_extension.BuildExtension是 PyTorch 提供的一个自定义构建扩展类，用于处理 CUDA/C++ 扩展的编译。
# torch.utils.cpp_extension.CUDAExtension是 PyTorch 提供的一个类，用于定义 CUDA 扩展模块,指定了需要编译的 CUDA 源文件（.cu 文件）以及编译选项setup(# 包的名称name="polynomial_cuda", ext_modules=[# CUDAExtension 参数：第一个是扩展模块的名称（'polynomial_cuda'），即编译后的模块在 Python 中导入时的名称。# 第二个参数是一个列表，包含需要编译的 CUDA 源文件（这里是 polynomial_cuda.cu）CUDAExtension('polynomial_cuda',['polynomial_cuda.cu',]),],# 自定义构建命令，这里使用 BuildExtension 来构建 CUDA 扩展。cmdclass={# BuildExtension 的作用：自动配置编译器选项，确保 CUDA 代码能够与 PyTorch 正确链接。# 处理 CUDA 的依赖关系（如 CUDA 路径、库文件等）,支持跨平台编译（Linux、Windows 等）'build_ext':BuildExtension}
)
运行 python setup.py install 即可安装扩展。
polynomial_activation.py
安装扩展后，就可以使用了，如下：
import torch
import torch.nn as nn
import time
import polynomial_cuda # This assumes you've built the extension using setup.pyclass CUDAPolynomialActivation(torch.autograd.Function):@staticmethoddef forward(ctx, x):return polynomial_cuda.polynomial_activation(x)@staticmethoddef backward(ctx, grad_output):# Implement backward pass if neededraise NotImplementedError("Backward pass not implemented")class PolynomialActivation(nn.Module):def __init__(self, implementation='pytorch'):super().__init__()self.implementation = implementationdef forward(self, x):if self.implementation == 'pytorch':return x**2 + x + 1elif self.implementation == 'cuda':return CUDAPolynomialActivation.apply(x)else:raise ValueError(f"Unknown implementation: {self.implementation}")# Benchmark function
def benchmark(func, x, name, num_runs=1000):start_time = time.time()for _ in range(num_runs):func(x)torch.cuda.synchronize()end_time = time.time()return f"{name}: {(end_time - start_time) / num_runs * 1000:.4f} ms"# Main function to run benchmarks
def main():torch.manual_seed(0)x = torch.randn(1000000, device='cuda')pytorch_activation = PolynomialActivation(implementation='pytorch').cuda()cuda_activation = PolynomialActivation(implementation='cuda').cuda()out = cuda_activation.forward(x)print(out)pytorch_time = benchmark(pytorch_activation, x, "PyTorch built-in")cuda_time = benchmark(cuda_activation, x, "CUDA extension")print(pytorch_time)print(cuda_time)if __name__ == "__main__":main()
这段代码展示了如何使用 PyTorch 自定义扩展（CUDA）来加速操作，并与 PyTorch 内置实现进行性能对比。这种方法在实际应用中非常有用，尤其是当你需要将某些操作的计算从 CPU 转移到 GPU 并且希望最大化性能时。
学习资料：
- PyTorch C++扩展教程
- PyTorch扩展开发文档
- 自定义CUDA操作教程
参考资料：https://github.com/Infatoshi/cuda-course/tree/master/09_PyTorch_Extensions