customPytorchCudaFunctor.md · 史朗清/MyBlog

# 自定义Pytorch CUDA算子
开始之前：

我还是想说，多种语言混合编程着实自讨苦吃，但是没办法，谁让用CUDA自定义Pytorch算子是CUDA一个重要应用呢。

## 自定义Pytorch算子的目的
Pytorch中已经提供了大量较为完备的张量和标量算子，但一些较特殊或者较新的算子在Pytorch中没有提供或者需要用已有的算子做复杂的组合，这样导致性能上的额外开销。因此在这种场景下，直接使用CUDA编写自定义算子是一种选择。当然如果性能不敏感或者重新编写提升有限的话，也没必要。

## 自定义Python扩展的方式
自定义Pytorch算子实际上是通过Python的C++扩展功能实现的。通过CUDA C/C++自定义Pytorch算子实际上是将编写好的CUDA C/C++编译为能够被Python解释器加载执行的动态链接库。这一点和MatLab/C++混合编程基本上一致，毕竟都是解释型语言去调用编译型语言。在Windows环境下，Python的动态链接库是.pyd格式，而Linux下则是.so格式。

与Matlab的入口函数```mexFunction```所包含的数据交互逻辑不同，Python的C++扩展不需要调用API进行复杂的数据格式转换，而是直接使用PYBIND将C++端的函数进行绑定，PYBIND支持C++端的许多特性，包括基本数据类型，STL，迭代器等，这可能与常用的CPython解释器就是由C/C++实现有关。示例代码如下：
```C++
int add(int j, int j){
    return i + j;
}

PYBIND11_MODULE(example, m){
    m.doc() = "pybind test";
    m.def("add", &add, "add 2 inputs");
}
```
Python以导入模块的形式使用C++扩展，通过```PYBIND11_MODULE```将C++函数绑定到Python函数。其中，```PYBIND11_MODULE```的参数```example```是该模块的名称，```m```是代表该模块的变量。```m.def```的第一个参数是该函数在Python中的调用名称，第二个参数是以函数引用的形式将函数的返回值返回Python，第三个参数是该函数的描述，可以缺省。```m.doc()```是该模块的描述，也可以缺省。示例：
```python
import example
ans = example.add(1, 2)
```

## 自定义Pytorch算子的开发环境
Pytorch实际上是Torch的一系列Python接口，Pytorch使用Torch作为后端引擎。而Torch本身是用CUDA和C++编写的一个张量计算引擎，其中包含了以```Tensor```类为核心的张量数据类，以及张量自动求导```auto_grad```等方法。

仅限于使用CUDA编写Pytorch自定义算子这个应用场景而言，conda中安装的torch运行环境就足够了。如果要在C++平台上使用部署Pytorch模型，则需要另外下载```libtorch```，其比conda中安装的torch运行环境多了一些内容，组成了完整的开发环境，即更加完整的头文件和动/静态库文件，如果下载的是Debug版本还会有调试所需的符号文件，Release版本则没有。
> 实际上libtorch并不复杂，如果Pytorch是Torch的Python接口，那么libtorch也可以理解为Cpptorch

Pytorch的自定义算子需要编译才能运行，通常有两种编译方式——通过setuptool直接编译，或者通过JIT方式即时编译。前者通过定义一个setup.py文件先将C++/CUDA源文件编译为pyd动态链接库，而后运行主函数调用。而JIT方式则较为方便，直接在主函数中调用Pytorch提供的即时编译模块编译并直接调用运行，缺点则是比前一种多出了编译的时间。这里使用的是JIT编译模式。通过调用```torch.utils.cpp_extension```模块中的```load```方法，指定编译所需的源文件，以及目标文件的名称，也就是这个模块在Python中的名称。```load```方法默认使用Ninja作为构建工具，因此需要先安装Ninja并将其路径添加到系统环境变量中。Windows平台可以直接下载exe版本，无需自行编译。Ninja在Windows平台默认使用的编译器是Visual Studio的cl.exe，因此需要额外将cl.exe所在的目录添加到系统环境变量中。以VS2019为例，路径为```vsroot\VC\Tools\MSVC\14.29.30133\bin\Hostx64\x64```

## 自定义Pytorch算子的过程
使用CUDA开发自定义Pytorch算子在C++端需要3个文件。以自定义一个两个张量相加的算子为例，add.cu定义了核函数的计算逻辑以及启动核函数的入口函数，add.cpp通过函数封装调用了入口函数，通过Torch的数据类型传入参数的方式，以及通过```PYBIND11_MODULE```将上述的封装函数绑定到Python，该封装函数就是最终在Python中调用的函数。add.cpp通过包含add.h使用add.cu中定义的入口函数。另有一个main.py作为顶层文件。上述4个文件都放在同一个目录下。示例如下：
```C++
//add2.cu
#include "cuda_runtime.h"
#include "add2.h"

__global__ void add2_kernel(float* c, const float* a, const float* b, int n){
    for(int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x){
        c[i] = a[i] + b[i];
    }
}

void launch_add2(float* c, const float* a, const float* b, int n){
    dim3 grid((n + 1023) / 1024);
    dim3 block(1024);
    add2_kernel<<<grid, block>>>(c, a, b, n);
}

//add2.h
void launch_add2(float* c, const float* a, const float* b, int n);

//add2.cpp
#include<torch/extension.h>
#include "add2.h"

void torch_launch_add2(torch::Tensor &c, const torch::Tensor &a, const torch::Tensor &b, int n){
    launch_add2((float*)c.data_ptr(), (const float*)a.data_ptr(), (const float*)b.data_ptr(), n);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("torch_launch_add2", &torch_launch_add2, "add2 kernel warpper");
}
```
在Python中通过调用```torch.utils.cpp_extension```模块中的```load```方法即时编译该模块并直接使用，示例如下：
```python
# main.py
import numpy as np
import torch
from torch.utils.cpp_extension import load
import time

def run_cuda():
    cuda_module.torch_launch_add2(cuda_c, a, b, n)
    return cuda_c

def show_time(func):
    times = []
    res = []
    for _ in range(10):
        torch.cuda.synchronize(device="cuda:0")
        start = time.time()
        r = func()
        torch.cuda.synchronize(device="cuda:0")
        stop = time.time()
        times.append((stop - start) * 1e6)
        res.append(r)
    return times, res

def run_torch():
    a + b
    return None

if __name__ == '__main__':
    cuda_module = load(name="add2",
                       sources=["add2.cpp", "add2.cu"],
                       verbose=True)
    n = 1024 * 1024
    a = torch.rand(n, device="cuda:0")
    b = torch.rand(n, device="cuda:0")
    cuda_c = torch.rand(n, device="cuda:0")
    cuda_time, _ = show_time(run_cuda)
    print("Cuda time: {:.3f}us".format(np.mean(cuda_time)))
```
至此，一个简单的CUDA自定义Pytorch算子构建完毕。在终端直接运行main.py，输出如下：
```shell
$ D:/Anaconda/envs/torch-gpu/python.exe d:/SLQ/Learn/Cuda/CudaTorchTest/main.py
Using C:\Users\admin\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu102 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file C:\Users\admin\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu102\add2\build.ninja...
Building extension module add2...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module add2...
Cuda time: 197.363us
```

史朗清/MyBlog

简介

发行版

贡献者

近期动态

史朗清/MyBlog .gitee-modal { width: 500px !important; }

简介

发行版

贡献者

近期动态

搜索帮助

史朗清/MyBlog