tracel-ai/cubecl

何を解決するか

CubeCLは、Rustで一度だけ高性能なコンピュートカーネルを記述し、NVIDIA、AMD、Apple、Vulkan、WebGPU、CPUなど複数のハードウェアプラットフォームで実行できるようにする仕組みを提供します。CUDAやWGSLなどの異なるシェーダ言語で別々のコードを書く必要がありません。

動作方法

これはRust言語拡張および即時（JIT）コンパイラとして機能します。#[cube]属性を使用することで、Rust関数を中間表現（IR）に変換し、必要に応じてターゲットプラットフォームのネイティブ言語（例：CUDA C++、HIP、Metal、SPIR-V）にコンパイルします。

多様なハードウェアで最高のパフォーマンスを維持するために、4軸並列モデル（Vector、Plane、CubeDim、CubeCount）を採用しており、ワープサイズなどの特定のハードウェア定数にハードコードされるのではなく、カーネルを適応可能にしています。

対象ユーザー

科学計算ライブラリ、ディープラーニングフレームワーク（例：Burn）、高性能コンピュートカーネルを開発する開発者向けです。低レベルのGPU/CPUプログラミングのパフォーマンスを犠牲にすることなく、Rustの安全性と組み合わせのしやすさを享受したい方におすすめです。

特徴

• マルチプラットフォーム対応：CUDA、HIP、Metal、SPIR-V、WGSL、CPU SIMDにコンパイル可能。
• コンパイル時（comptime）：最初のコンパイル時にコンパイラIRを変更することで、命令の特殊化やループ展開が可能。
• 自動チューニング：実行時に異なるカーネル設定をベンチマークし、現在のハードウェアに最適なバージョンを自動的に検出。
• 自動ベクトル化：高レベルのベクトル型をプラットフォームのネイティブSIMD命令に下げる際、自動的にブロードキャストを処理。
• JITコンパイル：実際に起動された特定のカーネルバリアントのみをコンパイルするため、事前コンパイルに比べてバイナリサイズを削減。