CUB 和 Thrust 等 C++ 库提供高级构建块，使 NVIDIA CUDA 应用和库开发者能够编写跨架构可移植的光速代码。许多广泛使用的项目 (例如 PyTorch、TensorFlow、XGBoost 和 RAPIDS) 都使用这些抽象来实现核心功能。
Python 中缺少相同的抽象。其中包括 CuPy 和 PyTorch 等高级数组和张量库，以及 numba.cuda 等低级内核创作工具。但是，由于缺乏“基础模组”，Python 库开发者不得不使用 C++ 来实现自定义算法。
介绍 cuda.cccl
cuda.cccl 为 CUDA 核心计算库 CUB 和 Thrust 提供 Python 接口。现在，您无需使用 C++ 或从头开始编写复杂的 CUDA 核函数，即可编写跨不同 GPU 架构提供最佳性能的算法。
cuda.cccl 由两个库组成：
parallel提供作用于整个数组、张量或数据范围 (迭代器) 的可组合算法。cooperative通过提供作用于块或线程束的算法，支持您快速、灵活地编写numba.cuda核函数。
本文将介绍 parallel 库。
一个简单的示例：自定义归约
为了说明 cuda.cccl 的功能，以下是一个玩具示例，它结合了 parallel 中的功能片段，计算了 1 – 2+ 3 – 4+ … N 之和。
查看完整内容代码示例。
# define some simple Python functions that we'll use later
def add(x, y): return x + y
def transform(x):
return -x if x % 2 == 0 else x
# create a counting iterator to represent the sequence 1, 2, 3, ... N
counts = parallel.CountingIterator(np.int32(1))
# create a transform iterator to represent the sequence 1, -2, 3, ... N
seq = parallel.TransformIterator(counts, transform)
# create a reducer object for computing the sum of the sequence
out = cp.empty(1, cp.int32) # holds the result
reducer = parallel.reduce_into(seq, out, add, initial_value)
# compute the amount of temporary storage needed for the
# reduction, and allocate a tensor of that size
tmp_storage_size = reducer(None, seq, out, size, initial_value)
tmp_storage = cp.empty(temp_storage_size, cp.uint8)
# compute the sum, passing in the required temporary storage
reducer(tmp_storage, seq, out, num_items, initial_value)
print(out) # out contains the result
速度快吗？
让我们使用 parallel 对刚刚构建的算法进行计时，同时使用 CuPy 的数组运算进行原生实现。这些计时在 NVIDIA RTX 6000 Ada 架构上完成。请参阅全面的基准测试脚本。
以下是使用数组运算的计时：
seq = cp.arange(1, 10_000_000)
%timeit cp.cuda.runtime.deviceSynchronize(); (seq * (-1) ** (seq + 1)).sum(); cp.cuda.runtime.deviceSynchronize()
690 μs ± 266 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
以下是使用我们使用 parallel 构建的算法的计时：
seq = TransformIterator(CountingIterator(np.int32(1)), transform_op)
def parallel_reduction(size):
temp_storage_size = reducer(None, seq, out_tensor, size, initial_value)
temp_storage = cp.empty(1, dtype=cp.uint8)
reducer(temp_storage, seq, out_tensor, size, initial_value)
return out_tensor
%timeit cp.cuda.runtime.deviceSynchronize(); parallel_reduction(); cp.cuda.runtime.deviceSynchronize()
28.3 μs ± 793 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
我们看到，结合 parallel 迭代器和算法的方法比天真的方法更快。
加速从何而来？
CuPy 中的许多核心 CUDA 运算都使用 CUB 和 Thrust，即 parallel 向 Python 提供的相同 C++ 库。为什么使用 parallel 可以看到更好的性能？
这里没有魔法。parallel 为您编写通用算法提供了更多的控制力和灵活性。具体而言：
- 内存分配更少：
parallel的一个主要优势是能够使用 tg_ 18 和 tg_ 19 等迭代器作为reduce_into等算法的输入。迭代器可以表示序列，而无需为其分配内存。 - 显式内核融合：以这种方式使用迭代器，将所有工作“融合”到一个内核中 — — 原生 CuPy 代码段会启动四个内核 (看看能否找到所有内核) 。我们使用
parallel构建的自定义算法将它们合并为单个归约。请注意，这不同于例如通过 tg_ 22 完成的融合。它是显式的，这意味着您可以控制事物的融合方式，而不是隐式的，编译器控制融合。此控制可让您融合编译器可能无法融合的运算。 - Python 用度更低：最后，
parallel是一个较低级别的库，是基于底层 CUB/ Thrust 功能的薄层。通过使用parallel，您无需在调用设备代码之前翻阅多个 Python 层。
cuda.ccl 适用于哪些人？
cuda.cccl 的目标不是取代 CuPy、PyTorch 或任何现有的 Python 库。相反，它旨在更轻松地实现此类库，或扩展这些库，并更高效地使用 CuPy 数组或 PyTorch 张量实现自定义运算。具体而言，请在以下情况下查看 cuda.cccl：
- 构建可由更简单的算法 (例如
reduce、tg_ 29、transform等) 组成的自定义算法。 - 创建序列并对其进行操作，无需为其分配任何内存 (使用迭代器) 。
- 定义和操作由更简单的数据类型组成的自定义“结构化”数据类型。我们提供了一个如何执行此操作的示例。
- 使用 CUDA C++ 并编写自定义 Python 绑定到 Thrust 或 CUB 抽象。借助
cuda.cccl，您可以直接从 Python 使用这些功能。
cuda.cccl API 有意设置为低级别，并与底层 C++ 设计非常相似。这使得它们尽可能保持轻量级和低用度，同时提供 CuPy 和 PyTorch 等许多库内部使用的强大构建块。
后续步骤
现在您已经了解了 cuda.cccl 及其功能，不妨尝试一下。安装是单个 pip 命令。
pip install cuda-cccl
接下来，查看我们的文档和示例，并在我们的 GitHub 存储库中报告任何问题或功能请求。
