warp-compile-time-optimizer

Warp のコールドスタート時のコンパイル時間

実行可能なコマンドから始める

プローブは対象をサブプロセスで実行し、一時ディレクトリにのみ書き込みます。ネットワーク、外部ツールサーバー、Warp のチェックアウトは不要です。

スクリプト	目的	引数
scripts/warp_compile_probe.py	分離したコールド／ウォーム時のコンパイルと起動を測定します。	measure [OPTIONS] -- COMMAND...。--help を使用してください。

利用可能な場合は run_script("scripts/warp_compile_probe.py", args=[...]) を使用してください。それ以外の場合は、以下の Python コマンドを使用してください。対象は最後まで実行される必要があります。

コンパイルの仕組み

Warp は個々のカーネルではなく、モジュールをコンパイルします。モジュールの識別情報は次のとおりです。

(live kernel & function set) x (module options) x (CUDA block_dim) x (generic instances)

各識別情報について、モジュール全体のコード生成とネイティブコンパイルが必要です。

コールドスタートのコストは、概ね次のとおりです。

number of distinct module identities you touch  x  size of each module

次の 2 つの方法で削減できます。

1. 識別情報の頻繁な変化を防ぎます。読み込み後にモジュールが変化すると、再コンパイルされます。
2. ソースの重複を防ぎます。3 種類のブロック次元で使用されるモジュールでは、すべてのカーネルが 3 回コンパイルされます。

ビルドが続行されるモジュールからカーネルを 1 つ削除しても、節約できるのはビルドの一部だけです。不要なモジュール識別情報を削除すれば、ビルド全体を省けます。

どちらも当てはまらない場合は、独立した CUDA ビルドを並行して実行します（CS-13）。これにより処理のタイミングは変わりますが、コンパイル量は変わらないため、経過時間で評価してください。

モジュールのオプションが固定されるタイミング

オプションには設定期限が 2 つあります。

1. モジュール作成時（Python のインポート時）。 モジュールは warp.config から enable_backward、max_unroll、lineinfo、deterministic、deterministic_max_records、compile_time_trace をコピーします。後からグローバル設定を変更しても、そのモジュールでは黙って無視されます。例外は default_grid_stride です。
2. 初回ロード時。 コンパイル前に、wp.set_module_options() または wp.get_module(name).options で既存のモジュールを変更します。ロード後にオプションを変更すると、新しい識別情報が作成され、モジュールが再ビルドされます（CS-3）。

期限を逃した場合	症状	コスト
インポート後に wp.config.* を設定	ハッシュが変わらず、オプションが黙って適用されない	効果がすべて失われます。しかも気づきにくい状態です
ロード後にモジュールのオプションを設定	2 つ目のハッシュが生成され、モジュールが 2 回ビルドされる	ビルドが 1 回余分に発生します。トレースで確認できます

オプションの変更後は、対象となるすべてのモジュールでハッシュが変わったことを確認してください。ハッシュが変わらない場合、オプションは適用されていません。

動作を維持する

ワークロードではなく、Warp のコードのコンパイル方法を変更してください。

効果を主張するために、カーネルを削除したり統合したりしないでください。一見重複しているように見えるステージも、所有権、エイリアシング、保持される出力、数値境界、API の動作を維持している場合があります。重複はモジュールレベルで解消してください。

すべての起動とその順序、次元、データ型、デバイス、ブロック次元、勾配、数値モード、動的動作／プラグイン動作、公開 API のシグネチャを維持してください。定義をモジュールスコープに移動する場合も、ログ、キャッシュアーティファクト、外部ツールに表示されるため、カーネル名は変更しないでください。

手順

1. 実際のコストを調べ、コンパイルが原因だと確認する

ユーザーが実際に待っているコマンドを尋ね、その後、キャッシュが空の状態で計測します。

python scripts/warp_compile_probe.py measure --samples 3 \
    --json baseline.json -- <the user's command>

