NVIDIA、GPUカーネルをRustへ自動移植するAI手法を公開

米NVIDIAは、GPUカーネルライブラリ「TileGym」のCUDAタイル演算コード（PythonおよびTriton実装）を、型安全性とメモリ安全性を備えたRust言語環境（cuTile Rust）へ自動変換・検証するマルチエージェント型AIスキルを技術ブログで公開した。共通の中間表現による構造差分検査と機械的な判定器を組み合わせることで、人間によるコード確認に頼らず、高い精度と実行性能を両立した。

共通中間表現による構造検証とエージェント分業

今回の移植パイプラインでは、Python（cuTile Python）、Triton（Triton-TileIR）、Rust（cuTile Rust）の3つのフロントエンドが、共通の低レベル中間表現「CUDA Tile IR」を介して同一の最適化コンパイラ「tileiras」に接続するアーキテクチャが採用されている。

この設計により、移植作業はアルゴリズムの再最適化ではなく「同一のタイル処理をより安全なホスト言語で型安全に再表現する」問題へと単純化される。変換の正当性は、生成された中間表現と元の参照実装の中間表現の差分を直接比較することで、テスト実行前から構造的に検証できる。

エージェント群の分業体制は、自然言語による曖昧なやり取りを完全に排除し、定義された成果物ファイルと終了コードのみで連携するステートマシンとして構築された。

• Analyzer（仕様解析）​: 参照カーネルの中間表現をダンプし、データ型や定数、許容誤差を定義した仕様書を出力する。
• Kernel Writer（カーネル記述）​: ホスト側のコードから隔離された環境でデバイスコードのみを記述し、49項目のコーディング規則と構造差分検査を自律実行する。
• Host/FFI Builder（連携基盤構築）​: C言語ABIを介してPythonテストフレームワークからRustカーネルを呼び出すラッパーを構築し、テストスイートを実行する。
• Performance Validator（性能評価）​: CUPTI（NVIDIAのGPUプロファイリングおよび性能計測ツール群）を用いてカーネル単体の実行時間を測定し、性能基準の適合を判定する。
• ​不具合診断スペシャリスト​: 失敗時のみ起動し、中間表現の不一致原因がコード生成ミスか上流コンパイラの不具合かを診断する。

実証性能と開発者への影響

NVIDIA DGX B200環境でTileGymの全24演算子（約40種類のGPUカーネル）を移植したベンチマーク結果では、cuTile Python実装と比較して幾何平均で99.5%という同等の実行性能を達成した。

全演算子が目標基準（参照比95%以上）をクリアしたほか、正規化処理や要素単位の演算など約3分の1のカーネルではPython版を上回る実行速度を記録した。さらに、エージェントのコンテキストを各工程に必要な最小限の参照文書に絞り込んだことで、変換に要したトークン費用は平均して約半分に抑制された。

項目	実績数値
対象演算子数	24演算子（約40カーネル）
対Python版 幾何平均性能比	99.5%
性能目標達成率（95%以上）	100%（24/24演算子）
トークン消費量削減	約50%削減

本手法が開発者にもたらす最大の価値は、AIによるプログラミング支援を「人間が目視でレビューして信頼する」方式から、「中間言語の差分検査と自動計測によって数学的・機械的に正当性を保証する」方式へと転換した点にある。

システム開発におけるAIエージェント活用の新地平

ハイパフォーマンスコンピューティングやAI推論基盤の領域では、メモリ安全性や並行処理性能に優れたRustの採用が進んでいるが、高度にチューニングされたGPUカーネルを手作業でRustへ書き換える作業は極めて高い専門知識を要求され、普及の障壁となっていた。

今回の取り組みは、LLMエージェントに厳格なルールセットと機械的検証器を与えることで、専門性の高いシステムプログラミングや低レイヤーの最適化コード移植を無人で安全に遂行できることを実証した。今後、大規模な既存コードベースの安全な言語への移行を後押しするモデルケースになるとみられる。

引用元：NVIDIA Technical Blog