【CUDA和C++混合编译实现Python扩展】

---

源码

在深度学习任务中，如果使用PyTorch实现神经网络中的一些功能模块，其运行速度是比较慢的；为了能够加快这些模块的运行速度，通常会使用CUDA或者Cython完成这些功能模块的逻辑编写，并编译成动态链接库，之后在Python文件中实现调用，并封装成相应的功能模块，实现网络的前向和反向传播，从而大大提高网络整体的执行速度。

本文主要是想梳理一下实现上述功能的整体流程。

1. 如何使用CUDA完成某一功能模块的编写；
2. 编写好的CUDA扩展如何能够和Python进行连接/交互 (后面会提到，这一连接主要是靠C++文件)；
3. 如何将写好的CUDA代码，C++代码进行混合编译；
4. 编写好的CUDA扩展，如何被Python封装成网络的一个模块，并实现网络的前向和反向传播；

1.文件整体结构

想要实现CUDA和C++的混合编译，并能够被Python调用，执行相应的功能主要需要4个文件，如下所示。主要参考Yulong：扩展Pytorch：利用CUDA实现算子（二）

（一） CPP文件
在CPP文件中主要完成4个功能。

1）对要调用的CUDA函数进行声明，这里提到的CUDA函数的定义在cu文件中；
2）完成一些宏定义，这些宏定义主要作用是为了进行数据的形状检查或张量检查，防止运行过程中出现问题；
3）定义Python调用扩展功能的函数接口，也就是CPP文件中定义的C++函数；
4）PYBIND11_MODULE(TORCH_EXTENSION_NAME)函数封装Python调用扩展函数的C++接口。

（二）CU文件
在CU文件中一般包括2个文件。

1）对CPP文件中的声明的CUDA函数进行定义；在该函数中一般会调用真正实现模块逻辑功能的kernel函数；
2）定义实现模块真正逻辑功能的kernel函数。

注：在实际项目中，一般这种编译的扩展文件都会被集中放在一个名为src的文件夹中，在这个src文件夹中通常根据不同的功能包含很多的xxx文件夹，比如本文扩展想要实现sigmoid()函数的运算，那么CPP文件和CUDA文件就会放在sigmoid文件夹下。
那如何实现各个功能模块的编译呢，就是靠下面的setup.py文件啦~

（三）setup.py文件
实现CPP文件和CUDA文件的混合编译。

（四）test.py文件
在Python代码中实现对扩展的调用，并将相应功能封装成模块，定义相应的forward()函数和backward()函数，实现网络的前向和反向传播。

下面以实现一个非常简单的sigmoid()运算为例，介绍下如何编写上述的4个文件。

2. CPP文件

由于上面已经整体介绍过了，这里主要说一下PYBIND11_MODULE函数中m.def()3个变量是什么含义：

1）第一个变量"forward"：实际就是在Python文件中调用CPP文件中定义的C++函数的接口；
2）第二个变量&sigmoid_forward：这个位置的名称一定要和C++定义的函数名称一致，并前面加上引用&，这样在Python文件中调用forward函数接口，才会调用CPP文件中定义的sigmoid_forward()函数，从而完成后续一系列函数的调用；
3）第三个变量"sigmoid forward (CUDA)"：相当于注释，方便别人阅读或自己知道这个函数的功能。

注：在编译这个CPP代码时，使用torch::Tensor或者at::Tensor命名空间都是可以通过编译的。而且即使上面没有include Aten.h相关的库，直接使用at::Tensor也没有问题 ，希望有大佬可以解答一下，为啥行呀？ 。

#include
#include

// declare cuda forward function
torch::Tensor sigmoid_cuda_forward(torch::Tensor input);

// declare cuda backward function
torch::Tensor sigmoid_cuda_backward(torch::Tensor output);

// macro definition
#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), "data must be a CUDA Tensor.")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), "data must be contiguous.")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// define c++ forward function
torch::Tensor sigmoid_forward(torch::Tensor input){
    CHECK_INPUT(input);
    return sigmoid_cuda_forward(input);
}

// define c++ backward function
torch::Tensor sigmoid_backward(torch::Tensor output){
    CHECK_INPUT(output);
    return sigmoid_cuda_backward(output);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
  m.def("forward", &sigmoid_forward, "sigmoid forward (CUDA)");
  m.def("backward", &sigmoid_backward, "sigmoid backward (CUDA)");
}

3. CUDA文件

上述也提到过，CUDA文件中主要包含两种函数，一种是呼应CPP里声明的CUDA函数，另一种是真正实现功能的kernel函数。

在本文的例子中，sigmoid_cuda_forward()以及sigmoid_cuda_backward()函数就是呼应CPP的函数，并在这些函数中调用真正实现模块逻辑功能的kernel函数，也就是sigmoid_cuda_forward_kernel()和sigmoid_cuda_backward_kernel()函数；

其中AT_DISPATCH_FLOATING_TYPES()函数的功能，参考 Pytorch拓展C++和CUDA：讲解和细节-爱代码爱编程