このプローブでは、各サンプルに専用の WARP_CACHE_PATH、WARP_CACHE_ROOT、CUDA_CACHE_PATH ディレクトリを割り当て、モジュールタイマーを有効にして、カーネル起動を記録します。システムの一時ディレクトリ内にこれらのディレクトリを作成し、プローブ自身が削除するため、サンプルを分離する際にプロジェクト内へキャッシュを書き込んだり、再計測のために何かを削除したりする必要はありません。手書きのサンプルも同じ方法で分離してください。wp.clear_kernel_cache() や wp.clear_lto_cache() を使って、稼働中のキャッシュを消去しないでください。消去は分離されず、ほかのプロセスに影響する可能性があります。

ソースコードより先に、プローブの出力を確認してください。壁時計時間に占めるコンパイル時間が小さい場合は、実際のボトルネックを報告して終了します。ライブラリやテストでは、対象ワークロードのモジュールをコンパイルする最小限のコマンドを使ってください。

各モジュールがそれぞれ一度だけコンパイルされ、ハッシュの重複、ブロック次元のバリアント、LTO がない場合、構造上の無駄な再コンパイルはありません。これは冗長なビルドがないことを示すものであり、ビルドが過大でないことを示すものではありません。引き続き、キャッシュの再利用（CS-2）、逆方向コード生成（CS-10）、ループ展開（CS-11）、プリコンパイル済みヘッダー（CS-12）、複数の CUDA モジュールが残る場合の並列実行（CS-13）を確認してください。

各サンプルでは、生成したキャッシュを使ってコマンドも再実行します。ウォーム状態のモジュール処理時間がほぼゼロでない場合は、モジュール構造を変更する前に、キャッシュの再利用（CS-2）を調査してください。

2. 必要に応じて実行時のトレードオフを確認する

順序の修正、ライフサイクルごとのグループ化、オプションの事前設定は、確認せずに行ってください。fast_math、max_unroll、MathDx／タイル実装を変更する前に、確認してください。

これらの設定の一部はコンパイル時間を短縮できますが、コンパイル後のカーネルが遅くなったり、数値が変わったりすることがあります。編集から実行までの反復を速くすること（通常、カーネルが遅くなっても問題ありません）と、本番環境での起動を速くすること（通常、カーネルが遅くなるのは望ましくありません）のどちらを優先しますか？

ユーザーに確認できない場合：

• 数値を変えるオプションや実装の置き換えは行わないでください。
• 逆方向コード生成などの機能を削除する前に、実際に使われているかを確認してください。削除する内容、計測した効果、元に戻す方法を報告してください。
• グローバルな wp.config.* オプションは、ライブラリコードではなく、アプリケーションのエントリーポイントで設定してください。

辞退されたオプションと計測した効果を、ステップ 6 の記録に残してください。

変更の範囲を、根拠の範囲に合わせる

プロファイルから正当化できる変更は、計測対象のアプリケーションに対するものであり、共有ライブラリのあらゆる利用者に対するものではありません。また、リポジトリ内の検索では、リポジトリ外の呼び出し元や将来の呼び出し元を見落とします。たとえば、順方向のみのアプリケーションがあるからといって、別のアプリケーションが微分処理に使うソルバーライブラリ内で勾配を無効にしてよいことにはなりません。

ライブラリをインポートする前に、計測対象のプロセスに限ってオプションを設定してください。

import warp as wp

wp.config.enable_backward = False   # must precede the library import

import the_library

これにより、アプリケーションが読み込むすべてのモジュールにも設定が適用されます。インポート順序による予期しにくい問題については CS-10 で説明しています。ライブラリの変更でしか解決できない場合は、計測結果をメンテナーに共有し、契約をどうするか判断してもらってください。

3. ソースを読むのではなく、測定結果から診断する

プローブは、コンパイル済みモジュールの識別情報をすべて、名前、ハッシュ、デバイス、ブロック次元とともに出力し、複数回ビルドされたモジュールを示します。表示された内容に照らし合わせてください。

