前言

Pytorch拓展C语言并不难，因为我们有torch.util.ffi模块；Pytorch拓展cuda语言也不难，因为pytorch的前身为torch，torch是使用lua语言进行编写的，lua语言最大的特点就是和C语言可以有很好的互动。因此，pytorch这个python版的torch，使用的cuda底层其实和torch是相似的。

torch的底层是由c语言编写。而到了pytorch这里，底层大部分的语言修改了一些，但大部分依然是C,只是编译的环境由C转变为C++。我们都知道C++对C语言是兼容的，也就是说，在pytorch中，官方更推荐使用C++接口来拓展pytorch，因为这个是发展趋势，下图在Pytorch官方对这个看法的大致说明。

总结就是一句话，Pytorch正在将底层由C慢慢转化为C++(并不说底层变为C++，C语言接口就不能用了),但是目前来说，底层大部分还是C语言，移植于原来的torch底层代码。

为什么要拓展

原因很简单，在初始阶段，我们只需要利用pytorch中提供的接口组合搭配自己设计的神经网络即可。但是我们如果进一步发展，例如需要实现一个自己的算法，或者自己的自定义层，光使用pytorch自带的模块就不够用了。

关于自定义层：探讨Pytorch中nn.Module与nn.autograd.Function的backward()函数

Pytorch虽然已经使用了NVIDIA cuDNN、Intel MKL和NNPACK这些底层来加快训练速度，但是在某些情况下，比如我们要实现一些特定算法，光靠组合Pytorch已有的操作是不够的。这是因为Pytorch虽然在特定操作上经过了很好的优化，但是对于Pytorch已经写好的这些操作，假如我们组合起来，组成我们的新的算法，Pytorch才不管你的算法的具体执行流程，一般Pytorch只会按照设计好的操作去使用GPU的通道，然后通道不能充分利用或者直接超负载，然后python解释器也不能对此进行优化，导致程序执行速度反而变慢了。

只有自己深度底层设计算法，与GPU亲密接触，才能够充分利用GPU的grid、block和thread，大幅提升算法的速度。

C语言拓展

C语言拓展需要上文提到的torch.util.ffi模块，这个模块在安装pytorch的时候就已经自带了。

另外还有个需要注意的问题,那就是我们要编写C语言的底层接口：TH、THC、THCS…，这边有个最大的问题，就是，这些接口没有官方说明(或者说我没有找到)..需要我们自己去阅读源码，去分析，所幸我们可以通过阅读一些其他人的实现来学习。

这里有一个简明的torch库教程：https://apaszke.github.io/torch-internals.html

编写C语言代码

编写C接口的代码很简单，以官方的为例，其实只要稍微对pytorch或者其他深度学习框架熟悉的，见名知意，从函数名称就可以大概知道其函数作用：

以下代码实现一个自己的加法层。

/* src/my_lib.c */
#include <TH/TH.h>

int my_lib_add_forward(THFloatTensor *input1, THFloatTensor *input2,
THFloatTensor *output)
{
    if (!THFloatTensor_isSameSizeAs(input1, input2))
        return 0;
    THFloatTensor_resizeAs(output, input1);
    THFloatTensor_cadd(output, input1, 1.0, input2);
    return 1;
}

int my_lib_add_backward(THFloatTensor *grad_output, THFloatTensor *grad_input)
{
    THFloatTensor_resizeAs(grad_input, grad_output);
    THFloatTensor_fill(grad_input, 1);
    return 1;
}

头文件TH就是pytorch底层代码的接口头文件，关于头文件的一些简单介绍，TH就是cpu下的pytorch底层文件，GPU下测试THC：

有了.c当然我们还需要头文件.h：

/* src/my_lib.h */
int my_lib_add_forward(THFloatTensor *input1, THFloatTensor *input2, THFloatTensor *output);
int my_lib_add_backward(THFloatTensor *grad_output, THFloatTensor *grad_input);

最后在同目录下创建一个.py文件，内容如下：

