pybind机制

使用场景

• python里面调用C++代码

举例说明：小demo拆解

• 链接：datacopy_nvshmem_pybind

目标

• 因为跨服务器的传输，不用MPI的话就得用pytorch提供的分布式通信接口
• pytorch的分布式通信接口是python的
• 所以必须得用python调用底层编写的C++代码
• 这个项目还有一个麻烦的事情是必须要使用到nvshmem库，一些链接过程也写在这里

代码结构

├─include
└─src
    ├─datacopy.cu
    └─pybind_data_copy.cpp
└─build.sh
└─CMakeLists.txt
└─setup.py
└─test.py

include文件夹

• 没什么特别的，include里面包含了一些头文件

src文件夹

• datacopy.cu：主要是实现了数据传输的功能，用的是普通的C++语法和CUDA语法
• pybind_data_copy.cpp：主要是实现了python和C++的接口，用的是pybind11的语法

  #include <pybind11/pybind11.h>
  #include <pybind11/stl.h>
  #include "double_buffer_manager.h" 
  
  // 这两个函数实现在datacopy.cu里面
  std::vector<uint8_t> get_unique_id();
  int init_nvshmem(const std::vector<uint8_t>& root_unique_id, int rank, int world_size);
  
  namespace py = pybind11;
  
  PYBIND11_MODULE(pydatacopy, m) {
      py::class_<DoubleBufferManager>(m, "DoubleBufferManager")
          .def(py::init<>())
          .def("init", &DoubleBufferManager::init, "Initialize double buffer")
          .def("cleanup", &DoubleBufferManager::cleanup, "Cleanup double buffer")
          .def("test_bandwidth", &DoubleBufferManager::test_bandwidth, "Test bandwidth");
      
      m.def("get_unique_id", &get_unique_id, "Generate a unique ID for NVSHMEM initialization");
      m.def("init_nvshmem", &init_nvshmem, "Initialize NVSHMEM",
            py::arg("root_unique_id"), py::arg("rank"), py::arg("world_size"));
  }

• 主要是用PYBIND11_MODULE宏定义了一个名为pydatacopy的模块
• 这个模块包含了一个名为DoubleBufferManager的类和两个函数get_unique_id和init_nvshmem
• 这样就可以在python中import pydatacopy，然后使用DoubleBufferManager类和这两个函数

build.sh

#!/bin/bash

# if no build dir, create one
mkdir -p build
cd build

export NVSHMEM_HOME=/usr/local/nvshmem
export Python_ROOT_DIR=$(python3.12 -c "import sys; print(sys.prefix)")

export TORCH_CUDA_ARCH_LIST="9.0"

cmake .. \
    -G Ninja \  # use ninja instead of make to get faster
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_CUDA_ARCHITECTURES=90 \
    -DPython_EXECUTABLE=$(which python3.12) \
    -DTorch_DIR=/usr/local/lib/python3.12/dist-packages/torch/share/cmake/Torch \
    -Dpybind11_DIR=$(python3.12 -c "import pybind11; print(pybind11.get_cmake_dir())")

ninja -j$(nproc)

• 第一行的#!/bin/bash表示这个脚本是用bash解释器来执行的，这个指令必须放在脚本的第一行
• mkdir -p build：创建一个名为build的目录，-p选项表示如果要创建的一个目录路径下任何一个目录不存在，都会自动创建出来，一次可以创建多个目录
• export NVSHMEM_HOME=/usr/local/nvshmem：设置环境变量NVSHMEM_HOME为/usr/local/nvshmem，这是nvshmem库的安装路径
• export Python_ROOT_DIR=$(python3.12 -c "import sys; print(sys.prefix)")：设置环境变量Python_ROOT_DIR为python3.12的安装路径
  • python3.12 -c "import sys; print(sys.prefix)"：这条命令会输出python3.12的安装路径
  • $(...)：这是命令替换的语法，会把命令的输出结果赋值给变量
  • 当 CMake 执行 find_package(Python COMPONENTS Interpreter Development)时：会首先检查 Python_ROOT_DIR，使用该路径查找 Python 开发文件，验证找到的 Python 版本是否符合要求
  • 如果在标准安装中，返回python的安装根目录；如果是虚拟环境，返回虚拟环境的根目录
  • 不设置可能导致找到错误版本的python
