# 機能設計書 105-TPUレプリケーション

## 概要

本ドキュメントは、TPU上でのデータ並列レプリケーション制御機能の設計を記述する。TPUレプリカのメタデータ管理、入出力のレプリケーション、コンパイル結果の管理を提供する。

### 本機能の処理概要

**業務上の目的・背景**：TPUでの大規模モデル学習では、データ並列処理のためにモデルを複数のレプリカに複製して実行する。本機能は、レプリカ数、コアごとのレプリカ配置、トポロジ設定、SPMD（Single Program Multiple Data）パーティショニングなどのレプリケーション制御を担う。

**機能の利用シーン**：TPU Pod上でのデータ並列学習、レプリカ間のテンソル入出力マッピング、TPUコンパイル結果の管理、SPMD方式のXLAパーティショニングなどで利用される。

**主要な処理内容**：
1. レプリケーションメタデータの設定（TPUReplicateMetadata）
2. レプリカ入力の統合（TPUReplicatedInput）
3. レプリカ出力の分配（TPUReplicatedOutput）
4. コンパイル結果の取得（TPUCompilationResult）
5. TPUレプリケーション実行（_TPUReplicate）
6. ダミー入力生成（TPUDummyInput）

**関連システム・外部連携**：XLAコンパイラ、TPUトポロジ、Shardyパーティショナー。

**権限による制御**：TPUデバイスへのアクセス権限が必要。

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | - | - | 本機能に関連する画面は登録されていない |

## 機能種別

計算処理（分散レプリケーション制御）

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| num_replicas | int | Yes | レプリカ数 | >= 0（Metadata）、>= 1（_TPUReplicate） |
| num_cores_per_replica | int | No | レプリカあたりのコア数（デフォルト: 1） | 正の整数 |
| topology | string | No | TPUトポロジ文字列（デフォルト: ""） | - |
| use_tpu | bool | No | TPU使用フラグ（デフォルト: true） | - |
| device_assignment | list(int) | No | デバイス割り当て | - |
| inputs (ReplicatedInput) | N * T | Yes | 各レプリカからの入力（N >= 1） | 全入力が同一形状 |
| input (ReplicatedOutput) | T | Yes | 分配するテンソル | - |
| computation | func | Yes（_TPUReplicate時） | 実行する計算関数 | - |
| use_spmd_for_xla_partitioning | bool | No | SPMD方式XLAパーティショニング（デフォルト: false） | - |
| use_shardy_partitioner | bool | No | Shardyパーティショナー使用（デフォルト: false） | - |
| is_mirrored_variable | bool | No | ミラー変数フラグ（デフォルト: false） | - |
| is_packed | bool | No | パック入力フラグ（デフォルト: false） | - |
| tpu_compile_options_proto | string | No | コンパイルオプションプロト（デフォルト: ""） | - |

### 入力データソース

TPUレプリカの入力テンソル、計算関数。

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| output (ReplicatedInput) | T | 統合された入力テンソル |
| outputs (ReplicatedOutput) | num_replicas * T | レプリカごとの出力テンソル |
| output (CompilationResult) | string | コンパイル結果（スカラー文字列） |
| outputs (_TPUReplicate) | output_types | レプリケーション実行結果 |
| output (DummyInput) | dtype | ダミーテンソル（float/bfloat16） |

### 出力先

TPUレプリカの計算結果テンソル。

## 処理フロー

### 処理シーケンス

```
1. TPUReplicateMetadata設定
   └─ レプリカ数、トポロジ、デバイス割り当て等を指定
2. TPUReplicatedInputで入力をレプリカに分配
   └─ 各レプリカの入力テンソルを統合
3. _TPUReplicateで計算を実行
   └─ computation関数を各レプリカで並列実行
4. TPUReplicatedOutputで出力を収集
   └─ 統合テンソルを各レプリカの出力に分配
5. TPUCompilationResultでコンパイル状態を確認
```

### フローチャート

