Pythonのcopy-and-patch JITの実装を読む

Transcript

1. copy-and-patch JITを書く Pythonのcopy-and-patch JITの 実装を読む 2025/10/03 コンパイラのコンパの部分 #01 @Arata
2. 2025/10/02 コンパイラのコンパの部分 #01 @Arata 背景: Python 3.13での試験的なJITの導入 2
3. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Python 3.11: PEP659 - Specializing Adaptive Interpreter • 汎用的な命令から特定の型に特化した命令へ実行時に書き換える ◦ e.g. x + 1: BINARY_OP(ADD) -> BINARY_OP_ADD_INT • データの型、よく実行される経路などのプロファイルも取る Python 3.12: DSLからのインタプリタ生成 • バイトコードのセマンティクスがより明確に Python 3.13ではこれらのプロファイラやDSLがJIT実装の基礎となった 背景: Faster CPythonの道のり 3
4. 2025/10/02 コンパイラのコンパの部分 #01 @Arata copy-and-patchと呼ばれる仕組みを用いてJITを実装している Pythonビルド時: インタプリタのコードから各命令のステンシルを作成 1. インタプリタのコードから各命令の実行処理を抽出 2. 各命令の実行処理をそれぞれコンパイル 3. コンパイル結果からステンシル作成 Pythonコード実行時: ステンシルを使ってJITコンパイル 1. 各命令に対応するステンシルをコピー 2. ステンシル中の一部にパッチを当てる PythonのJIT実装 4
5. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Python/bytecodes.c (DSL)からPython/executor_cases.c.h (C)を生成 巨大なswitch-caseになる Pythonビルド時: インタープリタのコード生成 5 Python/bytecodes.c Python/executor_cases.c.h
6. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Python/executor_cases.c.hから正規表現で各命令の実装コードを抽出 Pythonビルド時: インタプリタからコード抽出 6 Tools/jit/_targets.py Python/execuor_cases.c.h 抽出対象
7. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Pythonビルド時: ステンシル作成 7 抽出した各命令の実装コードをテンプレートに挿入・コンパイル 実行時に決まる値はテンプレートで外部シンボルとして定義 コード挿入場所→ Tools/jit/template.c
8. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Pythonビルド時: ステンシルの最適化 8 ステンシルをアセンブリレベルで最適化 • 条件分岐命令でホットパスをフォールスルー側に配置 • 冗長な条件分岐命令の削除 Tools/jit/_optimizers.py
9. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Pythonビルド時: ステンシルのホール作成 9 実行時に決まる値=外部シンボルの参照には再配置エントリがあるはず それらの参照位置と参照シンボルを記録 Tools/jit/_stencils.py
10. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Pythonビルド時: ヘッダーファイルの作成 10 ここまででステンシルの情報が集まったので、JITコンパイルに使う 関数たちを生成
11. 2025/10/02 コンパイラのコンパの部分 #01 @Arata Pythonビルド時: ヘッダーファイルの作成 11 ここまででステンシルの情報が集まったので、JITコンパイルに使う 関数たちを生成 ここではUNARY_NOTの例を紹介 ステンシルのアセンブリ UNARY_NOTの例 ステンシルのバイト列 code + data ステンシルのコピー&パッチ False,Trueへの参照を解決
12. 2025/10/02 コンパイラのコンパの部分 #01 @Arata 結局何ができたのか？ • 「あるバイトコード命令を実行する機械語」を生成する関数が全ての 命令について手に入った • これを命令ごとに呼び出せばバイトコードを機械語に変換できる → JITコンパイル達成 copy-and-patchの名前の由来は？ • JITコンパイル時にステンシルをコピー&パッチするため • パッチは実行時に決まるオペランド、外部シンボルの参照等で必要 まとめ 12
13. 2025/10/02 コンパイラのコンパの部分 #01 @Arata 余談: PyTorchのJITコンパイラ(TorchDynamo)も面白い 概要 • PEP 523で評価器を差し替えて、実行されたテンソル演算をトレース • トレースから計算グラフを構築してコンパイル • 特殊化のため、コンパイルされたコードから参照する変数の値が変化 していないことを実行時に検証する仕組みがある 終わりに 13
14. 2025/10/02 コンパイラのコンパの部分 #01 @Arata • https://peps.python.org/pep-0659/ • https://peps.python.org/pep-0744/ • https://github.com/python/cpython/tree/main/Tools/jit • https://www.youtube.com/watch?v=HxSHIpEQRjs • https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_ov erview.html 参考文献 14