Skip to content
ForeverYoung
Go back

pybind11: Python Bindings for C++/CUDA Code

There are several ways to call C++/CUDA from Python — boost.python, Cython, pybind11, and others. pybind11 is a lightweight, header-only library. Its ease of use has made it the binding layer of choice for many projects, including PyTorch and TVM. Additionally, Python’s buffer protocol exposes the internal storage of custom data types, which means Python matrix types (numpy.ndarray, torch.Tensor) can be converted to their C++ counterparts (Eigen, cv::Mat, vector, etc.) without any copying. This post walks through a sequence summation example to show how to use pybind11 to create Python bindings for C++/CUDA code.

We implement sequence summation in C++ and CUDA, bind the implementations as Python functions, call them from Python, and verify the results.

Environment

GCC: 7.5
CUDA: 10.2
CUDNN: 7.6.5
Python: 3.7
pybind11: 2.5

”Hello, world!” for pybind11

Rather than the CMake-based approach recommended in the pybind11 documentation, this post uses PyTorch’s extension build system: compiling C++ and CUDA code when building the Python package via setup.py.

Code: "Hello, world!"

Project structure and build:

.../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)},
)

Quick test:

.../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.

Generated module documentation from 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.

Sequence summation


Share this post:

Previous Post
Learning Strategies for Patch-Based Local Descriptors
Next Post
Numba: Accelerating Python Code with Simple Decorators