注：在CUDA文件中，就只能使用at::Tensor而不能使用torch::Tensor，一旦使用torch::Tensor就会报错，不知道什么原因，希望有大佬可以解答下我的疑惑 。

#include
#include
#include
#include

template <typename scalar_t>
__device__ scalar_t sigmoid(scalar_t z){
    return 1.0 / (1.0 + exp(-z));
}

template <typename scalar_t>
__device__ scalar_t d_sigmoid(scalar_t z){
    return (1.0 - z) * z;
}

template <typename scalar_t>
__global__ void sigmoid_cuda_forward_kernel(const scalar_t * __restrict__ input, scalar_t * __restrict__ output){
    const int index = blockIdx.x * blockDim.x + blockIdx.y;
    output[index] = sigmoid(input[index]);
}

template <typename scalar_t>
__global__ void sigmoid_cuda_backward_kernel(const scalar_t* __restrict__ output,
                                             scalar_t* __restrict__ new_grad_output){
    const int index = blockIdx.x * blockDim.x + blockIdx.y;
    new_grad_output[index] = d_sigmoid(output[index]);
}

// only using at::Tensor in .cu file
// not using torch::Tensor
at::Tensor sigmoid_cuda_forward(at::Tensor input){
    auto output = at::zeros_like(input);
    dim3 blocks(input.size(0), input.size(1));
    int threads = 1;

    AT_DISPATCH_FLOATING_TYPES(input.type(), "error in sigmoid_cuda_forward", ([&]
        {sigmoid_cuda_forward_kernel<scalar_t> <<<blocks, threads>>> (input.data<scalar_t>(), output.data<scalar_t>());
        }));

    return output;
}

at::Tensor sigmoid_cuda_backward(at::Tensor output){
    auto new_grad_output = at::zeros_like(output);
    dim3 blocks(output.size(0), output.size(1));
    int threads = 1;

    AT_DISPATCH_FLOATING_TYPES(output.type(), "error in sigmoid_cuda_backward", ([&]{
    sigmoid_cuda_backward_kernel<scalar_t> <<<blocks, threads>>> (output.data<scalar_t>(),
                                                                  new_grad_output.data<scalar_t>());}));

    return new_grad_output;
}

4. setup.py文件

该文件主要就是完成上述代码的混合编译，了解的不是很多，就不多介绍了。

# setup()函数中变量的含义
name: 在python中，import该模块的名称
sources: 源代码文件的名称
laungage: 默认为c，可以改成c++
include_dirs: 传递给gcc或者g++编译器，include的头文件目录
library_dirs: 传递给gcc或者g++编译器，指定链接文件的目录
libraries: 传递给gcc或者g++编译器，指定链接的文件
extra_compile_args: 额外的gcc或者g++编译参数
extra_links_args: 额外的传给g++或者gcc编译器的链接参数
define_macros: 定义的宏
undefine_macros: 取消宏

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

setup(
    name='_CUDA',
    ext_modules=[
        CUDAExtension('sigmoid_cuda', [
            'src/sigmoid_cuda.cpp',
            'src/sigmoid_cuda_kernel.cu',
        ]),
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

5. test.py文件

该文件就是将扩展功能封装到一个模块中，并在该模块中定义前向传播和反向传播函数。之所以要定义反向传播函数是因为模块是自己用CUDA写的，无法使用PyTorch自带的反向求导，所以必须要自己实现反向传播的逻辑。具体的实现思路可以参考 科技猛兽：PyTorch 74.自定义操作torch.autograd.Function

这里主要呼应一下上面提到的m.def()函数第一个参数的"用武之地"。

from torch.nn import Module
from torch.autograd import Function
import torch
import sigmoid_cuda


class SigmoidFunction(Function):

    @staticmethod
    def forward(ctx,
                input):

        # 在cpp文件中定义为"forward",
        # 在这个位置就用.forward()实现对CPP文件中的C++函数进行调用
        output = sigmoid_cuda.forward(input)

        print(f'forward result: {output}')
        ctx.save_for_backward(output)
        return output

    @staticmethod
    def backward(ctx,
                 grad_output):
        output = ctx.saved_tensors
        output = output[0]
        grad_sigmoid = sigmoid_cuda.backward(output.contiguous())
        grad_result = grad_sigmoid * grad_output
        print(f'backward result: {grad_result}')
        return grad_result


class Dense(Module):

    def __init__(self):
        super(Dense, self).__init__()

    def forward(self, input):
        return SigmoidFunction.apply(input)


if __name__ == '__main__':
    a = torch.tensor([[1.], [3.]], dtype=torch.float32, requires_grad=True).cuda()
    print(f'input data: {a}')
    m = Dense()
    result = m(a)
    sum_ = result.sum()
    sum_.backward() 

6. 代码的具体实现链接

理论说了这么多，毕竟“纸上得来终觉浅，绝知此事要躬行”，我将上面的所有代码整理到了repo中，这个repo中还有我收集的其他扩展功能，欢迎大家PR、Star~。
https://github.com/HsLOL/ExtensionOPs/tree/master/demo​
github.com/HsLOL/ExtensionOPs/tree/master/demo

编辑于 2022-05-16 19:07