プローブの表示	意味	確認する箇所
1つのモジュール名に対して複数の ハッシュ	識別情報の変化：最初にロードされた後に、カーネルセット、オプション、またはジェネリックインスタンスが変わりました	CS-1、CS-3、CS-6
1つのモジュール名に対して複数の block_dim 値	ブロック次元ごとにモジュール全体が再コンパイルされています（CUDA）	CS-5
1つの機能に多数の1カーネルモジュール	モジュールごとの固定コストが繰り返されています	CS-4
カーネルごとにハッシュ名のモジュールがある	安定したカーネルで module="unique" が使われています	CS-9
モジュール時間とネイティブコンパイル時間の大きな差に加え、.lto 成果物がある	MathDx/LTO のセットアップ	CS-7
ウォーム状態になるはずの実行で (compiled) と表示される	キャッシュが再利用されていません	CS-2
モジュールがロードされた後に「Failed to find module」と表示される	CPU JIT の初回使用時に同時実行の競合が発生しています	CS-8
生成されたソースが大きく、再ビルドの問題はない	アンロールの上限	CS-11
何も微分しないモジュールに随伴コードがある	逆方向のコード生成	CS-10
全般的にコンパイルが遅い、または CUDA でツールキット 13 未満の場合に小さなモジュールがいくつかある	プリコンパイル済みヘッダーが無効になっているか、効果が出ていません	CS-12
独立した複数のモジュールがそれぞれ1回ずつビルドされ、overlap_factor が 1.0 に近い	ビルドが一度に1つずつ実行されています。並列ロードはデフォルトで無効です	CS-13
設定したオプションが何も変えず、そのモジュールのハッシュも変わらない	モジュール作成後に割り当てられたため、反映されていません	「モジュールのオプションが固定されるタイミング」
上記のいずれにも当てはまらない	ビルドの重複はありません。コストはビルド自体の規模によるものです	ステップ 6

references/mechanisms.md には、各メカニズムについて、確認方法、修正方法、制限事項、失敗モードを説明するセクションがあります。測定結果から該当するセクションだけを読んでください。

4. モジュールの境界を慎重に選ぶ

次の条件を共有する場合に限り、カーネルを1つのモジュールにまとめてください。

• ライフサイクル：同時に定義、ロード、無効化される。
• オプションセット：同じ fast_math、enable_backward、max_unroll、MathDx の設定が必要である。
• CUDA でブロック次元が安定している。

ライフサイクルが同じでも、安定したブロック次元が異なるカーネルは、共有するとそれぞれが2回コンパイルされるため、同じモジュールにまとめないでください。ライフサイクルが独立しているカーネルも分けてください。

ブロック次元が実行時に変わるカーネル（たとえば入力サイズから選ばれるカーネル）には安定した対応関係がないため、共有モジュール全体を追加のバリアントに巻き込むのではなく、専用のモジュールに分けてください。

ビルドを1つ減らせる変更のうち、最も影響の小さいものを優先してください。順序の修正やオプションの先行設定は、モジュール構成の再設計より低コストで安全です。モジュールごとの固定コストやブロック次元の重複が支配的な場合に限り、再編成してください。

wp.set_module_options() は、明示的な module="pkg.name" が指定されたカーネルではなく、呼び出し元の Python モジュールを対象とします。実際の Python モジュールを使うか、名前付きモジュールがロードされる前に更新してください。

wp.set_module_options({"enable_backward": False})  # at module scope

wp.get_module("pkg.name").options.update({"enable_backward": False})

wp.get_module() を wp.set_module_options(module=...) に渡さないでください。また、module="unique" を指定せずに @wp.kernel(module_options={...}) を使わないでください。カーネルごとの enable_backward=False には、CS-10 で説明するタイルモジュールの例外があります。オプションを変更するたびに、モジュールのハッシュを確認してください。

