cuda-oxide: Rustの安全性と表現力をNVIDIA GPUへ

長年、GPUプログラミングのゴールドスタンダードはCUDA C++でした。これは強力ですが、単一のポインタエラーがサイレントなデータ破損や致命的なカーネルパニックを引き起こす可能性のある、危険な言語です。Rustエコシステムでは、cudarcのようなクレートを介してCUDAとインターフェースする様々な試みがなされてきましたが、これらは通常、FFI（Foreign Function Interface）やCMake、nvccといった外部ビルドシステムに依存しており、開発体験に摩擦が生じていました。

NVIDIAは現在、このギャップを埋めるために設計された実験的なRust-to-CUDAコンパイラであるcuda-oxideを発表しました。cuda-oxideは、ラッパーやドメイン固有言語（DSL）として機能するのではなく、カスタムのrustc codegen backendであり、純粋なRustを直接PTX（Parallel Thread Execution）にコンパイルします。これにより、開発者はRustの型システムと所有権モデルを使用してSIMT（Single Instruction, Multiple Threads）カーネルを記述できるようになります。

GPUカーネルのための新しいパラダイム

cuda-oxideの核心は、従来のRust-CUDAワークフローに関連する「シリアライゼーションの障壁」とビルド時のオーバーヘッドを排除することを目指しています。コンパイラパイプラインに直接統合されることで、ホストコードとデバイスコード間のよりシームレスな移行が可能になります。

技術的な主要な柱

  • ネイティブSIMTコンパイル: C++カーネルのための高レベルなオーケストレーターとしてRustを使用する以前の試みとは異なり、cuda-oxideは標準的なRustコードをコンパイルします。これは、開発者がトレイト、ジェネリクス、およびクロージャをGPUカーネル内で直接活用できることを意味します。
  • 非同期GPUプログラミング: このプロジェクトはDeviceOperationモデルを導入し、GPUの作業を遅延グラフ（lazy graphs）として構成できるようにします。これらの操作はストリームプールを横断してスケジューリングされ、Rustネイティブの.await構文を使用して待機（await）することができ、非同期プログラミングの利便性をGPUにもたらします。
  • コンパイラパイプライン: アーキテクチャは洗練されており、Pliron（MLIRに似たもの）と呼ばれるカスタムIRと、RustのMid-level Intermediate RepresentationをGPU実行に適した形式に変換するためのMIRインポーターを利用しています。このパイプラインは、最終的にrustc-codegen-cuda backendを介してコードをPTXに低減（lower）します。

安全性の挑戦: Rust vs. GPUハードウェア

cuda-oxideにおいて最も議論されている側面の一つは、、どのようにRustの安全性保証を、大規模な並列性と共有メモリへのアクセス用に設計されたハードウェア上で扱うかという点です。プロジェクトのドキュメントに記載されているように：

GPUカーネルは、数千ののスレッドが同時に同じメモリを参照する環境で動作します... 所有権チェッカー（borrow checker）は、このような設計にはなっていません。

これに対処するため、cuda-oxideは階層的な安全性モデルを実装しています：

  1. 構成による安全性: 一般的なケース（一つのスレッドが特定の要素に書き込む場合）は、unsafeブロックを必要とせずに安全に処理されます。
  2. 文書化された契約: 共有メモリへのアクセスやwarp shuffleなどのより複雑な操作にはunsafeブロックが必要ですが、開発者を導くための文書化された契約が伴います。
  3. 手動制御: Tensor Memory Accelerators (TMA)やtensor coresのような最先端の機能は、基盤となるハードウェアの極端な複雑さを反映して、完全に手動での制御が残ります。

コミュニティの視点とトレードオフ

この発表は、GPU上でのRustの有用性とパフォーマンスについて、システムプログラマの間で大きな議論を巻き起こしています。

パフォーマンスとオーバーヘッド

一部の開発者は、Rustのランタイムオーバーヘッドについて懸念を表明しています。具体的には、Rustの安全性の要である配列の境界チェック（bounds checking）が、レジスタ圧迫を引き起こし、カーネルの占有率（occupancy）や並列実行（concurrency）を低下させるのではないかという疑問があります。パフォーマンスが重要なカーネルでは、これらのチェックを回避する（opt-out）機能が極めて重要になります。

「クローズドソース」のジレンマ

批判的な意見を持つ人々は、言語インターフェースは改善されるものの、基盤となるインフラストラクチャは依然としてプロプライエタリ（独占的）なままであると指摘しています。あるコメントでは、「ツールチェーンは最終的にNVIDIAのクローズドソースのドライバーとランタイムバイナリに依存しており、APIの『Rust化』は、ベンダーロックインやオープンソースのコンパイラスタックの欠如という根本的な問題を解決しない」と述べられています。

代替手段との比較

コミュニティは、cuda-oxideも他の新興技術と比較しています：

  • Mojo: 一部の人は、Mojoがオープンソースのコンパイラとマルチバックエンドサポートを約束しているため、AIハードウェアにおいてより野心的な代替案になると主張しています。
  • Slang: 他の人は、cuda-oxideがGPU作業のための現代的で汎用的な言語を提供することで、Slangのような特殊なシェーダー言語の必要性を減らすのではないかと考えています。
  • Cudarc: すでにcudarcを使用しているユーザーにとって、cuda-oxideは、CMakeを介した高コストなnvccの呼び出しを不要にし、ビルド時間を大幅に削減できる潜在的な「ドロップイン・リプレイスメント」となる可能性があります。

結論

Cuda-oxideは現在、初期段階のアルファ版（v0.1.0）であり、チームはユーザーに対してバグやAPIの破壊的変更を予想しておくよう警告しています。しかし、このプロジェクトは、GPUプログラミングをよりアクセシブルでエラーの起こりにくいものにするための重要な一歩です。Rustの安全性モデルの境界をGPUへと広げることで、NVIDIAは、現代的な型システムによる生産性の向上は、そのシステムをSIMTハードウェアにマッピングすることの複雑さよりも上回るということに賭けています。

Sources