```mermaid
flowchart TD
    A[TPUReplicateMetadata設定] --> B[TPUReplicatedInput<br/>レプリカ入力統合]
    B --> C[_TPUReplicate<br/>計算実行]
    C --> D[TPUReplicatedOutput<br/>出力分配]
    D --> E[TPUCompilationResult<br/>コンパイル結果確認]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-105-01 | 入力形状統合 | TPUReplicatedInputの全入力テンソルは同一形状にマージ可能である必要がある | レプリカ入力時 |
| BR-105-02 | リソースハンドル統合 | DT_RESOURCE型の入力はハンドルの形状・型もマージされる | リソース変数使用時 |
| BR-105-03 | デバイス割り当て | device_assignmentがtopologyと整合性がある必要がある | デバイス指定時 |
| BR-105-04 | ステップマーカー | step_marker_locationでステップ境界の位置を制御（デフォルト: STEP_MARK_AT_ENTRY） | パフォーマンス最適化 |

### 計算ロジック

TPUReplicatedInputのShapeFnは全入力テンソルの形状をマージし、不一致がある場合はエラーを返す。DT_RESOURCE型の場合はハンドルのshapes_and_typesも統合する。

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

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

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| InvalidArgument | 形状不一致 | レプリカ入力の形状がマージ不可能 | 全レプリカの入力形状を統一 |
| InvalidArgument | トポロジ不整合 | device_assignmentとtopologyの不一致 | 設定を確認 |

### リトライ仕様

自動リトライは行われない。

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

レプリケーション操作はTPU_SYSTEMデバイス上でステートフルに実行される。

## パフォーマンス要件

- レプリカ数に比例したスケーラビリティ
- SPMDパーティショニングによるモデル並列の最適化

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

- TPUデバイスへのアクセス制御が必要

## 備考

- computation_shape属性は非推奨（num_cores_per_replicaを使用）
- padding_map属性は非推奨
- TPUDummyInputはSetDoNotOptimizeが設定されており、最適化の対象外

---

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

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

### 推奨読解順序

#### Step 1: Op登録を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | tpu_replication_ops.cc | `tensorflow/core/ops/tpu_replication_ops.cc` | 全6個のレプリケーションOp登録定義 |

**主要処理フロー**:
1. **28-43行目**: TPUReplicateMetadata Op - num_replicas、topology、use_spmd等の属性定義
2. **45-85行目**: TPUReplicatedInput Op - 全入力の形状マージ、DT_RESOURCE時のハンドル統合ロジック
3. **87-97行目**: TPUReplicatedOutput Op - 入力テンソルをnum_replicas個の出力にコピー
4. **99-101行目**: TPUCompilationResult Op - スカラー文字列出力
5. **103-128行目**: _TPUReplicate Op - computation関数、Tinputs/Tbroadcast_inputs/NumVariables等の多数の属性
6. **130-142行目**: TPUDummyInput Op - float/bfloat16型、shape属性からの出力形状推論

#### Step 2: ShapeFn実装を理解する

TPUReplicatedInputのShapeFnが最も複雑で重要。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | tpu_replication_ops.cc | `tensorflow/core/ops/tpu_replication_ops.cc` | TPUReplicatedInputのShapeFn（55-85行目） |

**読解のコツ**:
- **56-62行目**: 全入力テンソルの形状を逆順にマージ（最後の入力から）
- **66-83行目**: DT_RESOURCE型の特殊処理 - input_handle_shapes_and_typesのマージ
- MergeInputHandleShapesAndTypesの戻り値は「形状が精緻化された」を示し、エラーではない

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

```
TPUReplicateMetadata (メタデータ設定)
    │
    ├─ TPUReplicatedInput (入力統合)
    │      └─ 形状マージ + ハンドル統合
    │
    ├─ _TPUReplicate (計算実行)
    │      ├─ inputs: レプリカ入力
    │      ├─ broadcast_inputs: ブロードキャスト入力
    │      ├─ variables: リソース変数
    │      ├─ guaranteed_constants: 定数
    │      └─ computation: 計算関数
    │
    ├─ TPUReplicatedOutput (出力分配)
    │      └─ 入力を num_replicas 個にコピー
    │
    └─ TPUCompilationResult (コンパイル結果)
```

### データフロー図

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

レプリカ0入力 ─┐
レプリカ1入力 ─┼─▶ TPUReplicatedInput ──▶ _TPUReplicate ──▶ TPUReplicatedOutput
レプリカN入力 ─┘    (形状マージ)          (並列計算)         (出力分配)
                                                              ├─▶ レプリカ0出力
                                                              ├─▶ レプリカ1出力
                                                              └─▶ レプリカN出力
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| tpu_replication_ops.cc | `tensorflow/core/ops/tpu_replication_ops.cc` | ソース | レプリケーションOp登録定義（6 Ops） |
