技術要約: KeyMemRT コンパイラおよびランタイム

問題提起

完全準同型暗号（FHE）は、暗号化されたデータに対するプライバシー保護計算を可能にするが、過度な計算量とメモリ要件により、実世界での導入には大きな障壁がある。主要なボトルネックの一つは、**回転鍵（rotation keys）**の管理である。バッチ化されたFHEスキーム（例：CKKS, BGV, BFV）において、回転操作はベクトル内計算やブートストラップに不可欠である。しかし、これらのスキームでは、特定の回転インデックスごとに固有の回転鍵を必要とする。

FHEプログラムがスケールするにつれ、必要な回転インデックスの数は増加し、すべての鍵を同時にメモリに格納するために、数百ギガバイト、あるいはテラバイト単位のメモリが必要になることがよくある。既存のFHEコンパイラは、一般に回転鍵を単一の静的な実体として扱い、プログラムの実行期間中、すべての鍵をメモリにロードしておく。このアプローチはメモリの拡張性に欠ける。対照的に、メモリ削減を試みる既存の戦略であるFhelipeは、任意の回転インデックスを2のべき乗の連鎖へと分解することに依存している。これは、格納される鍵の数を減らす一方で、連鎖によるオーバーヘッドのために計算レイテンシを大幅に増加させる。メモリ使用量と実行速度のバランスを取るための、きめ細かな自動鍵管理を提供するコンパイラは存在しない。

手法

著者らは、回転鍵のライフタイムを個別に管理するように設計されたコンパイラおよびランタイムフレームワークであるKeyMemRTを提案する。本システムはMLIR (Multi-Level Intermediate Representation) インフラストラクチャ上に構築されており、新しいダイアレクトとランタイムアーキテクチャを導入している。

1. kmrt MLIR ダイアレクト

KeyMemRTは、IR内で鍵管理を直接表現するためのkmrtと呼ばれる新しいダイアレクトを導入している。

• 新しい型: 回転インデックスによってパラメータ化されたrot_key型を定義し、コンパイラが異なるインデックス間の鍵を区別できるようにしている。
• 新しい操作: ダイアレクトには、明示的な鍵管理のための操作が含まれる：
  • kmrt.load_key: 特定の回転鍵をロードする。
  • kmrt.clear_key: メモリから鍵を削除する。
  • kmrt.prefetch_key: 非同期ロードのためのヒントをランタイムに提供する。
  • kmrt.use_key / kmrt.assume_key: 鍵のスコープと生存期間（liveness）を管理する。
• 統合: コンパイラは、高レベルの操作（例：HEIRを介したckksダイアレクトからの操作）を、kmrt操作を挿入しながらopenfheダイアレクトへと低レベル化（lowering）する。

2. コンパイラ最適化パイプライン

Keyもたらされる性能を犠牲にすることなく、メモリフットプリントを削減するために、KeyMemRTコンパイラは一連の低レベル化および最適化パスを実行する：

• CKKSからOpenFHEへの低レベル化: すべての回転操作を特定のrot_key値に関連付け、問題を「生存範囲の作成」から「生存範囲の統合」へと反転させる。
• Baby-Step Giant-Step (BSGS) 分解: 連続する回転インデックスに対して、回転を連鎖させることで、必要な鍵の数を$N$から$\approx 2\sqrt{N}$へと削減する。これは行列ベクトル乗算のパターンにおいて特に効果的である。
• 回転鍵の統合 (Merge Rotation Keys): データフロー解析を用いて、ある鍵がクリアされた直後に同じインデックスで再ロードされるシーケンスを検出し、冗長なロード/クリアのペアを削除する。
• 回転のホイスティング (Rotation Hoisting): 共通の入力に適用される回転パターンを特定し、それらをOpenFHEライブラリがサポートするより高速な代替手段へと変換する。
• ブートストラップ鍵の管理: 暗号文のレベルを更新するために使用される内部回転鍵をIR内に実体化する。これにより、コンパイラはブートストラップ鍵を、明示的な回転鍵と同じきめ細かな制御で管理できるようになる。これは、これまで対処されていなかった機能である。
• 共通部分式除去 (Common Subexpression Elimination, CSE): 鍵の最適化によって生じるデッドコードを削除する。
• ヒントの配置: ランタイムが鍵のロードと計算をオーバーラップ（並列実行）できるように、prefetch_keyヒントを挿入する。

