tracel-ai/cubecl

解決的問題

CubeCL 提供了一種在 Rust 中一次撰寫高性能計算內核，並在多個硬體平台（NVIDIA、AMD、Apple、Vulkan、WebGPU 和 CPU）上執行的方法，而無需為 CUDA 或 WGSL 等不同著色語言撰寫獨立程式碼。

工作原理

它作為 Rust 語言擴充和即時（JIT）編譯器運作。透過使用 #[cube] 屬性，將 Rust 函數轉換為中間表示（IR），然後按需編譯為目標平台的原生語言（例如 CUDA C++、HIP、Metal 或 SPIR-V）。

為了在多種硬體上保持最佳效能，它採用四軸平行模型（Vector、Plane、CubeDim 和 CubeCount），使內核能夠自適應，而非硬編碼為特定硬體常數（如 warp 大小）。

適用對象

專為開發科學計算函式庫、深度學習框架（如 Burn）以及高性能計算內核的開發者設計，他們希望在不犧牲低階 GPU/CPU 程式設計原始效能的前提下，享受 Rust 的安全性與可組合性。

主要亮點

• 多平台支援：可編譯為 CUDA、HIP、Metal、SPIR-V、WGSL 和 CPU SIMD。
• 編譯時（comptime）：允許在首次編譯時修改編譯器 IR，實現指令特化與迴圈展開。
• 自動調校：在執行時自動對不同內核設定進行基準測試，以找到目前硬體最高效的版本。
• 自動向量化：將高階向量類型降低為平台原生 SIMD 指令，並自動支援廣播。
• JIT 編譯：僅編譯實際啟動的特定內核變體，相比提前編譯可大幅縮小二進位體積。