前言

坑啊坑！在小小的工位上坑啊坑

必须要把这种巨坑的编程风格公开处刑，这篇论文：Multiple Video Frame Interpolation via Enhanced Deformable Separable Convolution，真把我坑惨了。先是以为cupy不兼容，遂学CUDA+Torch；复以为驱动问题，用上docker，最后以为是硬件问题，搞个压力测试，通通失败！枉费我与师姐交流如此之久也，若不是今天一早忽然想到设备的问题，加上师姐一提点，这破代码不知道还要接着坑多少人呢！

这位代码的祖宗，也是祖宗之法不可变也，经典屎山；论文对它是一点不改，一动不动直接照搬，有设备问题的Bug也不解决，坑了我整整两天的时间！对它的厌恶是无以复加！

分割线以下是正常部分。

Pytorch是当前最为流行的深度学习框架，其与CUDA的深度兼容性也极大地方便了使用者在不需要了解任何GPU底层架构和编程知识的前提下写出高性能的解决方案，可谓是前无古人，后…来者也。因此当下所有的网络模型、甚至是新的网络框架（oneflow、spikingjelly、slayer等）都在积极兼容torch。

但是，再好用的框架也有其局限性，Pytorch也并不能一站式实现我的全部模型构想 ，对于如上述论文所提出的“弹性可分离卷积”代码，其对每一个像素位置都将产生一个独有的偏移，该偏移将作用于卷积的位置上，因此直接使用torch自带API将无法完成该算法的构建。论文采用的方法是CUDA+cupy继承自祖传代码，实际上这种方法简直是多此一举，cupy就是为了方便不用CUDA的，你都用上CUDA了，还不如直接写CUDA呢。本文将介绍我在使用CUDA+Torch算子方式复现该论文过程中所使用的一般方法及某个奆坑。

温馨提示：CUDA编程，一定要深度掌握主机与设备、设备和设备之间的关系，尤其是多卡服务器

安装说明

除非一定要在C++环境下写单测并使用cmake完成CUDA算子的编译构建，否则没有必要单独安装libtorch。如果一定要安装，请在官网上仔细查看自己的CUDA版本和驱动，然后安装C++版本的稳定发布版本，如下图：
安装libtorch
解压至需要的目录并将其添加至环境变量中，然后在CMakeLists.txt中添加如下代码：

1
2
3
4
set(PYTHON_DIR /path/to/your/Python.h)
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
include_directories(${PROJECT_NAME} PUBLIC_HEADER ${PYTHON_DIR})

当然也可以手工添加该路径变量，如下：

1
2
3
4
5
set(PYTHON_DIR /path/to/your/Python.h)
set(TORCH_DIR /path/to/your/libtorch)
include_directories(${PROJECT_NAME} PUBLIC_HEADER ${TORCH_DIR}/include)
include_directories(${PROJECT_NAME} PUBLIC_HEADER ${PYTHON_DIR})
include_directories(${PROJECT_NAME} PUBLIC_HEADER ${TORCH_DIR}/include/torch/csrc/api/include)

如果仅需要编写CUDA算子，请注意

一定要安装CUDA版本的Pytorch

具体可以使用conda或者pip查看，后缀后带有+cuXXX的即为CUDA版本。

编写CUDA算子

本文默认读者都有一定的CUDA编程基础（虽然作者本人的CUDA编程基础很差劲），因此这里不再赘述CUDA语法知识。根据官方教程，编写好的CUDA算子可以采用三种方式内嵌至Pytorch中，分别是使用即时编译JIT、使用setuptools构建为第三方库或者使用cmake构建为第三方可执行文件。

这三种的好处与缺点大概陈列如下：

方式 优点 缺点
JIT 无需额外的构建步骤，直接在Python中调用 速度慢，高度依赖python环境中的torch等库版本
setuptools 预先构建，不用每次都编译，可以直接在Python中调用 需要额外的构建步骤，需要额外的构建工具，同时也高度依赖torch库版本
cmake 速度快，不依赖Python环境 需要额外的构建步骤，需要额外的构建工具，不依赖Python环境

