# 機能設計書 79-TF2XLA変換

## 概要

本ドキュメントは、TensorFlowにおけるTF2XLA変換機能の設計について記述する。TF2XLAは、TensorFlowオペレーションをXLA HLOオペレーションに変換するためのライブラリであり、HLO Builder APIを通じてXLAの演算プリミティブを直接操作する低レベルインターフェースを提供する。

### 本機能の処理概要

**業務上の目的・背景**：XLAコンパイラを活用するためには、TFオペレーションをXLAのHLO中間表現に変換する必要がある。TF2XLAは、各TF OpにHLO変換パターン（XLAカーネル）を登録し、グラフ全体のTF→HLO変換を実現する。また、jax2tf変換で使用されるXLA Opの直接的なPython APIも提供する。

**機能の利用シーン**：tf.function(jit_compile=True)によるXLAコンパイル時の内部変換、AOTコンパイル（tfcompile）時のグラフ変換、jax2tfによるJAX→TF変換時のHLO Op利用。

**主要な処理内容**：
1. TF Op → HLO Op変換カーネルの提供
2. XLA Op Registry: デバイスタイプごとのXLA対応Op管理
3. Python XLA API: HLO演算を直接操作するPythonインターフェース
4. 定数分析: XLAコンパイル時にコンパイル時定数が必要な引数の特定
5. リソース操作テーブル: リソース変数操作のXLA互換性分類

**関連システム・外部連携**：XLAコンパイラ、HLO Builder、JAX

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

## 関連画面

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

## 機能種別

Op変換ライブラリ / HLO Builder API

## 入力仕様

### 入力パラメータ（Python XLA API - xla.py）

Python XLA APIは、HLOオペレーションに対応するPython関数群を提供する。主要な関数カテゴリ：

**単項演算子**: abs, cos, ceil, exp, floor, log, neg, sin, sqrt, tanh等（math_opsへの直接委譲）

**二項演算子**: add, sub, mul, div, pow, max, min等

**XLA固有オペレーション**: gen_xla_opsモジュール経由で提供

### 入力データソース

TensorFlowの計算グラフ内のTensorFlow Op

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| HLO Computation | XlaComputation | 変換されたHLO計算 |
| Tensor | tf.Tensor | Python XLA API使用時の出力テンソル |

### 出力先

XLAコンパイラへの入力（HLO Computation）、または計算グラフ内の後続Op

## 処理フロー

### 処理シーケンス

```
1. Op登録
   └─ XLA Op Registryにデバイスタイプごとの変換カーネルを登録
2. 定数分析
   └─ const_analysis.hで各Opのコンパイル時定数引数を特定
3. GraphDef → XlaComputation変換
   └─ tf2xla.hのConvertGraphDefToXla関数
   └─ 各TF Opに対応するXLAカーネルを呼び出しHLO Opに変換
4. Python XLA API
   └─ 単項/二項演算はTF math_opsへ委譲
   └─ XLA固有OpはgenXlaOps経由でC++カーネルを呼び出し
```

### フローチャート

