Pytorch C Extension(C拓展)(即时编译JIT)(pybind11)(torch.utils.cpp_extension)

在Pytorch中也可以自定义C拓展，并调用

并且可以不用setuptools, 采用即时编译(JIT)的方法来构建C++扩展。即时编译机制提供了一种简单调用PyTorch API中名为torch.utils.cpp_extension.load()的函数，来编译和加载扩展的方法：

这里还是采用pybind11的例子

hxx_test.cpp

#include <pybind11/pybind11.h>
 
int add( int i, int j ){
    return i+j;
}
 
PYBIND11_MODULE(python_example, m ){
    m.doc() = "pybind11 example";
    m.def("add", &add, "add two number" );
}

test.py

import torch
from torch.utils.cpp_extension import load

cpp_ext = load(
     name = "python_example",
     sources = ["hxx_test.cpp",], 
     verbose=True,
)

print(cpp_ext.add(2,3))

 

 

注意torch.utils.cpp_extension.load中的name参数一定要和PYBIND11_MODULE中的参数对应起来