偶然にも Jax 用の LLVM コンパイラを構築してしまった

ソフトウェアを構築していると、時々、非常に深くヤクの毛を剃ることで、驚くほど素敵なセーターができあがることがあります（ヤクは、なぜ自分が突然とても隙間風を感じるのか、不思議に思っていることでしょう）。過去しばらくの間、私のチームと私は、量子ソフトウェアライブラリ PennyLane のための量子コンパイラ Catalyst を構築してきました。私たちの目標は比較的単純でした。MLIR を使用して、大規模なハイブリッド量子古典ワークフローをスケーラブルに最適化することです。これを行うために、古典的な Python 処理（NumPy や関連する科学ライブラリを含む）をキャプチャし、それを私たちの内部表現に表現するための、高速で堅牢な方法が必要でした。私たちは JAX を選択しました。その理由はいくつかあります。Python 関数をトレースして計算グラフをキャプチャする能力¹、比較的良好なカバレッジで NumPy および SciPy API をサポートしていること、そしてすでに MLIR に低下（lower）しているという事実です。また、ネイティブ Python 制御フロー² のキャプチャや、動的形状配列³ のサポートといった、いくつかの品質向上策も目指していました。 しかし、JAX を量子パイプラインに供給するために配線している過程で、私たちはあるくだらないことに気づきました。実際には、Catalyst の @qjit ワークフローに量子命令を一切含める必要はありません。純粋な標準 JAX NumPy コードを、ネイティブ Python 制御フローと一緒に渡すことができます。そして、そうすると、Catalyst は XLA コンパイラバックエンドを完全にバイパスし、JAX 表現を直接標準 MLIR に低下させ、LLVM を介して機械語にコンパイルします（バックプロパゲーションサポート付き）。つまり、私たちは偶然、JAX 用の MLIR コンパイルパイプラインを構築したのです。 ```python import jax.numpy as jnp from catalyst import qjit @qjit(autograph=True) def iterative_layer(weights, inputs, threshold): x = inputs while jnp.mean(x) > threshold: x = jnp.sin(jnp.dot(weights, x)) return x ``` 私たちがここにたどり着いた経緯（つまり、元々剃っていたヤク）を理解するには、量子コンパイルと量子勾配の世界に少し寄り道する必要があります。私たちの目標は 3 つありました。 既存の成熟した古典コンパイルツール（LLVM）を活用し、それに量子サポートを追加することです。独自の量子インフラストラクチャを個別に構築し、古典世界からの機能性をゆっくりと追加するのではなく。私たちは量子コンピューティングの専門家です。古典ソフトウェアとインフラストラクチャに関して、なぜ車輪の再発明をする必要があるのでしょうか？⁴ 量子プログラムを構造的に表現できるようにすることです。つまり、配列操作、量子命令、ループ、if 文を自然に絡み合わせる能力です。これは、プログラムのコンパクトな表現を維持するのに役立ちます（私たちは自然にループや if 文で量子アルゴリズムを記述しますが、それらを維持したままコンパイルすべきです。そうでなければコンパイルは決してスケールしません）。また、ネイティブに動的な量子プログラムの部分（例：量子測定・成功まで繰り返すパターンを表す `while` ループ）をサポートします。 最後に、量子自動微分を引き続きサポートすることです。量子ハードウェア自体で量子プログラムの勾配を計算する方法を見つけるために多大な努力を払ってきましたが、それが（必要な古典自動微分の隣で）引き続き機能することを確認したかったのです。 上記の 3 つの点を、簡単な声明で要約できるかもしれません。量子アルゴリズムは決して純粋に量子だけではありません⁵。入力を準備し、出力を処理し、正直に言って回路を構築するためでさえ、大量の古典処理が必要です。これは、量子エラー訂正の前段階であり、量子処理装置（QPU）を取り囲む古典制御ハードウェアは、大規模にパラメータとゲートを計算し、QPU に供給し、測定を実行し、超低遅延で決定論的なレイテンシでリアルタイムに反応する必要があります。自然な結論は、量子命令だけをコンパイルするのではなく、ハイブリッド実行グラフ全体を、ベアメタルに可能な限り近い場所で実行される、単一のスタンドアロン実行可能ファイルにコンパイルする必要があるということです。 私たちはこのコンパイラインフラストラクチャを MLIR 上に構築することにしました。JAX は Python レイヤーにとって自然な選択肢です（微分可能プログラミングでの使用と、すでに古典コードをキャプチャして MLIR に低下させているため）。そこで、量子命令セットといくつかの品質向上策（Python 制御フロー、動的形状配列）をサポートするように JAX を拡張しました。次に、バックプロパゲーションを処理するために、Enzyme を配線して古典 LLVM コードを直接バックプロパゲートできるようにしました。一方、カスタムで MLIR パスを記述して量子勾配を生成します。 古典 JAX コードを標準 LLVM IR に直接コンパイルするため、Enzyme のような標準 LLVM パスを「無料で」利用できます。これは、XLA の内部パイプラインに直接配線しようとすると、非常に厄介な作業になります。XLA は、線形代数を最適化し、GPU や Google 独自の TPU をターゲットとするための堅固なエンジニアリングです。内部的には、JAXpr（Python プログラムの表現）を取得し、それを MLIR StableHLO ダイアレクトに変換し、独自の重い最適化を実行し、最終的に LLVM を使用して CPU や GPU の機械語にコンパイルします（または TPU を直接ターゲットにします）。 では、XLA がすでに LLVM を使用しているのに、なぜ私たちの方法は違うのでしょうか？線形代数を StableHLO に低下させるための重い作業は、引き続き JAX に依存しています。しかし、その表現を従来の XLA パスに送る代わりに、Catalyst はそれをインターセプトし、標準の汎用 MLIR ダイアレクト（linalg、arith、scf など）に直接低下させます。そこから、それを直接 LLVM と Enzyme に供給し、XLA ランタイムを完全にドロップします。（余談ですが、いつか HLO を完全にバイパスして、NumPy のセマンティクスを直接キャプチャしたいと思っています。しかし、今のところ、線形代数の低下を JAX に任せることで、不必要な苦痛の多くを回避できます！） フロントエンドのトレーシングがバックエンド実行から完全に分離されているため、量子部分はオプションです。量子命令がゼロの純粋な古典関数を私たちの `@qjit`（Quantum Just-In-Time）デコレータに渡すと、それは機能します。 ```python import jax.numpy as jnp from catalyst import qjit import catalyst # ここには量子コードはありません。純粋な JAX + ネイティブ Python 制御フローのみです！ @qjit(autograph=True) def iterative_layer(weights, inputs, threshold): x = inputs # 標準の Python ループと条件分岐はシームレスに機能します # jax.lax.while_loop や jax.lax.cond は不要です！ while jnp.mean(x) > threshold: x = jnp.sin(jnp.dot(weights, x)) catalyst.debug.print('x={x}', x=x) # ランタイムでの印刷 return x ``` では…何がポイントなのでしょうか？正直に言って？まだ完全にはわかっていません。はっきり言っておきます。これは、標準的なディープラーニングワークロードにおいて XLA に勝るものではありません。XLA は、GPU および TPU 上の線形代数に関して、長年の超特化型最適化を持っています。大規模なトランスフォーマーをトレーニングしている場合は、標準 JAX を使用してください。 しかし、JAX をより広範な LLVM エコシステムに直接接続し、重い XLA ランタイムをドロップしたときに何が起こるかは、クールだと考えています。（さらに、XLA を Bazel でビルドする必要もありません！どういたしまして。）このブリッジは、いくつかの奇妙で素晴らしいハッキングの機会と、品質向上策を開きます。 スタンドアロン AOT バイナリ：XLA は、PJRT のような重い Python/C++ ランタイムによって管理される JIT（Just-In-Time）実行環境を期待します。それをバイパスすることで、実行に外部 ML 依存性を必要としない、純粋なスタンドアロン AOT（Ahead-of-Time）バイナリを生成できます。 動的形状配列：JAX を使用する人なら誰でも知っているように、XLA はテンソル形状に関して非常に厳格です。入力ディメンションが変更されると、通常はイライラするほど遅い再コンパイルが発生します。標準 MLIR と LLVM に直接マッピングするため、コンパイル済みバイナリで動的形状をすぐにサポートできます。可変長シーケンスを処理している場合、これは大規模な QoL 向上です。 ネイティブ Python 制御フローと NumPy インデックス処理：`while`、`if`、`for`、条件分岐、ブール演算、さらには NumPy 配列の代入もすべて、Autograph/Diastatic Malt を介してすぐに機能します（そして微分可能です）。 ベアメタルおよびエッジデプロイメント：JAX でトレーニングされたモデルを実行したいですか？