基于CMake构建3D Gaussian Splatting的CUDA Python扩展遇阻求助

我在开发基于3D Gaussian Splatting的CUDA/C++ Python扩展，希望通过CMake在setup.py中完成编译。我的思路是先用CMake生成.so文件，再基于该文件构建Python扩展（参考了相关方案），但实际执行后setup.py无法正确使用生成的.so文件，最终只生成了仅含__init__.py的无效扩展。以下是我的setup.py和CMakeLists.txt代码，请帮忙排查问题：

from setuptools import setup,Extension
from setuptools.command.build_ext import build_ext
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
import subprocess
import os
import shutil

current_dir=os.path.dirname(os.path.abspath(__file__))

class CMakeBuildExt(BuildExtension):
    def run(self):
        build_dir = os.path.join(current_dir, 'build')

        subprocess.check_call(['cmake', '..'], cwd=build_dir)
        subprocess.check_call(['make', 'clean'], cwd=build_dir)
        subprocess.check_call(['make'], cwd=build_dir)
        shutil.copy(os.path.join(build_dir,  'libCudaRasterizer.so'), 
                    os.path.join(current_dir, 'diff_surfel_rasterization', 'libCudaRasterizer.so'))
        super().run()



setup(
    name="diff_surfel_rasterization",
    packages=['diff_surfel_rasterization'],
    version='0.0.1',
    cmdclass={
        'build_ext': CMakeBuildExt
    },
    ext_modules=[
        CUDAExtension(
            name="diff_surfel_rasterization._C",
            sources=[],
            extra_compile_args={"nvcc": ["-I" + os.path.join(os.path.dirname(os.path.abspath(__file__)), "third_party/glm/")]})
        ],
    
)

#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

cmake_minimum_required(VERSION 3.20)

project(DiffRast LANGUAGES CUDA CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_EXTENSIONS OFF)
set(CMAKE_CUDA_STANDARD 17)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")

add_library(CudaRasterizer
    cuda_rasterizer/backward.h
    cuda_rasterizer/backward.cu
    cuda_rasterizer/forward.h
    cuda_rasterizer/forward.cu
    cuda_rasterizer/auxiliary.h
    cuda_rasterizer/rasterizer_impl.cu
    cuda_rasterizer/rasterizer_impl.h
    cuda_rasterizer/rasterizer.h
)

set_target_properties(CudaRasterizer PROPERTIES CUDA_ARCHITECTURES "70;75;86")

target_include_directories(CudaRasterizer PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/cuda_rasterizer)
target_include_directories(CudaRasterizer PRIVATE third_party/glm ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})

核心问题分析 #

1. 空sources导致扩展被忽略：你在CUDAExtension中传入了空的sources列表，setuptools会判定该扩展无编译内容，直接跳过扩展构建流程，仅打包packages参数指定的目录（也就是仅含__init__.py的空包）。
2. CMake生成的库未被纳入构建流程：即使手动复制libCudaRasterizer.so到源码目录，setuptools也不会自动将其关联到Python扩展，也不会默认打包到最终安装包中。

解决方案 #

提供两种可行修改方案，可根据需求选择：

方案一：让CMake直接生成Python扩展模块（推荐） #

利用pybind11将CUDA核心库包装为Python扩展，直接让CMake生成符合Python规范的.so模块，再通过setup.py完成安装。

