# 機能設計書 76-MLIR変換

## 概要

本ドキュメントは、TensorFlowにおけるMLIR（Multi-Level Intermediate Representation）変換機能の設計について記述する。MLIR変換は、TFダイアレクトからHLO（XLA中間表現）への変換パイプラインを提供し、TF→MLIR TF→MLIR HLO→XLAの最適化チェーンを実現する。また、StableHLOポータブルAPIも提供する。

### 本機能の処理概要

**業務上の目的・背景**：従来のTF→XLA変換は直接的なOp→HLO変換を行っていたが、MLIRを導入することで、段階的な変換（Multi-Level）が可能となった。これにより、各レベルで独立した最適化パスを適用でき、TF→HLOの変換品質と保守性が向上する。また、StableHLOの導入により、異なるMLフレームワーク間での中間表現の互換性が確保される。

**機能の利用シーン**：XLAコンパイル時のTF→HLO変換パイプライン、TPUコンパイル時のクラスタリングとシャーディング、TFLiteへの変換時のグラフ最適化、JAX→TFの変換（jax2tf）。

**主要な処理内容**：
1. TF Dialect → TF Executor Dialect変換: TF計算グラフをMLIR表現に変換
2. クラスタリングパス: TPU/XLA向けクラスタの形成と検証
3. TF Dialect → MHLO Dialect変換（Legalization）: TF OpをHLO Opに変換
4. StableHLOポータブルAPI: フレームワーク非依存のHLO操作インターフェース
5. Executor Dialect → Graph変換: MLIR表現からTFグラフへの逆変換

**関連システム・外部連携**：XLAコンパイラ、LLVM MLIR基盤、StableHLO、JAX

**権限による制御**：特段の権限制御はない。

## 関連画面

本機能はバックエンドのコンパイラ変換パイプラインであるため、直接的な関連画面は存在しない。

## 機能種別

コンパイラ変換 / 中間表現処理

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| TF Graph/Module | Graph/MLIRModule | Yes | 変換対象のTFグラフまたはMLIRモジュール | 有効なTF Opで構成されること |
| device_type | string | No | ターゲットデバイス | CPU/GPU/TPU |
| enable_mlir_bridge | bool | No | MLIRブリッジ有効化 | True/False |

### 入力データソース

TensorFlowの計算グラフ（GraphDef）またはMLIRモジュール

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| HLO Module | XlaComputation | XLAが理解できるHLO中間表現 |
| TF Graph | GraphDef | 逆変換時のTFグラフ |
| StableHLO Module | StableHLO | フレームワーク非依存のHLO表現 |

### 出力先

XLAコンパイラまたはTFLiteコンバータへの入力

## 処理フロー

### 処理シーケンス

```
1. Graph → TF Executor Dialect変換
   └─ TFグラフのisland/graphノードをMLIR表現に変換
2. TF Dialectへの正規化
   └─ TF Executor → TF Dialectへのリフティング
3. クラスタリングパス（TPU/XLA向け）
   ├─ TPUClusterFormation: TPUクラスタの形成
   ├─ TPUValidateInputs: 入力の検証
   ├─ TPUShardingIdentification: シャーディングの特定
   ├─ MarkOpsForOutsideCompilation: コンパイル外配置Opのマーク
   └─ VerifyClusteringPass: クラスタリング結果の検証
4. Legalization（TF → HLO変換）
   ├─ legalize_tf.cc: 直接的なパターンマッチ変換
   ├─ legalize_tf_with_tf2xla.cc: TF2XLAカーネル経由の変換
   └─ legalize_tf_communication.cc: 通信Opの変換
5. 検証パス
   └─ verify_tfxla_legalization.cc: 全TF Opが変換済みか検証
6. TF Executor → Graph逆変換（必要時）
```

### フローチャート

