NVlabs/CuTe

解決的問題

PyCuTe 提供了 NVIDIA 的 CUTLASS 3.x 中使用的階層式佈局與張量代數的純 Python 參考實作。它讓開發者能夠學習複雜的代數運算、原型化新的資料轉換，並為 C++ 和 DSL 實作產生測試向量，而無需依賴 GPU。

工作原理

核心上，PyCuTe 將 Layout 定義為一個函數，使用 Shape（定義域）和 Stride（映射邏輯）將座標映射到偏移量。它實現了一套代數運算來操作這些佈局，包括：

• 合併 (Coalesce)：將佈局簡化至最少的模式數量。
• 組合 (Composition)：透過另一個佈局索引一個佈局。
• 補集 (Complement)：識別缺失的模式以填滿對應域。
• 邏輯除法/乘法 (Logical Divide/Product)：處理圖案的分塊與重複。
• 反轉/零空間 (Inversion/Nullspace)：分析並反轉映射。

它也支援 F2（XOR-swizzle）位移，以協助視覺化和管理共用記憶體銀行衝突。一個輕量的 Tensor 層將這些佈局與 Accessor 配對，提供資料模型。

適用對象

• 使用 CUTLASS 3.x 或 CuTe DSL 的開發者。
• 原型化 GPU 張量運算與記憶體佈局的工程師。
• 研究階層式佈局代數的學者。

特色

• 與 GPU 無關：在純 Python（3.10+）中執行，無需任何硬體需求。
• 視覺化工具：內建工具可將佈局以 ASCII 表格、彩色 SVG 或 TikZ/PDF 格式呈現。
• 完整的代數實作：實現了 CuTe 白皮書中描述的完整佈局代數。
• 支援 swizzle：特別處理 XOR-swizzling 以視覺化銀行衝突。