# build.py
from torch.utils.ffi import create_extension
ffi = create_extension(
name='_ext.my_lib',        # 输出文件地址及名称
headers='src/my_lib.h',    # 编译.h文件地址及名称
sources=['src/my_lib.c'],  # 编译.c文件地址及名称
with_cuda=False            # 不使用cuda
)
ffi.build()

上面的程序就是使用ffi这个模块对我们刚才写的c语言代码进行编译，具体过程不用我们操心。

调用c语言代码

写好之后调用就比较简单了，在编译过后pytorch就会创建一个_ext的目录然后生成.so链接文件。

然后我们通过torch.autograd.Function自定义层，引用我们之前生成的my_lib：

# functions/add.py
import torch
from torch.autograd import Function
from _ext import my_lib


class MyAddFunction(Function):
    def forward(self, input1, input2):
        output = torch.FloatTensor()
        my_lib.my_lib_add_forward(input1, input2, output)
        return output

    def backward(self, grad_output):
        grad_input = torch.FloatTensor()
        my_lib.my_lib_add_backward(grad_output, grad_input)
        return grad_input

将自定义层包装一下：

# modules/add.py
from torch.nn import Module
from functions.add import MyAddFunction

class MyAddModule(Module):
    def forward(self, input1, input2):
        return MyAddFunction()(input1, input2)

然后调用执行：

# main.py
import torch
import torch.nn as nn
from modules.add import MyAddModule

class MyNetwork(nn.Module):
    def __init__(self):
        super(MyNetwork, self).__init__()
        self.add = MyAddModule()

    def forward(self, input1, input2):
        return self.add(input1, input2)

model = MyNetwork()
input1, input2 = torch.randn(5, 5), torch.randn(5, 5)
print(model(input1, input2))
print(input1 + input2)

Cuda拓展

C语言拓展比较简单，我们只需要注意一些接口层以及函数的定义知道如何使用他们就可以了。但是光编写C语言利用只能在cpu上跑。我们实际需要的还是使用GPU显卡来跑，因此我们需要编写cuda代码，然后连接c语言，然后用pytorch导入使用。

编写cuda代码

cuda是并行语言，运行在英伟达的显卡GPU上，是nvidia公司结合C/C++编写出来的语言，后缀为.cu，头文件为.h。语法支持大部分C/C++,然后根据需要设计了一些其他的语法。

现在我们编写一段cuda代码，实现broadcast-sum，就是element-wise的相加，其实在pytorch中已经实现了，这里只是用作cuda的演示：

// mathutil_cuda_kernel.cu
// 头文件，最后一个是cuda特有的
#include <curand.h>
#include <stdio.h>
#include <math.h>
#include <float.h>
#include "mathutil_cuda_kernel.h"

// 获取GPU线程通道信息
dim3 cuda_gridsize(int n)
{
    int k = (n - 1) / BLOCK + 1;
    int x = k;
    int y = 1;
    if(x > 65535) {
        x = ceil(sqrt(k));
        y = (n - 1) / (x * BLOCK) + 1;
    }
    dim3 d(x, y, 1);
    return d;
}
// 这个函数是cuda执行函数，可以看到细化到了每一个元素
__global__ void broadcast_sum_kernel(float *a, float *b, int x, int y, int size)
{
    int i = (blockIdx.x + blockIdx.y * gridDim.x) * blockDim.x + threadIdx.x;
    if(i >= size) return;
    int j = i % x; i = i / x;
    int k = i % y;
    a[IDX2D(j, k, y)] += b[k];
}


// 这个函数是与c语言函数链接的接口函数
void broadcast_sum_cuda(float *a, float *b, int x, int y, cudaStream_t stream)
{
    int size = x * y;
    cudaError_t err;
    
    // 上面定义的函数
    broadcast_sum_kernel<<<cuda_gridsize(size), BLOCK, 0, stream>>>(a, b, x, y, size);

    err = cudaGetLastError();
    if (cudaSuccess != err)
    {
        fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
        exit(-1);
    }
}