• cmake .. \：运行cmake命令，..表示上一级目录，也就是项目的根目录
• -G Ninja \：指定使用Ninja作为构建系统，Ninja是一种更快的替代Make的构建工具
• -DCMAKE_BUILD_TYPE=Release \：指定构建类型为Release，表示生成优化过的代码
• -DCMAKE_CUDA_ARCHITECTURES=90 \：指定CUDA的计算能力为9.0，H20可用
• -DPython_EXECUTABLE=$(which python3.12) \：指定python解释器的路径
• -DTorch_DIR=/usr/local/lib/python3.12/dist-packages/torch/share/cmake/Torch \：指定pytorch的CMake配置文件路径，这个路径下面实际上存的是一些.cmake文件
• -Dpybind11_DIR=$(python3.12 -c "import pybind11; print(pybind11.get_cmake_dir())")：指定pybind11的CMake配置文件路径，编译好pybind11之后，这个目录下面存的也就是一些.cmake文件
• ninja -j$(nproc)：运行ninja，编译

CMakeLists.txt

• 解释见注释

  cmake_minimum_required(VERSION 3.18)
  
  # 设置项目名称和使用的语言，这里用到C++和CUDA
  project(pydatacopy LANGUAGES CXX CUDA)
  
  # 设置编译选项，关闭冗长的编译输出
  set(CMAKE_VERBOSE_MAKEFILE OFF)
  
  # ccache加速编译，在系统中查找ccache程序，把路径赋值给CCACHE_PROGRAM变量
  find_program(CCACHE_PROGRAM ccache)
  if(CCACHE_PROGRAM)
      set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
      set(CMAKE_CUDA_COMPILER_LAUNCHER ${CCACHE_PROGRAM})
      message(STATUS "Using ccache for compilation acceleration")
  endif()
  
  # 查找python这个外部依赖包，3.12版本，REQUIRED表示必须找到，否则报错，COMPONENTS后面的内容指定需要的部分，需要解释器和开发文件，这里会优先查找Python_ROOT_DIR变量指定的路径
  find_package(Python 3.12 REQUIRED COMPONENTS Interpreter Development)
  message(STATUS "Python include dirs: ${Python_INCLUDE_DIRS}")
  message(STATUS "Python libraries: ${Python_LIBRARIES}")
  
  # set()：CMake 命令，用于定义变量
  set(TORCH_DIR "/usr/local/lib/python3.12/dist-packages/torch/share/cmake/Torch")
  message(STATUS "Using Torch at: ${TORCH_DIR}")
  # CMAKE_PREFIX_PATH：CMake 的特殊变量，用于存储查找包的路径列表
  # list(APPEND ...)：向列表末尾添加新元素
  # 作用是将 TORCH_DIR 添加到 CMAKE_PREFIX_PATH 中，这样 find_package(Torch ...) 时就会在这个路径下查找TorchConfig.cmake
  list(APPEND CMAKE_PREFIX_PATH ${TORCH_DIR})
  
  set(CMAKE_CUDA_ARCHITECTURES "90")
  set(TORCH_CUDA_ARCH_LIST "9.0")
  
  # 使用torch的cmake配置文件，用于查找外部依赖包
  find_package(Torch REQUIRED CONFIG)
  
  find_package(pybind11 REQUIRED CONFIG)
  
  
  set(NVSHMEM_HOME "/usr/local/nvshmem")
  message(STATUS "Using NVSHMEM at: ${NVSHMEM_HOME}")
  set(NVSHMEM_LIB "${NVSHMEM_HOME}/lib/libnvshmem.a")
  
  
  find_package(CUDAToolkit REQUIRED)
  
  # 设置源文件
  set(SOURCES 
      src/pybind_data_copy.cpp
      src/data_copy.cu
  )
  
  # 包含头文件目录
  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
  
  # 7. 获取 Python 模块后缀
  execute_process(
      COMMAND python3.12 -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))"
      OUTPUT_VARIABLE PYTHON_MODULE_SUFFIX
      OUTPUT_STRIP_TRAILING_WHITESPACE
  )
  message(STATUS "Python module suffix: ${PYTHON_MODULE_SUFFIX}")
  
  # 8. 创建模块
  # 创建一个名为 pydatacopy的共享库（动态链接库），这将成为 Python 可导入的扩展模块
  add_library(pydatacopy SHARED ${SOURCES})
  
  
  # prefix：设置库文件名的前缀为空字符串，默认情况下，Linux/Mac 共享库有 lib前缀（如 libpydatacopy.so），Python 扩展模块不需要前缀，所以设为空
  # suffix：设置库文件的后缀为 Python 模块的标准后缀（如 .so 或 .pyd），确保生成的文件名符合 Python 的要求
  # POSITION_INDEPENDENT_CODE ON：生成位置无关代码（PIC），这对于共享库是必要的，多个程序共享同一段代码。位置无关的意思是代码可以加载到内存的任何位置，用的都是相对地址
  # CUDA_SEPARABLE_COMPILATION ON：启用 CUDA 的可分离编译，这允许 CUDA 代码跨多个文件进行编译和链接，这可以支持在修改的情况下只重新编译修改的部分，还可以支持启用设备链接函数，一个CUDA设备函数可以调用其他文件里的函数
  set_target_properties(pydatacopy PROPERTIES
      PREFIX ""
      SUFFIX "${PYTHON_MODULE_SUFFIX}"
      POSITION_INDEPENDENT_CODE ON
      CUDA_SEPARABLE_COMPILATION ON
  )
  
  # 9. 添加包含目录
  target_include_directories(pydatacopy PRIVATE
      ${CMAKE_CURRENT_SOURCE_DIR}/include
      ${Python_INCLUDE_DIRS}
      ${Torch_INCLUDE_DIRS}
      ${pybind11_INCLUDE_DIRS}
      ${NVSHMEM_HOME}/include
      ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
  )
  
  # 10. 添加链接库
  target_link_libraries(pydatacopy PRIVATE
      ${TORCH_LIBRARIES}
      ${Python_LIBRARIES}
      cudart
      ${NVSHMEM_LIB}
      dl
      pybind11::pybind11
  )
  
  # 11. 添加编译选项 - 强制指定架构
  target_compile_options(pydatacopy PRIVATE
      $<$<COMPILE_LANGUAGE:CUDA>:
          -Xcompiler=-fPIC
          --expt-relaxed-constexpr
          --ptxas-options=-O3  # 优化PTX汇编
          --maxrregcount=64    # 限制寄存器使用以提高并行度
          --use_fast_math      # 使用快速数学（精度略低但更快）
          --gpu-architecture=compute_90  # 强制指定计算能力
          --gpu-code=sm_90    # 强制指定目标架构
      >
      $<$<COMPILE_LANGUAGE:CXX>:
          -fPIC
          -O3
          -march=native        # 使用本地CPU架构优化
      >
  )
  
  # 12. 设置输出目录
  set_target_properties(pydatacopy PROPERTIES
      LIBRARY_OUTPUT_DIRECTORY ${CMAKE_SOURCE_DIR}
      CUDA_RESOLVE_DEVICE_SYMBOLS ON
  )
  
  # 13. 安装规则
  install(TARGETS pydatacopy 
      LIBRARY DESTINATION lib
      ARCHIVE DESTINATION lib
  )

setup.py

from setuptools import setup
from setuptools import Extension
import sys

setup(
    name="pydatacopy",
    version="0.1",
    # ​​ext_modules​​：包含 C/C++ 扩展模块的列表
    ext_modules=[
        Extension(
            "pydatacopy",
            # 这里如果sources写了东西，比如.cu，.cpp，就会编译，没写就用下面的package_data
            sources=[],
        )
    ],
    package_data={'': ['build/pydatacopy.cpython-312-x86_64-linux-gnu.so']},
    include_package_data=True,
)

• 然后如果sources里面有东西，就执行：

  python setup.py build

• 然后把编译好的pydatacopy.cpython-…so文件拷贝到python的site-packages目录下，以便可以直接import导入

  python setup.py install