deepseek-ai/DeepJIT

何を解決するか

DeepJITは、GPUおよびNPUカーネル用の軽量でヘッダーのみのC++20ランタイムを提供し、ランタイムでのJust-In-Time（JIT）コンパイルのプロセスを簡素化します。カーネルライブラリの作成者が、実行時におけるソースコードのコンパイル、バイナリのキャッシュ、ハードウェアデバイスへのロードのための独自インフラを構築する必要がなくなります。

動作方法

DeepJITは、主に2つのバックエンド（NVIDIA CUDA GPUとHUAWEI Ascend NPU）に対して共通のインターフェースを提供します。カーネルのライフサイクル全体を管理します：適切なツールチェーン（CUDA用にNVCC、Ascend用にBisheng/ld.lld）を使用してソースコードをコンパイルし、結果のバイナリをディスクおよびメモリにキャッシュして重複コンパイルを回避し、バックエンド固有のオプションでカーネルを起動します。

主な技術的メカニズムは以下の通りです：

• ハッシュ化とキャッシュ：ソースコード、トラッキングされたincludeファイル、コンパイラバージョン、構成オプションに基づいて一意のキャッシュキーを生成します。これにより、異なるプロセス、ユーザー、さらには共有ファイルシステムを介して異なるノード間でもバイナリを再利用できます。
• includeパーサー：行指向のスキャナが特定の角括弧形式のincludeをトラッキングし、依存ヘッダーの変更が再コンパイルをトリガーするようにします。
• PyTorch統合：PyTorchのCUDAまたはtorch_npuストリームと統合でき、pybind11を介してPythonから公開可能です。
• 遅延初期化：デバイスおよびコンパイラの検出は、ランタイムが実際に使用されるまで延期され、起動オーバーヘッドを削減します。

対象ユーザー

高性能なデバイスカーネルを書くC++およびPython拡張の作成者で、コンパイルおよびキャッシュの低レベル処理を実装せずに、ユーザーにJITコンパイル体験を提供したい人向けです。

特徴

• マルチバックエンド対応：NVIDIA CUDAとHUAWEI Ascend NPUの両方に対する統一されたワークフロー。
• 分散キャッシュ：POSIX準拠のファイルシステムを使用して、複数のユーザーおよびノード間でコンパイル済みカーネルを共有可能。
• PyTorch対応：PyTorchストリームとシームレスに統合され、get_jit()というPython APIを提供。
• コンパイル診断：デバッグおよび検査用にPTX/SASS（CUDA）またはアセンブリ（Ascend）をダンプできる機能。
• コンパイル後フック：CUDAバックエンドでは、CUBINファイルをキャッシュする前にカスタムPythonスクリプトを実行して変更可能。