就我个人而言，使用JIT是最方便的，但是别人使用它最不方便（除非环境这些都完全一样），我倾向于使用cmake，但是它的构建比较麻烦，目前很纠结，因此这两种方式都有所涉足，而至于setuptools，它基本上就类似于这两种的过渡。

本文将以构建一个矢量加法算子为例，主要参考自知乎博客，对该博客有所拓展（考虑了多卡运行的情况，这也是当前几乎所有博客都没有提到的一个点，谁让他们ctrl C+ctrl V呢）。

CUDA编程部分

首先编写CUDA代码，就与正常的CUDA程序一致，文件目录结构如下：

1
2
3
4
5
6
7
8
9
10
11
12
.
|-- include
|   `-- my_add.cuh
|
|-- src
|   |-- my_add.cpp
|   `-- my_add_kernel.cu
|
|-- test_my_add.py
|-- test_setup.py
|-- setup.py
|-- CMakeLists.txt

其中一定包括头文件目录和源文件目录，目录的名字倒是不一定非得叫include和src，不过编程习惯倒是决定了大多数情况下它都能适用。

在my_add.cuh中，我们定义了一个简单的矢量加法算子：

1
2
3
4
5
6
7
#pragma once
void launch_my_add( const int n, 
                    const float *input_a, 
                    const float *input_b, 
                    float *output_c, 
                    const int cuda_index
                  );

前四个参数不再赘述，请注意最后一个参数，该参数将用于指定使用的显卡序号，用于在多卡运行时避免一个非常诡异的错误。

在my_add_kernel.cu中，我们定义了一个简单的矢量加法算子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
#include "my_add.cuh"
#include <torch/torch.h>
/*算子的具体实现*/
__global__ void my_add_kernel(const int n, 
                              const float *vec_a, 
                              const float *vec_b, 
                              float *vec_out)
{
  // int tid = blockIdx.x * blockDim.x + threadIdx.x;
  for (int tid = blockIdx.x * blockDim.x + threadIdx.x;
       tid < n; tid += gridDim.x * blockDim.x)
  {
    vec_out[tid] = vec_a[tid] + vec_b[tid];
  }
}
void launch_my_add(const int n,
                   const float *a,
                   const float *b,
                   float *c, 
                   const int cuda_index)
{
  cudaSetDevice(cuda_index); // 设置当前使用的显卡
  dim3 grid((n + 1023) / 1024);
  dim3 block(1024);
  my_add_kernel<<<grid, block>>>(n, a, b, c);
}

此处的编程使用到了CUDA编程中较为重要的一个原则，即核函数封装，核 函数本身不能成为成员函数，但是其封装后的函数将如同一个普通的函数，可以与C++其他类型的函数一样被调用。封装函数的第一句代码是指定当前核函数的运行设备的，在别的博客中都没有提到过这一点（或许是我太菜了吧），封装的核 函数中不能够显式出现设备相关的函数（因为核函数不能调用主机函数）。

在my_add.cpp中，我们将封装好的算子使用Pybind11接口程序构建为可由Python程序调用的库（仅适用于JIT和setuptools方式，cmake方式不适用Pybind11）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
#include <torch/extension.h>
#include "my_add.cuh"
#include <iostream>
/*算子的封装*/
void torch_launch_my_add( const int n, 
                          const torch::Tensor &a,
                          const torch::Tensor &b,
                          torch::Tensor &c, 
                          const c10::Device device)
{
  launch_my_add(n, (const float *)a.data_ptr(), (const float *)b.data_ptr(), (float *)c.data_ptr(), device.index());
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
  m.def("torch_launch_my_add",
        &torch_launch_my_add,
        "description of the function, you can determine it by yourself");
}

到此准备性质的工作已经全部就绪。

使用JIT方式构建

