python调用C++/CUDA有不少的方法,如boost.python, cython, pybind11等。其中,pybind11是一个轻量级的仅标头的库。由于pybind11的易用性,pybind11被很多库用于于创建现有C++/CUDA代码的Python绑定,比如pytorch,tvm等。此外,由于Python的缓冲区协议可以公开自定义数据类型的内部存储,python的矩阵类型(如numpy.ndarry,torch.Tensor)可以快速转化到C++中对应矩阵类型(如Eigen,cv::Mat,vector等),不须额外的复制操作。本篇文章将通过一个数列求和的例子来讲解如何使用pybind11来将C++/CUDA代码进行Python绑定。
在通过C++和CUDA实现数列求和,在将其绑定为python函数,并在python中调用对应函数,验证结果。
基本环境
GCC: 7.5
CUDA: 10.2
CUDNN: 7.6.5
Python: 3.7
pybind11: 2.5
”hello, world!” for pybind11
不同于pybind11官网推荐的通过CMake编译的编译方式,本文将采用pytorch中pybind11的编译方式,即在setup创建python包时编译项目中C++和CUDA代码。
代码:"hello, world!"
工程目录及编译:
.../pybind$ tree
.
└── hello_world
├── cpp_extension.py
├── hello_world.cpp
└── setup_hello.py
.../pybind$ python setup_hello.py clean -a install
...
hello_world.cpp:
#include <iostream>
#include <pybind11/pybind11.h>
namespace py = pybind11;
class HelloRobot {
public:
HelloRobot(const std::string & robot_name) : robot_name_(robot_name) {};
void hello(std::string guest_name) {
if(guest_name == "") {
guest_name = "World";
}
std::cout << "Hello, " + guest_name + "! I'm " + robot_name_ + ".\n";
};
std::string get_robot_name() {return robot_name_;};
private:
std::string robot_name_;
};
PYBIND11_MODULE(hello_world, m) {
m.doc() = "pybind11 hello world";
py::class_<HelloRobot>(m, "HelloRobot")
.def(py::init<const std::string &>())
.def("hello", &HelloRobot::hello, "Provide your name...", py::arg("guest_name")="")
.def_property_readonly("robot_name", &HelloRobot::get_robot_name);
}
setup_hello.py:
import os
from setuptools import setup
from cpp_extension import BuildExtension, CUDAExtension
setup(
name='hello',
version='0.1',
description='pybind11 hello world',
python_requires='>=3.7',
setup_requires=['pybind11>=2.5.0'],
ext_modules=[CUDAExtension('hello_world',
["hello_world.cpp"],
extra_compile_args={
'cxx': ['-std=c++14', '-O2', '-Wall'],
'nvcc': [
'-std=c++14', '--expt-extended-lambda', '--use_fast_math', '-Xcompiler', '-Wall',
'-gencode=arch=compute_60,code=sm_60', '-gencode=arch=compute_61,code=sm_61',
'-gencode=arch=compute_70,code=sm_70', '-gencode=arch=compute_72,code=sm_72',
'-gencode=arch=compute_75,code=sm_75', '-gencode=arch=compute_75,code=compute_75'
],
},
include_dirs = [],
library_dirs = ['/usr/local/lib', '/usr/local/lib64/'],
)
],
cmdclass={'build_ext': BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)},
)
简单测试:
.../pybind/hello_world$ python
Python 3.7.7 (default, May 7 2020, 21:25:33)
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
ModuleNotFoundError: No module named 'hello'
>>> import hello_world
>>> hello_world.HelloRobot
<class 'hello_world.HelloRobot'>
>>> robot = hello_world.HelloRobot("Eva")
>>> robot.robot_name
'Eva'
>>> robot.hello()
Hello, World! I'm Eva.
>>> robot.hello("Young")
Hello, Young! I'm Eva.
通过pybind11生成模块的相关文档:
>>> import hello_world
>>> help(hello_world)
Help on module hello_world:
NAME
hello_world - pybind11 hello world
CLASSES
pybind11_builtins.pybind11_object(builtins.object)
HelloRobot
class HelloRobot(pybind11_builtins.pybind11_object)
| Method resolution order:
| HelloRobot
| pybind11_builtins.pybind11_object
| builtins.object
|
| Methods defined here:
|
| __init__(...)
| __init__(self: hello_world.HelloRobot, arg0: str) -> None
|
| hello(...)
| hello(self: hello_world.HelloRobot, guest_name: str = '') -> None
|
| Provide your name...
|
| ----------------------------------------------------------------------
| Data descriptors defined here:
|
| robot_name
|
| ----------------------------------------------------------------------
| Static methods inherited from pybind11_builtins.pybind11_object:
|
| __new__(*args, **kwargs) from pybind11_builtins.pybind11_type
| Create and return a new object. See help(type) for accurate signature.