上面的代码和C语言很像，但是多了一个__global__,这个是cuda中特有类型，这个函数实现向量a和b的element-wise的相加。

头文件当然也要有：

#ifndef _MATHUTIL_CUDA_KERNEL
#define _MATHUTIL_CUDA_KERNEL

#define IDX2D(i, j, dj) (dj * i + j)
#define IDX3D(i, j, k, dj, dk) (IDX2D(IDX2D(i, j, dj), k, dk))

#define BLOCK 512
#define MAX_STREAMS 512

#ifdef __cplusplus
extern "C" {
#endif

void broadcast_sum_cuda(float *a, float *b, int x, int y, cudaStream_t stream);

#ifdef __cplusplus
}
#endif

#endif

然后我们写与pytorch链接的C语言部分：

// mathutil_cuda.c
// THC是pytorch底层GPU库
#include <THC/THC.h>
#include "mathutil_cuda_kernel.h"

extern THCState *state;

int broadcast_sum(THCudaTensor *a_tensor, THCudaTensor *b_tensor, int x, int y)
{
    float *a = THCudaTensor_data(state, a_tensor);
    float *b = THCudaTensor_data(state, b_tensor);
    cudaStream_t stream = THCState_getCurrentStream(state);

    // 这里调用之前在cuda中编写的接口函数
    broadcast_sum_cuda(a, b, x, y, stream);

    return 1;
}

头文件：

int broadcast_sum(THCudaTensor *a_tensor, THCudaTensor *b_tensor, int x, int y);

编写好后，放到一个文件夹里：

cuda分别是：

• mathutil_cuda_kernel.cu
• mathutil_cuda_kernel.h

c语言：

• mathutil_cuda.c
• mathutil_cuda.h

这回我们不能一起编译了，因为pytorch的设计bug原因，在pytorch中使用.c和.cu文件时时，需要首先编译一下.cu然后将编译好的.cu包含在路径中和.c一起再用cffi进行编译才能使用。

所以首先我们编译.cu文件：

nvcc -c -o mathutil_cuda_kernel.cu.o mathutil_cuda_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52

编译好后，编写build.py，然后运行:

import os
import torch
from torch.utils.ffi import create_extension

this_file = os.path.dirname(__file__)

sources = []
headers = []
defines = []
with_cuda = False

if torch.cuda.is_available():
    print('Including CUDA code.')
    sources += ['src/mathutil_cuda.c']
    headers += ['src/mathutil_cuda.h']
    defines += [('WITH_CUDA', None)]
    with_cuda = True

this_file = os.path.dirname(os.path.realpath(__file__))

extra_objects = ['src/mathutil_cuda_kernel.cu.o']   # 这里是编译好后的.o文件位置
extra_objects = [os.path.join(this_file, fname) for fname in extra_objects]


ffi = create_extension(
    '_ext.cuda_util',
    headers=headers,
    sources=sources,
    define_macros=defines,
    relative_to=__file__,
    with_cuda=with_cuda,
    extra_objects=extra_objects
)

if __name__ == '__main__':
    ffi.build()

这个build.py与之前相比添加了编译好的cuda.o文件，添加进去与C文件一起编译出来。

编译好后，我们就可以使用了：

mathutil.broadcast_sum(a, b, *map(int, a.size()))
# 实现以下的效果：

a = torch.randn(3, 5)
b = torch.randn(3, 1)

a += b

后记

暂时先说到这，pytorch的拓展还有很多需要说的地方，之后的文章中会分析一下pytorch的源码以及如何对pytorch进行C++拓展进行讲解，请大家多留意~

参考：

https://pytorch.org/tutorials/advanced/c_extension.html
https://discuss.pytorch.org/t/how-to-understand-c-ffi-extensions-in-pytorch/7780
https://discuss.pytorch.org/t/compiling-an-extension-with-cuda-files/302/9
https://infinitescript.com/2018/03/build-pytorch-extensions-with-cuda-and-cffi/