【CUDA】Pytorch_Extensions

【CUDA】Pytorch_Extensions

为什么要开发CUDA扩展？

当我们在PyTorch中实现自定义算子时，通常有两种选择：

1. 使用纯Python实现（简单但效率低）
2. 使用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中调用 polynomial_activation。

这三个部分构成完整的polynomial_cuda.cu:

#include 
#include 
#include 

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.py

class CUDAPolynomialActivation(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return polynomial_cuda.polynomial_activation(x)

    @staticmethod
    def backward(ctx, grad_output):
        # Implement backward pass if needed
        raise NotImplementedError("Backward pass not implemented")

class PolynomialActivation(nn.Module):
    def __init__(self, implementation='pytorch'):
        super().__init__()
        self.implementation = implementation

    def forward(self, x):
        if self.implementation == 'pytorch':
            return x**2 + x + 1
        elif 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