pybind11实现python和c++之间的数据通信。这是文档。
Pytorch的文档CUSTOM C++ AND CUDA EXTENSIONS也提到了如何进行绑定，结合起来一个简单的绑定例子是
文件结构
1
2
3
4
extend_demo/
setup.py
op.cpp
op_test.py
1
2
3
4
5
6
7
8
# setup.py
from setuptools import setup, Extension
from torch.utils import cpp_extension
setup(name='simpleOp',
ext_modules=[cpp_extension.CppExtension('simpleOp', ['op.cpp'])],
cmdclass={'build_ext': cpp_extension.BuildExtension})
1
2
3
4
5
6
7
8
9
10
11
12
// op.cpp
#include <torch/extension.h>
#include <iostream>
int add(int x, int y){
return x + y;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("add", &add, "add two number");
}
执行python setup.py install，开始编译。
编译成功后，看看conda list里有
1
simpleOp 0.0.0 <pip>
然后就可以像导入一个普通的包来调用我们自定义的模块啦：
1
2
3
4
5
6
7
8
9
# op_test.py
import torch # torch should be imported first !
import simpleOp
a = 1
b = 1
c = simpleOp.add(a, b)
print(a, b, c) # 1 1 2