利用pybind11封装CUDA程序为python库

Pybind11是C++/Python混合编程的利器之一，是一个轻量级的只包含头文件的库，用于 Python 和 C++ 之间接口转换，可以为现有的 C++ 代码创建 Python 接口绑定。

Module声明

• 代码结构如下：

  ​ — qqq_gemm.cu

  ​ — qqq_gemm.h

  ​ — pybind.cpp

  ​ —setup.py

• 在使用pybind11封装我们的CUDA代码的时候，我们首先用到PYBIND11_MODULE的宏

  #include "qqq_gemm.h"
    
  PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("qqq_gemm", &qqq_gemm, "INT8xINT4 matmul based marlin FP16xINT4 kernel.");
  }
  

  • 第一个参数是想生成的模块名字，第二个其实是一个pybind11::module类型，m.doc后面填对于这个模块的解释，m.def后面第一个参数是在python里写的函数名，第二个参数是绑定的函数地址，后面是对函数的解释（可以不写）

Setup

• 通过pybind11可以把上述的qqq_gemm.cu（头文件为qqq_gemm.h）编译为.so方便使用

  from setuptools import setup, find_packages
  from torch.utils import cpp_extension
    
  setup(
      name='qqq_gemm',
      ext_modules=[
          cpp_extension.CUDAExtension(
              name='qqq_gemm',
              sources=[
                  'pybind.cpp',
                  'qqq_gemm.cu'
              ],
          ),
      ],
      cmdclass={
          'build_ext': cpp_extension.BuildExtension.with_options(use_ninja=False)
      },
      packages=find_packages(
          exclude=['notebook', 'scripts', 'tests']),
  )
  

  • cmdclass一行定义了一个 cmdclass 参数,它用于自定义 distutils 或 setuptools 命令。
    • 在这里,它指定了一个名为 build_ext 的自定义命令类,该类继承自 cpp_extension.BuildExtension。
    • with_options(use_ninja=False) 是一个选项,用于禁用 Ninja 构建系统,改为使用 distutils 的标准构建系统。
  • packages一行使用 find_packages() 函数自动发现要包含在包中的所有 Python 模块和包。
    • exclude 参数指定了要从搜索结果中排除的目录,在这里排除了 notebook、scripts 和 tests 目录。

Install

• 通过运行以下命令进行安装:

  python setup.py develop

• 然后我们可以找到build里（当前目录下也有）的 qqq_gemm.cpython-310-x86_64-linux-gnu.so文件
  • 在此文件目录下可以直接import qqq_gemm然后可以调用上述函数

Error

ImportError: libc10.so: cannot open shared object file: No such file or directory

• 这表示我们需要在import自己库之前先import一下torch