NVIDIA `cuda.cccl` 赋能 Python GPU 计算：CUB/Thrust 抽象直达，性能跃升

NVIDIA CUDA 应用程序和库的开发者现在可以通过 `cuda.cccl` 这个新引入的 Python 接口，直接利用 C++ 的 CUDA Core Compute Libraries CUB 和 Thrust，从而在 Python 环境中高效地编写跨架构的 CUDA 代码。该库旨在填补现有 Python 生态中，如 CuPy 和 PyTorch 等高级库与 Numba 等底层工具之间的“构建块”缺失。`cuda.cccl` 由 `parallel` 和 `cooperative` 两个子库组成，前者提供作用于整个数组或张量范围的可组合算法，后者则支持对数据块或线程束进行操作的算法，以编写更灵活的 Numba CUDA 内核。

通过使用 `cuda.cccl` 的 `parallel` 库，开发者能够以更精炼、更低开销的方式实现自定义算法。例如，一个自定义的交替加减序列求和示例显示，其性能远超直接使用 CuPy 数组操作的传统方法。这种性能提升主要归功于 `parallel` 库能够通过迭代器（如 `CountingIterator` 和 `TransformIterator`）代表序列而无需实际内存分配，并能将多个操作显式地融合进单个 CUDA 内核，从而减少了内存分配和 Python 调用栈的开销。

`cuda.cccl` 的目标并非取代现有 Python GPU 计算库，而是赋能开发者更便捷、更高效地实现自定义算法或扩展现有库的功能。当需要构建可组合的算法、无内存分配的序列操作，或直接从 Python 调用 C++ 的 Thrust/CUB 抽象时，`cuda.cccl` 是一个理想的选择。其低级别的 API 设计紧密映射底层 C++ 实现，确保了轻量级和高效率。

Delivering the Missing Building Blocks for NVIDIA CUDA Kernel Fusion in Python | NVIDIA Technical Blog

C++ libraries like CUB and Thrust provide high-level building blocks that enable NVIDIA CUDA application and library developers to write speed-of-light code that is portable across architectures.

NVIDIA Technical BlogAshwin Srinath