一、目标

在DNN训练的过程中为了从Python端调用C的代码，方便接下去的科研。

二、前文

参考：https://zhuanlan.zhihu.com/p/358778742

在PyTorch的框架中我们能在下图的文件夹中找到load函数：

image.png

在框架中是这么描述这个函数的：Loads a PyTorch C++ extension just-in-time (JIT).
即使用即时编译将Python与C联系起来，并且是在python代码运行的过程中系统自动编译。

这里要注意的地方是代码中要用pybind11进行呼应。下文细讲

三、内容

即时编译涉及如下文件：

1 Python文件（主文件）
2 C++头文件
3 Cpp文件
4 Cuda文件

在python中需要加入如下代码：

cuda_module = load(name="add2",
                           extra_include_paths=["/home/cping/PyTorch_scripts/Cuda_JIT_Test/NN-CUDA-Example/include"],
                           sources=["/home/cping/PyTorch_scripts/Cuda_JIT_Test/NN-CUDA-Example/pytorch/add2_ops.cpp", "/home/cping/PyTorch_scripts/Cuda_JIT_Test/NN-CUDA-Example/kernel/add2_kernel.cu"],
                           verbose=True)

其中add2就是个名字，可以换。但是extra_include_paths传入的是头文件的位置.h文件；sources要给定路径，这里传入cpp以及cu文件。

verbose代表开启日志。
下面是例子：

Example:
        >>> from torch.utils.cpp_extension import load
        >>> module = load(
                name='extension',
                sources=['extension.cpp', 'extension_kernel.cu'],
                extra_cflags=['-O2'],
                verbose=True)

之后在python文件中调用：cuda_module.torch_launch_add2(c, a, b, n)

其中torch_launch_add2()为cpp中的函数

现在来看cpp代码：

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

void torch_launch_add2(torch::Tensor &c,
                       const torch::Tensor &a,
                       const torch::Tensor &b,
                       int64_t 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");
}

这里我们要包含两个.h文件，其中add2.h是我们diy代码的头文件。

也就只有一句话;

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

回到cpp文件中，这里需要注意一个地方：

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("torch_launch_add2",
          &torch_launch_add2,
          "add2 kernel warpper");
}

这个地方必须要加，可以理解为python找c的入口（python中的cuda_module.torch_launch_add2(c, a, b, n)），而这个函数就是cpp中的函数。

之后调用launch_add2函数。而这个函数在cu文件中实现：

__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);
}

至此，逻辑到位，直接python3 xxx.py就可以运行，系统自动编译运算，如下图。

image.png