在test_my_add.py中，我们可以直接调用torch_launch_my_add函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch
from torch.utils.cpp_extension import load
import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "1"
# os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
cuda_module = load(
    name="my_add",
    extra_include_paths=["include"],
    sources=["src/my_add.cpp", "src/my_add_kernel.cu"],
    verbose=True,
)

# c = a + b (shape: [n])
n = 10
device=torch.device("cuda:5") # 请选定你所使用的设备序号，如果是单卡服务器就是"cuda:0"或者"cuda"
a = torch.ones(n, device=device)
b = torch.ones(n, device=device) * 5
cuda_c = torch.ones(n, device=device) * 2


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


def run_torch():
    # return None to avoid intermediate GPU memory application
    # for accurate time statistics
    print(a + b)
    return None


run_torch()  # 一个是直接跑torch
print(run_cuda())  # 一个是跑cuda算子

直接如往常一般运行该程序，在输出信息中将会查看到程序正在调用g++和nvcc对库文件进行编译，完成编译后将出现Loading extension module my_add...字样，说明编译正确后加载该库文件。该方法不愧是最简易的方法，不需要额外的构建步骤，只要保证所构建的环境是一致的，即可完成即时编译。

使用setuptools方式构建

在test_setup.py中，我们可以直接调用torch_launch_my_add函数：

1
2
3
4
5
6
7
8
9
import torch
import my_add

n = 10
a = torch.ones(n, device="cuda:5")
b = torch.ones(n, device="cuda:5") * 5
cuda_c = torch.ones(n, device="cuda:5") * 2
my_add.torch_launch_my_add(n, a, b, cuda_c, device=torch.device("cuda:5"))
print(cuda_c)

本方式构建的动态链接库类似于cmake编译得到的动态链接库，一定程度上能够脱离环境独立运行。

使用cmake方式构建

这一步较为复杂，我目前仅使用它来作为调试过程中的语法速查和初步结果查看，尚不能用于直接引入库。对于编写完成的工程文件，可以参考以下的CMakeLists.txt文件：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
cmake_minimum_required(VERSION 3.15)

project(dsconv LANGUAGES CXX CUDA)

set(CMAKE_CXX_STANDARD 17) #指定编译器的版本，这一步不可以省略
set(CMAKE_BUILD_TYPE Debug)
set(CMAKE_C_FLAGS_DEBUG "$ENV{CXXFLAGS} -O0 -Wall -g")
set(CMAKE_CXX_FLAGS_DEBUG "$ENV{CXXFLAGS} -O0 -Wall -g")
#正则式匹配全部的源文件，包括C、C++、CUDA
file(GLOB LIBRARY "src/*")

set(LIBRARY_NAME my_lib)
set(CUDA_DIR /usr/local/cuda-11.8)  #显式增加CUDA路径，仅适用于有显式安装CUDA的情况，对于使用conda安装的情况，直接调用其路径不能得到正确的结果，以下的torch和Python也是同样的。
set(TORCH_DIR /disk527/Datadisk/xdy_cbf/software/libtorch)
set(PYTHON_DIR /disk527/Datadisk/xdy_cbf/software/miniconda3/envs/ldm_t/include/python3.9/)
add_library(${LIBRARY_NAME} SHARED ${LIBRARY}) #编译为动态链接库，注意参数SHARED

include_directories(${LIBRARY_NAME}  ${CMAKE_CURRENT_SOURCE_DIR}/include)
include_directories(${LIBRARY_NAME}  ${CUDA_DIR}/include)
include_directories(${LIBRARY_NAME}  ${TORCH_DIR}/include)
include_directories(${LIBRARY_NAME}  ${PYTHON_DIR})
include_directories(${LIBRARY_NAME}  ${TORCH_DIR}/include/torch/csrc/api/include)

编译后将在build目录下生成my_lib.so文件，理论上可以使用ctypes库调用该库文件，但是我并没有成功。