1、简介
大家都知道，深度学习通过梯度反传去优化网络的参数。Pytorch作为具有代表性的深度学习库，对前向（Forward）和梯度反传（Backward）都封装的非常好。在日常的编程中，只需要调用Pytorch封装好的算子，就可以完成网络的设计和训练。
但是，在一些任务中，有些算法没法用Pytroch封装好的算子去实现，需要自己实现。这个时候，则需要自己去写CUDA代码，自己去写Forward和Backward的代码。这个博客就是基于简单的案例去讲解，如何基于Pytorch，用CUDA自己写算子。
在学习这个博客之前，需要具备一定的CUDA编程的基础。
2、简单案例
2.1 案例介绍
输入：一个张量x；
算子：
目的：基于Pytorch，自己实现该算子，完成forward和backward
这个案例虽然非常简单，非常的好理解，非常合适用于说明写CUDA和封装的过程。
假设我们的输入张量，那么经过该算子，得到的结果应该是
。
对于该算子，偏导为，则反传到x的梯度应该是
。
接下来，讲解如何实现代码，并对代码的意思进行简单的讲解。
2.2 案例代码
2.2.1 整体介绍
整个工程结构有以下几个文件：
1、test.py
作用：主函数，用于做测试
2、function_cuda_kernel.cu
作用：编写运行在GPU上的cuda代码，forward和backward的实现就是在这个文件里面
3、function_cuda.cpp
作用：将CUDA函数暴露给Python，python不能直接调用.cu中的C++函数，需要使用pybind11建立绑定关系
4、setup.py
作用：编译.cpp 和 .cu 文件，生成Python可以导入的扩展模块，生成 .so文件，供python调用
5、defined_functions.py
作用：把编译好的 CUDA 前向和反向函数接入 PyTorch 的自动求导系统。使用torch.autograd.Function 组织forward/backward
接下来，讲解每个函数的实现
2.2.2 test.py 文件
import torch
from defined_functions import AllFunctions
def main():
x = torch.tensor( [1.0, 2.0, -3.0, 4.0]).cuda()
x.requires_grad_(True)
function = AllFunctions()
y = function.squareApply(x)
print("前向的结果：", y)
loss = y.sum()
# 自动触发自定义 CUDA backward
loss.backward()
print("\n反向梯度 x.grad：")
print(x.grad)
if __name__ == "__main__":
main()这个main函数代码比较简单，就是定义了一个张量x，然后调用封装好了的前向函数。同时，自动启动梯度反传的计算，输出x的梯度。
在上面代码中，y = function.squareApply(x) 就是调用forward函数；loss.backward() 则会自己逐步的做梯度反传，并自动执行我们写好的backward函数。
2.2.3 defined_functions.py 文件
import torch
import square_cuda_ext
class AllFunctions:
def squareApply(self, inputTensor: torch.Tensor):
square = SquareFuction()
return square.apply(inputTensor)
class SquareFuction(torch.autograd.Function):
@staticmethod
def forward(ctx, inputTensor: torch.Tensor):
if not inputTensor.is_cuda:
raise RuntimeError("input_tensor must be a CUDA tensor")
input_contiguous = inputTensor.contiguous()
ctx.save_for_backward(input_contiguous)
output = square_cuda_ext.square_forward(input_contiguous)
return output
@staticmethod
def backward(ctx, grad_output):
(inputTensor,) = ctx.saved_tensors
grad_output_contiguous = grad_output.contiguous()
# 调用自己编写的 CUDA backward
grad_input = square_cuda_ext.square_backward(
grad_output_contiguous,
inputTensor,
)
return grad_input在这段代码里面，自己定义了一个类SquareFuction。这个类里面写了一个forward，一个backward，在这两个函数里面，分别调用了用cuda写的forward和backward计算过程。
SquareFuction类的基类是torch.autograd.Function，意思就是这个类是自定义前向和后向是怎么计算的，是pytorch提供的基类。
@staticmethod 表示函数是静态函数，只定义计算规则，不创建普通实例。通过ctx是保存数据。
ctx 是上下文信息，用于保存前向阶段的信息，并在反向阶段取回，用于计算梯度。是Pytorch的torch.autograd.Function专门设计的参数，用于连接前向传播和反向传播。
在forward和backward函数中，分别调用了 square_cuda_ext.square_forward(input_contiguous) 和grad_input = square_cuda_ext.square_backward(grad_output_contiguous,inputTensor) ，也就是我们自己封装的前向和后向函数。
2.2.4 function_cuda.cpp 文件
#include <torch/extension.h>
torch::Tensor square_cuda_forward(torch::Tensor input);
torch::Tensor square_cuda_backward(torch::Tensor grad_output, torch::Tensor input);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("square_forward", &square_cuda_forward);
m.def("square_backward", &square_cuda_backward);
}这个代码比较简单，首先是声明了两个cuda 的函数，这两个函数在function_cuda_kernel.cu文件中实现的。然后，利用PYBIND11_MODULE，把 C++ 函数暴露给Python。
在python中调用的时候，就只用调用square_forward和square_backward函数即可，实际执行的是.cu中的square_cuda_forward 和 square_cuda_backward函数。
2.2.5 function_cuda_kernel.cu 文件
这个文件是实现算子的核心文件，代码如下：
#include <torch/extension.h>
__global__ void square_forward_kernel(float* input, float* output, int64_t num_elements)
{
int64_t index = static_cast<int64_t> (blockIdx.x * blockDim.x + threadIdx.x);
if (index < num_elements)
{
float x = input[index];
output[index] = x * x;
}
}
__global__ void square_backward_kernel(float* grad_output, float* input, float* grad_input, int64_t num_elements)
{
int64_t index = static_cast<int64_t> (blockIdx.x * blockDim.x + threadIdx.x);
if (index < num_elements)
{
float x = input[index];
float grad_out = grad_output[index];
grad_input[index] = grad_out * 2.0f * x;
}
}
torch::Tensor square_cuda_forward(torch::Tensor input)
{
int64_t num_elements = input.numel();
int threads_per_block = 256;
int blocks = static_cast<int> ( (num_elements + threads_per_block - 1) / threads_per_block);
torch::Tensor output = torch::empty_like(input);
square_forward_kernel<<<blocks, threads_per_block>>>(input.data_ptr<float>(), output.data_ptr<float>(), num_elements);
return output;
}
torch::Tensor square_cuda_backward(torch::Tensor grad_output, torch::Tensor input)
{
torch::Tensor grad_input = torch::empty_like(input);
int64_t num_elements = input.numel();
int threads_per_block = 256;
int blocks = static_cast<int> ( (num_elements + threads_per_block - 1) / threads_per_block);
square_backward_kernel<<<blocks, threads_per_block>>>(grad_output.data_ptr<float>(), input.data_ptr<float>(), grad_input.data_ptr<float>(), num_elements);
return grad_input;
}这个就是用cuda实现了简单的平方计算，以及平方的梯度计算，没啥好多说的。
2.2.6 setup.py文件
这个文件的代码如下：
GCC_PATH = "/usr/bin/gcc-11"
GXX_PATH = "/usr/bin/g++-11"
from setuptools import setup
from torch.utils.cpp_extension import (
BuildExtension,
CUDAExtension,
)
setup(
name="function_project",
ext_modules=[
CUDAExtension(
name="square_cuda_ext",
sources=[
"function_cuda.cpp",
"function_cuda_kernel.cu",
],
extra_compile_args={
"cxx": [
"-O3",
],
"nvcc": [
"-O3",
],
},
)
],
cmdclass={
"build_ext": BuildExtension
},
)核心的是 setup() 函数，在里面定义了工程名叫什么（name = "function_project"），这个是定义了最后编译出来的.so文件的名字。
最核心的是ext_modules，这个里面的name决定了最后Python扩展模块的名字，也就是用的时候，是import square_cuda_ext、sources表示的是原文件。extra_compile_args 这个里面是一些编译的设置。
2.3 运行
先输入命令，编译我们写的算子：
python setup.py build_ext --inplace
然后，再直接
python test.py
一切顺利的话，就可以得到如下结果：
前向的结果： tensor( 1., 4., 9., 16., device='cuda:0', grad_fn=<SquareFuctionBackward>)
反向梯度 x.grad：
tensor( 2., 4., -6., 8., device='cuda:0')