像教女朋友一样教你用Cuda实现PyTorch算子

点击下方卡片，关注“AI前沿速递”公众号

各种重磅干货，第一时间送达

前段时间一直在做算子上的优化加速工作，在和其他同学的讨论中发现用Cuda编写算子存在一定的门槛。虽然知乎上有很多优秀的教学指南、PyTorch官方也给出了tutorial（具体地址会放在文章末尾），但是对于每个环节的介绍与踩坑点似乎没有详实的说明。

结合我当时入门踩坑的惨痛经验，一个简单明了的demo能够大大减小上手的时间成本。所以我在这里以数组求和（下称sum_single）、两数组相加（下称sum_double） 为例，详细介绍一下用Cuda实现PyTorch算子的完整框架，具体的代码详见CudaDemo

来源知乎：雅痞（已授权）   链接：https://zhuanlan.zhihu.com/p/595851188

框架结构

├── ops
│   ├── __init__.py
│   ├── ops_py
│   │   ├── __init__.py
│   │   └── sum.py
│   └── src
│       ├── reduce_sum
│       │   ├── sum.cpp
│       │   └── sum_cuda.cu
│       └── sum_two_arrays
│           ├── two_sum.cpp
│           └── two_sum_cuda.cu
├── README.md
├── setup.py
└── test_ops.py

demo结构如上，其中

• ops/src/是Cuda/C++代码
• setup.py是编译算子的配置文件
• ops/ops_py/是用PyTorch包装的算子函数
• test_ops.py 是调用算子的测试文件

Cuda/C++

对于一个算子实现，需要用到.cu(Cuda)编写核函数、.cpp(C++)编写包装函数并调用PYBIND11_MODULE对算子进行封装。

注意：Cuda文件和Cpp文件不能同名！！！否则编译不通过！！！

我们这里以src/sum_two_arrays/为例进行解释

// src/sum_two_arrays/two_sum_cuda.cu
#include <cstdio>

#define THREADS_PER_BLOCK 256
#define WARP_SIZE 32
#define DIVUP(m, n) ((m + n - 1) / n)


__global__ void two_sum_kernel(const float* a, const float* b, float * c, int n){
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n){
        c[idx] = a[idx] + b[idx];
    }
}


void two_sum_launcher(const float* a, const float* b, float* c, int n){
    dim3 blockSize(DIVUP(n, THREADS_PER_BLOCK));
    dim3 threadSize(THREADS_PER_BLOCK);
    two_sum_kernel<<<blockSize, threadSize>>>(a, b, c, n);
}

这里的关键是two_sum_kernel这一核函数实现数组相加功能。下面的two_sum_launcher函数负责分配线程块并调用核函数。

// src/sum_two_arrays/two_sum.cpp
#include <torch/extension.h>
#include <torch/serialize/tensor.h>

#define CHECK_CUDA(x) \
  TORCH_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
#define CHECK_CONTIGUOUS(x) \
  TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_INPUT(x) \
  CHECK_CUDA(x);       \
  CHECK_CONTIGUOUS(x)


void two_sum_launcher(const float* a, const float* b, float* c, int n);


void two_sum_gpu(at::Tensor a_tensor, at::Tensor b_tensor, at::Tensor c_tensor){
    CHECK_INPUT(a_tensor);
    CHECK_INPUT(b_tensor);
    CHECK_INPUT(c_tensor);

    const float* a = a_tensor.data_ptr<float>();
    const float* b = b_tensor.data_ptr<float>();
    float* c = c_tensor.data_ptr<float>();
    int n = a_tensor.size(0);
    two_sum_launcher(a, b, c, n);
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &two_sum_gpu, "sum two arrays (CUDA)");
}

在C++文件中实现算子的封装，文件开头的宏定义函数是为了保证传入的向量在cuda上（CHECK_CUDA）、传入的向量中元素地址连续（CHECK_CONTIGUOUS）。two_sum_launcher是对cuda文件中的声明。

two_sum_gpu是与Python的接口，传入的参数是PyTorch中的Tensor。在这一部分需要对Tensor做CHECK检验（可选），并通过.data_ptr得到Tensor变量的指针。对于Tensor在C++中的使用可查阅[2]。

最后PYBIND11_MODULE的作用是对整个算子进行封装，能够通过Python调用C++函数[3]。对于自定义的其他算子，只用改动m.def()中的三个参数

• "forward"：算子的方法名，假如算子的整个模块命名为sum_double，则在Python中通过sum_double.forward调用该算子
• &two_sum_gpu：进行绑定的函数，这里根据自己实现的不同函数进行更改
• "sum two arrays (CUDA)"：算子注释，在Python端调用help(sum_double.forward) 时会出现

可能有人会疑惑为什么要把算子和模块分开。假如整个sum_double有许多不同的功能，我就可以在一个模块中绑定多个算子，具体只用在PYBIND11_MODULE中写入多个m.def()，再通过sum_double.xxx调用不同的算子