5. 検証

python scripts/warp_compile_probe.py measure --samples 3 \
    --json candidate.json -- <the same command>
python scripts/warp_compile_probe.py compare baseline.json candidate.json

compare は起動トポロジーの変更を拒否し、 max(ベースラインの1%、ベースラインMADの2倍) 以内の結果を 判定不能として扱います。

BUILDS OVERLAPPED の場合は、合計されたモジュールタイマーではなく、 コンパイルの経過時間でスケジューリングの変更を判断してください。必要な ウォームパスがその計測値を提供します。references/measurement.md を参照してください。

次に、プローブでは確認できない項目をチェックします。

• 数値出力を差分比較し、プロジェクトのテストまたはエントリポイントを実行します。
• enable_backward や境界を変更した後は、勾配パスを検証します。
• fast_math、max_unroll、MathDx、または実装を変更した後は、 定常状態のランタイムをベンチマークします。
• カバーされていない動的カーネル、dtype、プロファイル、モードを試します。

6. 結果の報告

references/measurement.md の「結果の報告」を読んでください。報告内容は次のとおりです。

すべての最適化を平易な言葉で説明してください。振る舞い、根拠、 効果を記述します。たとえば、「CS-3 を適用した」ではなく、「冗長な再ビルドを 避けるため、最初のロード前にモジュールオプションを移動した」と書きます。 CS-* ラベルは内部のナビゲーション用であり、ユーザー向けの説明ではありません。

オプション	測定値	採用しなかった理由	採用するために
pkg.solver の enable_backward=False	コールド −38%	ライブテープがこれらのカーネルを通過する	エントリポイントで設定し、随伴を再確認する
max_unroll=4	−2%、ノイズ範囲内	測定された利得なしに生成コードを変更する	—

• 変更前後の中央値とサンプル数;
• 元の不満に対して測った削減量と残存コスト;
• 修正されたメカニズム、除外されたもの、測定して却下されたもの、到達しなかったもの;
• トレードオフと検証されていない振る舞い;
• 却下された、または不完全にしか検証されていないオプションに関する上記の台帳;
• 修正を妨げていた制約を緩和した場合の測定値。

根拠が支持する内容のみを述べてください。「構造的な無駄な変更がない」ことは 「最適」や「これ以上削減不能」を意味しません。実用的な場合は却下したレバーを測定してください。 推定値には未検証とラベルを付けてください。利用可能な修正がないと報告する前に、 CS-13 を確認してください。モジュール分割の可能性については、まず同じオプションで 単一カーネルモジュールを測定し、繰り返し発生する固定コストを明らかにしてください。

トラブルシューティング

プローブをデバッグする前に、対象のコマンドを直接実行してください。キャッシュやノイズの問題については references/measurement.md を、メカニズム固有の障害については references/mechanisms.md を参照してください。

制限事項

2つのルールがどのような利得よりも優先されます。

• すべてのコールドサンプルで、Warp と CUDA の両方のキャッシュを分離してください。
• ロードが CPU を対象とする可能性がある場合は、device=None や混在デバイスリストを含め、 max_workers <= 1 を維持してください。同時実行される CPU の初回ロードはカーネルを 失う可能性があり、リトライしても安全にはなりません。CUDA のみのロードは影響を受けません。

測定値は環境固有です。コールド時間は CPU、GPU、ドライバー、 ツールチェーン、Warp のバージョンによって変動します。メカニズムは移植できますが、数値はできません。 広範な主張をする前に、references/mechanisms.md の既知の未知事項を読んでください。

参照ファイル

• references/mechanisms.md: 13 のコンパイル時メカニズム。それぞれに 確認シグナル、修正、適用限界、故障モードがあります。測定が示す セクションを読んでください。
• references/measurement.md: 測定プロトコル、各メトリクスが意味することと 意味しないこと、報告のガイダンス、ログの例、手動測定。