1. 修改CMakeLists.txt：
   引入pybind11依赖，拆分核心库与Python绑定模块的构建：

   cmake_minimum_required(VERSION 3.20)
   
   project(DiffRast LANGUAGES CUDA CXX)
   
   set(CMAKE_CXX_STANDARD 17)
   set(CMAKE_CXX_EXTENSIONS OFF)
   set(CMAKE_CUDA_STANDARD 17)
   
   # 引入pybind11
   find_package(pybind11 REQUIRED)
   
   # 编译CUDA核心静态库
   add_library(CudaRasterizerCore STATIC
       cuda_rasterizer/backward.h
       cuda_rasterizer/backward.cu
       cuda_rasterizer/forward.h
       cuda_rasterizer/forward.cu
       cuda_rasterizer/auxiliary.h
       cuda_rasterizer/rasterizer_impl.cu
       cuda_rasterizer/rasterizer_impl.h
       cuda_rasterizer/rasterizer.h
   )
   
   set_target_properties(CudaRasterizerCore PROPERTIES CUDA_ARCHITECTURES "70;75;86")
   
   target_include_directories(CudaRasterizerCore PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/cuda_rasterizer)
   target_include_directories(CudaRasterizerCore PRIVATE third_party/glm ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
   
   # 用pybind11生成Python扩展模块
   pybind11_add_module(_C MODULE
       # 新增绑定代码文件，需自行实现接口暴露逻辑
       cuda_rasterizer/bindings.cpp
   )
   
   target_link_libraries(_C PRIVATE CudaRasterizerCore)
   # 设置输出路径与名称，匹配setup.py中的扩展标识
   set_target_properties(_C PROPERTIES
       OUTPUT_NAME "diff_surfel_rasterization._C"
       LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/diff_surfel_rasterization"
   )
   

   新增cuda_rasterizer/bindings.cpp文件，用pybind11暴露Python接口示例：

   #include <pybind11/pybind11.h>
   #include "rasterizer.h"
   
   namespace py = pybind11;
   
   PYBIND11_MODULE(_C, m) {
       // 根据实际业务补充类/函数绑定
       py::class_<Rasterizer>(m, "Rasterizer")
           .def(py::init<>())
           .def("forward", &Rasterizer::forward)
           .def("backward", &Rasterizer::backward);
   }
   

2. 修改setup.py：
   简化逻辑，仅负责调用CMake构建并打包生成的扩展：

   from setuptools import setup, find_packages
   from setuptools.command.build_ext import build_ext
   import subprocess
   import os
   
   current_dir = os.path.dirname(os.path.abspath(__file__))
   
   class CMakeBuildExt(build_ext):
       def run(self):
           build_dir = os.path.join(current_dir, 'build')
           os.makedirs(build_dir, exist_ok=True)
           # 执行CMake配置与构建
           subprocess.check_call(['cmake', '..'], cwd=build_dir)
           subprocess.check_call(['make'], cwd=build_dir)
   
   setup(
       name="diff_surfel_rasterization",
       packages=find_packages(),
       version='0.0.1',
       cmdclass={'build_ext': CMakeBuildExt},
       # 确保扩展模块被打包
       package_data={
           'diff_surfel_rasterization': ['_C.so']
       },
       include_package_data=True
   )
   

方案二：让setuptools链接CMake生成的库 #

若坚持保留原CMake生成动态库的逻辑，需调整setup.py让setuptools正确链接该库：

1. 修改CMakeLists.txt：
   将生成的库改为静态库，避免运行时依赖问题：

   # 将原add_library改为静态库
   add_library(CudaRasterizer STATIC
       # ... 原文件列表 ...
   )
   

2. 修改setup.py：
   添加绑定代码文件，让CUDAExtension有编译目标，并链接CMake生成的静态库：

   from setuptools import setup, Extension
   from torch.utils.cpp_extension import CUDAExtension, BuildExtension
   import subprocess
   import os
   
   current_dir = os.path.dirname(os.path.abspath(__file__))
   build_dir = os.path.join(current_dir, 'build')
   lib_path = os.path.join(build_dir, 'libCudaRasterizer.a')
   
   class CMakeBuildExt(BuildExtension):
       def run(self):
           os.makedirs(build_dir, exist_ok=True)
           subprocess.check_call(['cmake', '..'], cwd=build_dir)
           subprocess.check_call(['make', 'clean'], cwd=build_dir)
           subprocess.check_call(['make'], cwd=build_dir)
           super().run()
   
   setup(
       name="diff_surfel_rasterization",
       packages=['diff_surfel_rasterization'],
       version='0.0.1',
       cmdclass={'build_ext': CMakeBuildExt},
       ext_modules=[
           CUDAExtension(
               name="diff_surfel_rasterization._C",
               # 添加绑定代码文件
               sources=[os.path.join(current_dir, 'cuda_rasterizer', 'bindings.cpp')],
               extra_compile_args={
                   "nvcc": ["-I" + os.path.join(current_dir, "third_party/glm/")],
                   "cxx": ["-I" + os.path.join(current_dir, "cuda_rasterizer/")]
               },
               extra_link_args=[lib_path]
           )
       ],
   )
   

   同样需要编写bindings.cpp文件（同方案一示例）暴露Python接口。

内容的提问来源于stack exchange，提问作者Peppasaur