利用torch.utils.cpp_extension库编写支持CUDA的扩展模块

Torch提供了torch.utils.cpp_extension库，可以帮助我们编写支持CUDA的扩展模块。这个库可以方便地将C++代码编译成PyTorch的扩展模块，并且支持使用CUDA加速。

下面是一个使用例子，展示了如何使用torch.utils.cpp_extension编写一个简单的扩展模块，并使用CUDA加速。

首先，我们需要准备两个文件：add_cuda.cpp和add_cuda_kernel.cu。add_cuda.cpp文件是我们的C++源码文件，add_cuda_kernel.cu文件是我们的CUDA源码文件。

add_cuda.cpp的内容如下：

#include <torch/extension.h>

// 定义C++前向函数
torch::Tensor add_forward(torch::Tensor input);

// 定义C++反向函数
torch::Tensor add_backward(torch::Tensor gradOutput);

// 绑定C++函数到PyTorch接口
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &add_forward, "Add Forward");
    m.def("backward", &add_backward, "Add Backward");
}

add_cuda_kernel.cu的内容如下：

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

// CUDA前向函数
#define THREADS_PER_BLOCK 1024

__global__ void add_forward_kernel(float *input, float *output, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        output[idx] = input[idx] + 1.0;
    }
}

torch::Tensor add_forward(torch::Tensor input) {
    int size = input.numel();
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);

    torch::Tensor output = torch::empty({size}, options);
    float *input_data = input.data_ptr<float>();
    float *output_data = output.data_ptr<float>();

    const int blocks = (size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    add_forward_kernel<<<blocks, THREADS_PER_BLOCK>>>(input_data, output_data, size);

    return output;
}

// CUDA反向函数
__global__ void add_backward_kernel(float *gradOutput, float *gradInput, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        gradInput[idx] = gradOutput[idx];
    }
}

torch::Tensor add_backward(torch::Tensor gradOutput) {
    int size = gradOutput.numel();
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);

    torch::Tensor gradInput = torch::empty({size}, options);
    float *gradOutput_data = gradOutput.data_ptr<float>();
    float *gradInput_data = gradInput.data_ptr<float>();

    const int blocks = (size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    add_backward_kernel<<<blocks, THREADS_PER_BLOCK>>>(gradOutput_data, gradInput_data, size);

    return gradInput;
}

接下来，我们可以使用torch.utils.cpp_extension编译这两个文件，生成扩展模块。编译的代码如下：

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

ext_modules = [
    CUDAExtension('add_cuda', [
        'add_cuda.cpp',
        'add_cuda_kernel.cu',
    ]),
]

setup(
    name='add_cuda',
    ext_modules=ext_modules,
    cmdclass={'build_ext': BuildExtension}
)

最后，我们可以像使用其他Python扩展模块一样，直接导入并使用这个扩展模块。代码如下：

import torch
from add_cuda import forward, backward

# 创建输入Tensor
input = torch.tensor([1.0, 2.0, 3.0, 4.0], device='cuda')

# 使用扩展模块进行前向计算
output = forward(input)
print(output)

# 使用扩展模块进行反向计算
gradOutput = torch.tensor([1.0, 1.0, 1.0, 1.0], device='cuda')
gradInput = backward(gradOutput)
print(gradInput)

这个例子演示了如何使用torch.utils.cpp_extension编写一个支持CUDA加速的扩展模块，并在PyTorch中使用它。