基于Pytorch，如何用CUDA自己写算子？（一）

原创已于 2026-07-29 21:12:31 修改 · 216 阅读

本内容遵循CC 4.0 BY-SA版权协议

于 2026-07-23 14:41:59 首次发布

1、简介

大家都知道，深度学习通过梯度反传去优化网络的参数。Pytorch作为具有代表性的深度学习库，对前向（Forward）和梯度反传（Backward）都封装的非常好。在日常的编程中，只需要调用Pytorch封装好的算子，就可以完成网络的设计和训练。

但是，在一些任务中，有些算法没法用Pytroch封装好的算子去实现，需要自己实现。这个时候，则需要自己去写CUDA代码，自己去写Forward和Backward的代码。这个博客就是基于简单的案例去讲解，如何基于Pytorch，用CUDA自己写算子。

在我们自己写forward和backward时，pytroch对于cuda上自己写的代码并不关心。它只关心在forward时，输入和输出是什么。在backward的时候，梯度又是多少。我们需要对计算的正确性负责。

在学习这个博客之前，需要具备一定的CUDA编程的基础。

2、简单案例

2.1 案例介绍

输入：一个张量x；

算子： $y=x^{2}$

目的：基于Pytorch，自己实现该算子，完成forward和backward

这个案例虽然非常简单，非常的好理解，非常合适用于说明写CUDA和封装的过程。

假设我们的输入张量 $x=[1.0,2.0,-3.0,4.0]$ ，那么经过该算子，得到的结果应该是 $y=[1.0,2.0,9.0,16.0]$ 。

对于该算子，偏导为 $\frac{\partial y}{\partial x} = 2x$ ，则反传到x的梯度应该是 $x_{grad} = [2.0,4.0,-6.0,8.0]$ 。

接下来，讲解如何实现代码，并对代码的意思进行简单的讲解。

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')

标签

#pytorch #人工智能 #python