```mermaid
flowchart TD
    A[TF Op] --> B{XLA Op Registry登録あり?}
    B -->|Yes| C[XLAカーネル呼び出し]
    B -->|No| D[XLA非対応 → フォールバック]
    C --> E[HLO Op生成]
    E --> F[HLO Computation構築]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-79-01 | Op互換性 | TF OpにXLAカーネルが登録されている場合のみ変換可能 | Op変換時 |
| BR-79-02 | jax2tf後方互換 | jax2tfで使用されるOpのXLA APIは後方互換性を維持する | xla.py API |
| BR-79-03 | 単項演算委譲 | 単項演算はTF math_opsの対応関数に委譲 | Python XLA API |
| BR-79-04 | リソース操作分類 | リソース変数操作はread/write/read-write/unknownに分類 | リソース安全性分析 |

### 計算ロジック

特段の計算ロジックなし。各TF OpのHLO等価変換がパターンとして定義されている。

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

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

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| InvalidArgumentError | 非対応Op | XLA変換カーネルが未登録のOp | OpのXLAサポートを確認 |
| InternalError | 変換失敗 | HLO変換が失敗 | Op入力の形状・型を確認 |
| ValueError | 型不正 | unsupported型のテンソル | サポートされた型を使用 |

### リトライ仕様

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

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

該当なし。

## パフォーマンス要件

- Op変換は線形時間で実行
- 変換結果はXLAコンパイラのキャッシュで再利用

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

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

## 備考

- xla.pyのほとんどの単項/二項演算はTF標準Opのラッパー
- gen_xla_opsはC++ XLA OpのPythonバインディング（自動生成）
- roundはXLAとTFで丸め方向が異なる（XLA: 0方向、TF: 偶数方向）

---

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

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

### 推奨読解順序

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

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | xla.py | `tensorflow/compiler/tf2xla/python/xla.py` | Python XLA APIのエントリーポイント |

**読解のコツ**:
- **15-25行目**: モジュールドキュメント。jax2tf互換性維持の重要性
- **27行目**: gen_xla_opsのインポート（XLA固有Op）
- **70-76行目**: _unary_opラッパー関数パターン（シグネチャの標準化）
- **79-109行目**: 単項演算の一覧（abs, cos, exp, log等 → math_opsに委譲）
- **43-49行目**: TODO: 未実装のXLA Op一覧

#### Step 2: C++変換基盤を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | xla_op_registry.h | `tensorflow/compiler/tf2xla/xla_op_registry.h` | Op登録レジストリ |
| 2-2 | const_analysis.h | `tensorflow/compiler/tf2xla/const_analysis.h` | コンパイル時定数分析 |
| 2-3 | resource_operation_table.h | `tensorflow/compiler/tf2xla/resource_operation_table.h` | リソース操作分類テーブル |
| 2-4 | xla_argument.h | `tensorflow/compiler/tf2xla/xla_argument.h` | XLA引数データ構造 |
| 2-5 | tf2xla.h | `tensorflow/compiler/tf2xla/tf2xla.h` | GraphDef→XlaComputation変換 |

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

```
Python XLA API (xla.py)
    │
    ├─ 単項演算 (abs, cos, exp, ...) → math_ops
    ├─ 二項演算 (add, sub, mul, ...) → math_ops
    └─ XLA固有Op → gen_xla_ops → C++ XLA Ops

TF2XLA C++ 変換基盤
    │
    ├─ XlaOpRegistry (Op登録・検索)
    │      └─ 各Op用のXlaOpKernel
    │
    ├─ tf2xla::ConvertGraphDefToXla (グラフ変換)
    │      └─ XlaOpKernel::Compile → HLO Builder API
    │
    ├─ ConstAnalysis (定数引数特定)
    │
    └─ ResourceOperationTable (リソースOp分類)
```

### データフロー図

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

TF GraphDef ──▶  ConvertGraphDefToXla  ──▶  XlaComputation
                    │
                    ├─ XlaOpRegistry検索
                    ├─ 各Opの XlaOpKernel::Compile
                    └─ HLO Builder API呼び出し

TF Tensor ──▶  xla.py API関数  ──▶  TF Tensor（XLA Opバックエンド）
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| xla.py | `tensorflow/compiler/tf2xla/python/xla.py` | ソース | Python XLA API |
| gen_xla_ops | 自動生成 | ソース | C++ XLA OpのPythonバインディング |
| xla_op_registry.h | `tensorflow/compiler/tf2xla/xla_op_registry.h` | ヘッダ | Op登録レジストリ |
| tf2xla.h | `tensorflow/compiler/tf2xla/tf2xla.h` | ヘッダ | GraphDef→XlaComputation変換 |
| tf2xla_util.h | `tensorflow/compiler/tf2xla/tf2xla_util.h` | ヘッダ | ユーティリティ |
| const_analysis.h | `tensorflow/compiler/tf2xla/const_analysis.h` | ヘッダ | 定数分析 |
| resource_operation_table.h | `tensorflow/compiler/tf2xla/resource_operation_table.h` | ヘッダ | リソース操作テーブル |
| xla_argument.h | `tensorflow/compiler/tf2xla/xla_argument.h` | ヘッダ | XLA引数定義 |
| tf2xla.pb.h | 自動生成 | ヘッダ | TF2XLA Config Proto定義 |