setup.py编译配置

在整个项目的根目录新建setup.py文件配置编译信息，利用setuptools对算子打包

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

setup(
    name='CudaDemo',
    packages=find_packages(),
    version='0.1.0',
    author='xxx',
    ext_modules=[
        CUDAExtension(
            'sum_single', # operator name
            ['./ops/src/reduce_sum/sum.cpp',
             './ops/src/reduce_sum/sum_cuda.cu',]
        ),
        CUDAExtension(
            'sum_double',
            ['./ops/src/sum_two_arrays/two_sum.cpp',
             './ops/src/sum_two_arrays/two_sum_cuda.cu',]
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

文件中需要进行改动的有

• name：包名
• version：包版本号
• author：作者名称
• ext_modules：编译C/C++扩展，list类型，每个元素为一个模块的相关信息（这里的模块在讲Cuda/C++这一块的末尾有提到，一个模块可以含有多个具体的算子）

CUDAExtension

在ext_modules中采用CUDAExtension指明Cuda/C++的文件路径 ，其中第一个参数为对应模块的名字，第二个参数为包含所有文件路径的列表。

这里的模块名和Cuda/C++中m.def()定义的算子名共同决定了调用算子的方式。例如两数组相加的模块名是sum_double、算子方法名是forward， 所以在Python中调用该算子的方式为sum_double.forward()。

值得一提的是packages的值为list[str]，表示本地需要打包的package。这里find_packages()是找出本地所有的package。当然我们打包只用考虑ops/src/中的文件，所以packages=['ops/src']也能正常编译，不过为了方便还是采用find_packages()。

PyTorch包装

为了让自定义算子能够正常正向传播、反向传播，我们需要继承torch.autograd.Function进行算子包装。我们这里以sum_double为例进行介绍

#ops/ops_py/sum.py
import torch
from torch.autograd import Function
import sum_double

class SumDouble(Function):

    @staticmethod
    def forward(ctx, array1, array2):
        """sum_double function forward.
        Args:
            array1 (torch.Tensor): [n,]
            array2 (torch.Tensor): [n,]
        
        Returns:
            ans (torch.Tensor): [n,]
        """
        array1 = array1.float()
        array2 = array2.float()
        ans = array1.new_zeros(array1.shape)
        sum_double.forward(array1.contiguous(), array2.contiguous(), ans)

        # ctx.mark_non_differentiable(ans) # if the function is no need for backpropogation

        return ans

    @staticmethod
    def backward(ctx, g_out):
        # return None, None   # if the function is no need for backpropogation

        g_in1 = g_out.clone()
        g_in2 = g_out.clone()
        return g_in1, g_in2


sum_double_op = SumDouble.apply

文件开头import sum_double就是导入的setup.py中定义的模块名。

自定义的torch.autograd.Function类型要实现forward、backward函数，并声明为静态成员函数。

forward

前向传播的前半部分就是正常传入Tensor进入接口，如果传入向量在之前的代码里是索引出来的很可能非连续，所以建议在传入算子的时候使其连续。

如果算子不需要考虑反向传播，可以用ctx.mark_non_differentiable(ans) 将函数的输出标记不需要微分

backward

backward的输入对应forward的输出，输出对应forward的输入。例如这里backward的输入g_out对应forward输出ans，backward的输出g_in1, g_in2对应内联代码片forward输入array1, array2。

如果算子不需要考虑反向传播，则直接return None, None。否则就按照对应输入变量的梯度进行计算。

值得注意的是，如果反向传播需要用到forward的信息，可以用ctx进行记录存取。例如对一个数组求和，则反向传播的梯度为原数组长度的向量。就可以在forward中用ctx.shape=array.shape[0]记录输入数组长度，并在backward中通过n=ctx.shape进行读取。

如果存取的是Tensor则建议使用save_for_backward(x, y, z, ...)存储向量，并用x, y, z, ...=ctx.saved_tensors取向量，而不是直接用ctx。

To prevent incorrect gradients and memory leaks, and enable the application of saved tensor hooks.注：save_for_backward()只能存向量，标量用ctx直接存取。

最后用sum_double_op = SumDouble.apply获取最终的函数形式。

init.py

为了在外部调用包装好的PyTorch函数，通过ops/ops_py/__init__.py声明

from .sum import sum_single_op, sum_double_op
__all__ = ['sum_single_op', 'sum_double_op']

ops/__init__.py

from .ops_py import *

Build & Test

提前安装好PyTorch环境，并在demo的根目录下pip install -e .

通过python test_ops.py测试结果，没问题的情况应输出：

Average time cost of sum_single is 2.8257 ms
Average time cost of sum_double is 0.1128 ms

如果无法编译可能是没有将nvcc加入环境变量，ls /usr/local/看看是否有cuda文件夹。例如我这里是cuda-11.6文件夹，则进入~/.bashrc在文件末尾加入

export PATH=$PATH:/usr/local/cuda-11.6/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.6/lib64