CuPy 第三方库集成解析:CCCL、Jitify 与 DLPack 子模块的构建与应用
【免费下载链接】cupyNumPy & SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy
导读
CuPy 是面向 GPU 的 NumPy/SciPy 兼容加速库,其核心性能与跨框架互操作能力在很大程度上依赖三个上游第三方项目:CCCL(Thrust、CUB、libcu++ 的统一家园)、Jitify(运行时 CUDA C++ 代码 JIT 编译工具)以及DLPack(跨框架内存张量交换标准)。本文以仓库中 third_party/README.md 为骨架,深入 .gitmodules、install/cupy_builder/install_build.py、install/cupy_builder/_preflight.py 等源码,剖析这三个子模块如何通过 git submodule 被拉取、在构建期被探测与集成、并在运行期支撑 CuPy 的 CUB 原语、运行时 JIT 与零拷贝张量交换等关键能力。
一、Third Party 目录概览:子模块与捆绑头文件的分工
CuPy 仓库根目录下存在third_party/目录,其中每个文件夹都是一个git submodule(git 子模块)。当前仓库的实际状态如下:
| 子模块目录 | 对应上游项目 | 用途 |
|---|---|---|
third_party/cccl | NVIDIA CCCL | 提供 Thrust、CUB、libcu++ 头文件,构建期与运行期均被使用 |
third_party/jitify | NVIDIA Jitify | 为构建cupy.cuda.jitify模块提供 Jitify 头文件 |
third_party/dlpack | DLPack 项目 | 为构建cupy._core.dlpack模块提供 DLPack 头文件 |
third_party/xsf | SciPy 的 xsf | 特殊函数库(README 未列,但被构建系统一并管理) |
注意:从源码树看,
third_party/下的cccl、jitify、dlpack、xsf目录在仓库中均为空目录(子模块内容未随主仓库检出),这是 git submodule 的正常状态——内容需通过子模块命令单独拉取。README 中标注的xsf虽未在文中说明,但 .gitmodules 与构建脚本均将其纳入同一子模块管理框架。
从.gitmodules可以清晰看到子模块的登记信息:
[submodule "third_party/cccl"] path = third_party/cccl url = https://github.com/NVIDIA/cccl.git [submodule "third_party/jitify"] path = third_party/jitify url = https://github.com/NVIDIA/jitify.git [submodule "third_party/dlpack"] path = third_party/dlpack url = https://github.com/dmlc/dlpack.git [submodule "third_party/xsf"] path = third_party/xsf url = https://github.com/scipy/xsf.git四个子模块均指向上游官方仓库,CuPy 并不维护自己的 fork,而是通过**固定 commit(gitlink)**锁定版本,保证构建的可复现性。
二、CCCL:Thrust、CUB 与 libcu++ 的统一入口
2.1 CCCL 是什么
CCCL 是 NVIDIA 维护的统一开源项目,也是 Thrust、CUB、libcu++ 三个库的新家园。README 明确说明:cccl目录是 CCCL 项目的 git submodule,其头文件在 CuPy 的构建期和运行期(build- and run-time)都会被使用。
- CUB:提供 block/device 级并行原语(如
DeviceScan、DeviceSegmentedSort、DeviceReduce),是 CuPy 排序、扫描、归约类内核的高性能后端; - Thrust:提供 STL 风格的高层并行算法,被
cupy.cuda.thrust包装后供 Python 层调用; - libcu++:提供 CUDA 设备端的 C++ 标准库扩展,例如在 CuPy 编译内核时对
cuda::std命名空间的支持。
2.2 构建期如何被集成
CuPy 的构建系统在 install/cupy_builder/install_build.py 中显式设定了 CUB/Thrust 头文件的搜索优先级:
# For CUB, we need the complex and CUB headers. The search precedence for # the latter is: # - for CUDA: CuPy's CUB (and Thrust) bundle # - for ROCm: built-in CUB # Note that starting CuPy v8 we no longer use CUB_PATH, and starting v13 # we no longer use Thrust/CUB bundled in CUDA. # for <cupy/complex.cuh> cupy_header = os.path.join( cupy_builder.get_context().source_root, 'cupy/_core/include') global _jitify_path _jitify_path = os.path.join(cupy_header, 'cupy/_jitify') global _cub_path if rocm_path: _cub_path = os.path.join(rocm_path, 'include', 'hipcub') ... else: # all bundled together under cccl _cub_path = os.path.join(cupy_header, 'cupy/_cccl/cub') _thrust_path = os.path.join(cupy_header, 'cupy/_cccl/thrust') _libcudacxx_path = os.path.join(cupy_header, 'cupy/_cccl/libcudacxx')关键点:
- CUDA 平台:优先使用 CuPy 自带的、由
third_party/cccl生成的捆绑头文件(打包进cupy/_core/include/cupy/_cccl/),不再依赖 CUDA Toolkit 中捆绑的旧版 Thrust/CUB(自 v13 起废弃),也不再接受用户通过CUB_PATH环境变量指定(自 v8 起废弃); - ROCm/HIP 平台:改用
hipCUB(路径rocm/include/hipcub),若缺失则直接报错要求先安装 hipCUB; - C++ 标准:因为 CCCL 3.x(面向 CUDA)要求 C++17,构建脚本在 Linux 上追加
-std=c++17,在 Windows 上追加/std:c++17(见 install_build.py)。
2.3 版本探测机制
构建系统通过 check_cub_version 探测 CUB 版本:
- 优先尝试编译一段包含
<cub/version.cuh>的小程序,读取CUB_VERSION宏; - 若编译失败(例如子模块尚未检出),则回退到进入
third_party/cccl目录执行git describe --tags,从 git tag 解析版本号,并将 tag 归一化为与CUB_VERSION相同的格式(如v1.9.0→1.9.0); - 最终通过
CUPY_CUB_VERSION_CODE宏注入编译,同时缓存到全局变量,供get_cub_version()查询。
2.4 运行期如何被使用
构建出的扩展模块直接以 C++ 包装器调用 CUB/Thrust:
- cupy/cuda/cupy_cub.cu 通过
#include "cupy_cub.h"引入 CUB 模板,并被 cupy/cuda/cub.pyx 以cdef extern from 'cupy_cub.h'方式声明外部接口; - cupy/cuda/cupy_thrust.cu 通过
#include "cupy_thrust.h"包装 Thrust 算法,被 cupy/cuda/thrust.pyx 调用。
同时 setup.py 将cupy/cuda/cupy_thrust.cu、cupy/cuda/cupy_cub.cu等 .cu 源文件列入cupy_package_data,确保它们在构建 wheel/sdist 时被打包。
三、Jitify:运行时 CUDA 源码 JIT 编译
3.1 Jitify 的角色
Jitify 是 NVIDIA 提供的、在运行时对 CUDA C++ 源码进行 JIT(即时)编译的库。CuPy 将其作为子模块引入,用于构建cupy.cuda.jitify模块(见 install/cupy_builder/_features.py 中CUDA_jitifyfeature 的定义)。它支撑了 CuPy 的 JIT 融合内核、RawKernel 运行时编译等场景。
3.2 构建期集成细节
- 头文件路径:构建期将
cupy/_core/include/cupy/_jitify作为 include 目录(见 install_build.py),并在编译扩展模块时把./cupy/_core/include/cupy/_jitify/jitify.hpp注册为依赖文件(见 cupy_setup_build.py); - 版本探测:Jitify 没有 tag/branch 命名规范,check_jitify_version 只能进入
third_party/jitify执行git rev-parse --short HEAD,以短 commit hash 作为版本标识,并通过CUPY_JITIFY_VERSION_CODE宏注入; - 运行期暴露:
cupy.cuda.jitify模块(cupy/cuda/jitify.pyx)对外提供get_version()等 API,内部通过jitify::detail命名空间读取编译期写入的版本字符串,若探测失败则返回-1表示未知版本。该模块还会缓存 Jitify 解析出的 CUDA 头文件到磁盘,加速后续 JIT 编译。
四、DLPack:跨框架零拷贝张量交换
4.1 DLPack 的意义
DLPack 是一个开放的内存张量结构标准,用于在不同深度学习框架(PyTorch、MXNet、CuPy 等)之间零拷贝交换张量数据。README 将其描述为构建cupy._core.dlpack模块所需的子模块。
README 原文将 DLPack 一条误写为 "Including theJitifyheader...",从构建代码可见实际引入的是 DLPack 头文件,此处按构建事实理解为笔误。
4.2 实现与使用
- 构建期:
COMMON_dlpackfeature 定义了模块名与头文件路径cupy/_dlpack/dlpack.h(见 _features.py),编译时依赖./cupy/_core/include/cupy/_dlpack/dlpack.h(见 cupy_setup_build.py); - 运行期:核心实现在 cupy/_core/dlpack.pyx,对外提供:
toDlpack(array, ...):将cupy.ndarray导出为 DLPack 张量,底层由ndarray.__dlpack__协议支撑(见 dlpack.pyx);fromDlpack(dltensor):从 DLPack 张量零拷贝构造cupy.ndarray(见 dlpack.pyx);DLPackMemory:将 DLPack 内存块包装为 CuPy 的BaseMemory,并对过新的 DLPack 主版本号抛出BufferError,且同一 DLPack 对象只能被消费一次(见 dlpack.pyx)。
这套实现意味着用户可以通过 DLPack 协议在 CuPy 与其它支持该协议的框架间传递 GPU 张量而无需显式拷贝。
五、子模块的检出与构建前置检查
由于third_party/下是空的 git submodule,直接从源码构建前必须完成子模块初始化。仓库提供了强制的前置检查(preflight check),实现在 install/cupy_builder/_preflight.py:
- 构建脚本遍历
third_party/cccl、third_party/jitify、third_party/dlpack、third_party/xsf四个目录; - 若目录存在但为空、且当前处于 git 仓库中,则中止构建并输出提示:
$ git submodule update --init- 若当前不是 git 仓库(例如从 sdist 构建),则跳过该检查——因为 sdist 发布包中已包含编译所需的头文件,不再需要子模块。
注意:当前环境下的
third_party/子模块目录为空,若要在本地从源码构建 CuPy,必须先执行git submodule update --init(或--recursive)拉取全部子模块内容,再运行python setup.py之类的构建流程。
六、从源码到运行:第三方依赖的完整生命周期
综合以上各节,可将三个子模块的生命周期归纳如下:
- 拉取:
git submodule update --init按 .gitmodules 的 URL 与固定 commit 检出cccl、jitify、dlpack(以及xsf); - 前置检查:preflight_check 校验子模块是否就绪,空目录在 git 环境下直接报错引导初始化;
- 构建期集成:install_build.py 将 CCCL 的
cub/thrust/libcudacxx头文件目录加入 include 路径(CUDA 平台),将 Jitify 头文件拷入cupy/_core/include/cupy/_jitify,将 DLPack 头文件拷入cupy/_core/include/cupy/_dlpack; - 版本探测:通过编译探针(
cub/version.cuh)或 git 命令(git describe --tags/git rev-parse)解析版本,注入CUPY_CUB_VERSION_CODE、CUPY_JITIFY_VERSION_CODE等宏; - 打包分发:setup.py 将
cupy_cub.cu、cupy_thrust.cu等源码与全部cupy/_core/include头文件纳入 package_data,保证安装后的运行期可再编译(JIT 场景)或直接链接; - 运行期使用:
cupy.cuda.cub、cupy.cuda.thrust提供高性能并行原语;cupy.cuda.jitify提供运行时 CUDA 源码 JIT 编译;cupy._core.dlpack提供跨框架零拷贝张量互操作。
七、结语与进一步阅读
CuPy 通过 git submodule 机制优雅地复用了 NVIDIA 生态与 DLPack 社区的开源成果:CCCL 让 CuPy 获得与 CUDA 生态同步的高性能并行原语,Jitify 赋予其运行时 JIT 编译的灵活性,DLPack 则打通了 GPU 张量跨框架互操作的通道。理解这三者在构建期与运行期的分工,有助于排查源码构建问题(尤其是git submodule update --init缺失导致的空目录报错),也能更清晰地把握 CuPy 的性能与互操作能力边界。
若需深入,可继续阅读以下仓库文件:
- 子模块登记:.gitmodules
- 构建期头文件路径与 C++ 标准设置:install/cupy_builder/install_build.py
- 构建前置检查:install/cupy_builder/_preflight.py
- CUB 版本探测:install/cupy_builder/install_build.py#L487-L555
- Jitify 模块实现:cupy/cuda/jitify.pyx
- DLPack 模块实现:cupy/_core/dlpack.pyx
- CUB/Thrust C++ 包装器:cupy/cuda/cupy_cub.cu、cupy/cuda/cupy_thrust.cu
【免费下载链接】cupyNumPy & SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考