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 编译：仅编译实际启动的特定内核变体，相比提前编译可显著减小二进制体积。