在PyTorch中由Python端调用C端底层函数
今天看到一篇非常好的知乎blog，学到了如何在PyTorch下由python端的代码调用C代码。
一、目标
在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
