[CMake] 用CMake编译C++程序为Python库

[CMake] 用CMake编译C++程序为Python库

我有一些c++代码，想要将其编译成Python可调用的库，应该怎么做呢？

这里有2种做法

• 一种是使用python的setuptools，参见[CUDA] 手写一个PyTorch的算子。不过setuptools的API跟make比较像，如果你的项目比较大或者引用了很多外部库，在配置编译参数的时候可能会比较麻烦。
• 另一种是直接使用cmake。本人用cmake用惯了（觉得make的参数太多了懒得写），所以还是觉得用cmake编译最顺手。

所以这篇文章就简单讲一下如何用cmake编译c++代码成为python库。使用cmake与使用setuptools的相同之处在于，我们依然要使用pybind11这个库。

cpp文件

C++代码的编写与用setuptools差不多，都是定义一个PYBIND11_MODULE：

#include <pybind11/pybind11.h>

PYBIND11_MODULE(module_name, m) {
  m.def("python_func", &cpp_func, "Description");
}

CMakeLists.txt

在编写CMakeLists.txt时，在开头加入

find_package(pybind11 REQUIRED)

用于导入pybind11包。

在创建目标库的时候，使用

pybind11_add_module(module_name MODULE xxx.cpp)
target_link_libraries(module_name PRIVATE ...)

这样就可以使用pybind11创建一个python可调用的库了。

注意：这里的module_name是你编译的模块名字，它必须和上面c++代码中的module_name相同，否则在调用时会报错。

另外，上面的方法和下面的方法是等价的：

add_library(module_name MODULE xxx.cpp)
target_link_libraries(module_name PRIVATE pybind11::module ...)

Python

在你成功用cmake完成编译后，它会生成一个动态链接库文件（形如module_name.cpython-310-x86_64-linux-gnu.so）。你将这个动态链接库加入到path中，然后就可以在python当中import它了。

使用torch的tensor

如果你正在编写一个使用pytorch的项目，你可能希望你的c++模块可以支持torch的tensor作为参数。为此，你需要加入下面这些内容：

cpp

对于c++代码，导入pytorch库。

#include <torch/extension.h>

传递参数时直接用torch::Tensor即可

CMakeLists.txt

对于CMakeLists.txt，在开头加入

find_package(Torch REQUIRED)
find_library(TORCH_PYTHON_LIBRARY torch_python PATH "${TORCH_INSTALL_PREFIX}/lib")

在创建目标库时候，加入

target_link_libraries(module_name PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} ...)

在调用cmake的时候，还需加入以下选项：

cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` -DCMAKE_CUDA_COMPILER="/usr/local/cuda/bin/nvcc" ..

函数重载

在使用pybind11来创建Python接口的时候，我们可能想要指定一个C++中被重载的函数。此时，我们需要指定这个函数的参数和返回值类型，才能让接口找到具体是哪个被重载的函数。

比如我想指定那个参数为(int,int)，返回值为double的函数。

PYBIND11_MODULE(module_name, m) {
    m.def("python_func", static_cast<double (*)(int, int)>(&func), "Description");
}

如果要重载的是一个类里的成员函数，可以用如下方法：

PYBIND11_MODULE(module_name, m) {
    pybind11::class_<ClassName>(m, "Description")
    .def("python_func", static_cast<double (ClassName::*)(int, int)>(&ClassName::func), "Description");
}