```mermaid
flowchart TD
    A[TF Graph] --> B[Graph → TF Executor Dialect]
    B --> C[TF Executor → TF Dialect]
    C --> D{TPU/XLA向け?}
    D -->|Yes| E[クラスタリングパス]
    D -->|No| F[直接Legalization]
    E --> G[TPU Cluster Formation]
    G --> H[Sharding Identification]
    H --> I[Outside Compilation]
    I --> F
    F --> J[TF → HLO Legalization]
    J --> K[Legalization検証]
    K --> L[HLO Module出力]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-76-01 | MLIR Bridge Rollout | デバイスタイプに応じてMLIRブリッジの使用可否をポリシーで制御 | コンパイル時 |
| BR-76-02 | Legalization完全性 | 全てのTF OpがHLO Opに変換されること（未変換Opは検証エラー） | Legalization後 |
| BR-76-03 | Outside Compilation | XLA非対応OpはOutside Compilationとしてマークし、ホスト実行 | TPUクラスタリング時 |
| BR-76-04 | StableHLO互換性 | StableHLO APIはMLIRクラスに依存しないポータブル形式 | StableHLO使用時 |

### 計算ロジック

特段の計算ロジックなし。パターンマッチによるOp変換が主。

## データベース操作仕様

本機能はデータベース操作を行わない。

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| InvalidArgumentError | 変換不可 | TF OpにHLO変換パターンが存在しない | TF2XLAフォールバックまたはOutside Compilation |
| InternalError | 検証失敗 | Legalization後に未変換TF Opが残存 | 該当Opを除外または変換パターンを追加 |
| InternalError | クラスタ不正 | TPUクラスタ検証が失敗 | クラスタ構成を修正 |

### リトライ仕様

自動リトライ機能は提供されない。

## トランザクション仕様

該当なし。

## パフォーマンス要件

- MLIRパスは逐次的に適用されるが、各パスは効率的なパターンマッチで動作
- コンパイル時間はグラフのサイズに比例

## セキュリティ考慮事項

- 特段のセキュリティ考慮事項なし

## 備考

- MLIR TF2XLA API v1とv2が存在し、v2が推奨
- StableHLOはMLIR Python bindingsの完全なエクスポートを避け、ポータブルC++ APIのみを公開
- jax2tf変換で使用されるOpの後方互換性維持が重要

---

## コードリーディングガイド

本機能を理解するために参照すべきファイルと、推奨する読み解き順序を以下に示す。

### 推奨読解順序

#### Step 1: API層を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | legalize_tf.cc | `tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc` | v2 API エントリーポイント |
| 1-2 | cluster_tf.cc | `tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc` | クラスタリングAPI |
| 1-3 | compile_tf_graph.cc | `tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc` | v1コンパイルAPI |

#### Step 2: 変換パスを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | legalize_tf.cc | `tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf.cc` | TF→HLO直接変換パターン |
| 2-2 | legalize_tf_with_tf2xla.cc | `tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla.cc` | TF2XLAカーネル経由の変換 |
| 2-3 | tf2xla_rewriter.cc | `tensorflow/compiler/mlir/tf2xla/transforms/tf2xla_rewriter.cc` | Op書き換えロジック |
| 2-4 | verify_tfxla_legalization.cc | `tensorflow/compiler/mlir/tf2xla/transforms/verify_tfxla_legalization.cc` | 変換完了の検証 |

#### Step 3: クラスタリングパスを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | tpu_cluster_formation.cc | `tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_cluster_formation.cc` | TPUクラスタ形成 |
| 3-2 | tpu_sharding_identification_pass.cc | `tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_sharding_identification_pass.cc` | シャーディング特定 |
| 3-3 | mark_ops_for_outside_compilation.cc | `tensorflow/compiler/mlir/tf2xla/internal/passes/mark_ops_for_outside_compilation.cc` | Outside Compilationマーキング |

#### Step 4: StableHLOポータブルAPIを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 4-1 | stablehlo.py | `tensorflow/compiler/mlir/stablehlo/stablehlo.py` | StableHLO Pythonバインディング（stablehlo_extensionをラップ） |

### プログラム呼び出し階層図

```
TF Graph
    │
    ├─ [API v2]
    │   ├─ cluster_tf.cc (クラスタリング)
    │   ├─ legalize_tf.cc (TF→HLO変換)
    │   ├─ graph_to_tf_executor.cc (Graph→MLIR)
    │   └─ tf_executor_to_graph.cc (MLIR→Graph)
    │
    ├─ [API v1]
    │   ├─ compile_tf_graph.cc (グラフコンパイル)
    │   ├─ compile_mlir_util.cc (MLIRコンパイルユーティリティ)
    │   └─ tf_dialect_to_executor.cc (TF→Executor変換)
    │
    ├─ [Transforms]
    │   ├─ legalize_tf.cc (直接パターン変換)
    │   ├─ legalize_tf_with_tf2xla.cc (TF2XLAカーネル経由)
    │   ├─ legalize_tf_communication.cc (通信Op変換)
    │   ├─ legalize_tf_collective.cc (Collective Op変換)
    │   └─ verify_tfxla_legalization.cc (変換検証)
    │
    ├─ [Internal Passes]
    │   ├─ tpu_cluster_formation.cc
    │   ├─ tpu_sharding_identification_pass.cc
    │   ├─ mark_ops_for_outside_compilation.cc
    │   ├─ extract_outside_compilation.cc
    │   └─ verify_clustering_pass.cc
    │
    └─ [StableHLO]
        └─ stablehlo.py → stablehlo_extension (C++)
```

### データフロー図

```
[入力]                      [処理]                              [出力]

TF GraphDef ──▶  Graph → TF Executor Dialect  ──▶  MLIR Module
                      │
                      └─ クラスタリング（TPU向け）
                      │
                      └─ TF Dialect → HLO Legalization  ──▶  HLO Module
                      │
                      └─ [必要時] Executor → Graph逆変換  ──▶  TF GraphDef
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| legalize_tf.cc (api/v2) | `tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc` | ソース | v2 Legalization API |
| cluster_tf.cc (api/v2) | `tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc` | ソース | v2 クラスタリングAPI |
| legalize_tf.cc (transforms) | `tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf.cc` | ソース | TF→HLO直接変換パターン |
| legalize_tf_with_tf2xla.cc | `tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla.cc` | ソース | TF2XLAカーネル経由変換 |
| verify_tfxla_legalization.cc | `tensorflow/compiler/mlir/tf2xla/transforms/verify_tfxla_legalization.cc` | ソース | 変換完了検証 |
| tpu_cluster_formation.cc | `tensorflow/compiler/mlir/tf2xla/internal/passes/tpu_cluster_formation.cc` | ソース | TPUクラスタ形成 |
| mlir_bridge_rollout_policy.cc | `tensorflow/compiler/mlir/tf2xla/mlir_bridge_rollout_policy.cc` | ソース | MLIRブリッジロールアウトポリシー |
| stablehlo.py | `tensorflow/compiler/mlir/stablehlo/stablehlo.py` | ソース | StableHLO Pythonバインディング |