3. KeyMemRT ランタイム

ランタイムは、コンパイルされたC++コードとFHEライブラリ（OpenFHE）との間の中継器として機能する。これは、コンパイラが挿入したヒントに基づき、ディスクからRAMへの鍵の動的なロードおよびアンロードを管理する。

• 低メモリモード (Low-Memory Mode): 指示を逐次的に実行し、メモリ使用量を最小限に抑えるために、厳密に必要とされる時にのみ鍵をロードする。
• バランスモード (Balanced Mode): 計算と並行して鍵をロードするプリフェッチメカニメントを利用する。内部キューを維持し、ユーザー定義の制限まで事前に鍵をロードし、スペースが必要になったときにそれらをクリアすることで、レイテンシのオーバーヘッドを軽減する。

主な貢献

1. 新しいMLIRダイアレクト (kmrt): 回転鍵を特定の型とメモリ管理操作を持つIR要素として公開する初のダイアレクトであり、きめ細かな解析を可能にする。
2. 自動的なきめ細かな鍵管理: 鍵の生存期間を自動的に追跡し、ブートストラップ鍵を最適化し、鍵をグローバルな集合としてではなく個別に管理する初のFHEコンパイラである。
3. ランタイムフレームワーク: コンパイラのヒントを活用して、I/Oと計算をオーバーラップさせるオンデマンドの鍵ロードおよびクリアをサポートするランタイム。
4. 統合: 既存のFHEフロントエンド（OrionおよびHEIRで実証）がターゲットにできる、最適化後のコンパイラツールチェーン。

結果

著者らは、KeyMemRTを2つの最先端のアプローチ、すなわちANT-ACE（すべての鍵をメモリに保持する）およびFhelipe（鍵を減らすために分解を使用するが、レイテンシが増加する）と比較評価した。ベンチマークには、リング次元$2^{16}$、乗法的深度30のFHEパラメータを設定した様々なMLモデル（MLP, LoLa, LeNet, ResNetのバリアント, AlexNet）が含まれる。

• メモリ削減:
  • ANT-ACEと比較して、KeyMemRTは幾何平均で1.74倍のメモリ削減を達成した。LeNetのような特定のケースでは、メモリ使用量は2.59倍削減された。
  • Fhelipeと比較して、KeyMemRTは1.16倍のメモリ削減を達成した。
• レイテンシ/スピードアップ:
  • ANT-ACEと比較して、KeyMemRT（バランスモード）は1.20倍のスピードアップを達成した。
  • Fhelipeと比較して、KeyMemRTは（幾何平均で）1.73倍の顕著なスピードアップを達成し、AlexNetでは最大2.5倍の改善を示した。これは、KeyMemRTがFhelipeで必要とされる回転の連鎖に伴う計算オーバーヘッドを回避しているためである。
• スケーラビリティ: KeyMemRTは、ANT-ACEの下では100 GBを超えるメモリ要件を持つ大規模なプログラム（例：ResNet-50）を実行できる能力を示した。

意義と主張

本論文は、KeyMemRTが自動的かつきめ細かな鍵管理およびブートストラップ鍵管理を提供する初のFHEコンパイラであることを主張している。著者らは、この研究がFHEのメモリのスケーラビリティを解き放ち、すべての回転鍵を同時に保持するという法外なメモリ負担なしに、任意の規模のプログラムを実行することを可能にすると断言している。

メモリ使用量を総鍵数から切り離し、代わりに鍵の生存期間に基づいて管理することで、KeyMemRTはFHEプログラムを既存のハードウェア上でより良くスケールさせることを可能にする。著者らは、本手法がパフォーマンスを損なわないことを強調している。メモリと速度をトレードオフにする分解ベースの手法とは異り、KeyMemRTは、必要な鍵をまさに必要な瞬間にロードすることにより、メモリ削減とレイテンシ改善の両方を実現している。これは、リソース制約のある環境や大規模な環境におけるFHE採用の決定的な障壁に対処するものである。