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 laterdef add(x, y): return x + ydef transform(x): return -x if x % 2 == 0 else x# create a counting iterator to represent the sequence 1, 2, 3, ... Ncounts = parallel.CountingIterator(np.int32(1))# create a transform iterator to represent the sequence 1, -2, 3, ... Nseq = parallel.TransformIterator(counts, transform)# create a reducer object for computing the sum of the sequenceout = cp.empty(1, cp.int32) # holds the resultreducer = parallel.reduce_into(seq, out, add, initial_value)# compute the amount of temporary storage needed for the# reduction, and allocate a tensor of that sizetmp_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 storagereducer(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 存儲庫中報告任何問題或功能請求。
?
?
