# 機能設計書 104-TPUエンベディング

## 概要

本ドキュメントは、TPU上での大規模エンベディングテーブルの操作機能の設計を記述する。TPU専用のエンベディングルックアップ・勾配更新・パラメータLoad/Retrieveメカニズムを提供する。

### 本機能の処理概要

**業務上の目的・背景**：推薦システムや自然言語処理など多くの機械学習モデルでは、大規模なエンベディングテーブルが必要となる。TPUは専用のエンベディングエンジンを持ち、ホストメモリに分散配置されたエンベディングテーブルを効率的にルックアップ・更新する機構を提供する。本機能はこのTPU固有のエンベディング処理パイプラインを実現するOp群を提供する。

**機能の利用シーン**：TPU上での推薦モデル学習（大規模語彙のエンベディング）、TPUエンベディングのパラメータ初期化（Load）、チェックポイント保存前のパラメータ取得（Retrieve）、スパース・密・RaggedTensorベースのバッチエンキュー、エンベディング活性化の受信と勾配送信。

**主要な処理内容**：
1. エンベディングインデックスのエンキュー（EnqueueTPUEmbeddingIntegerBatch、EnqueueTPUEmbeddingSparseBatch等）
2. エンベディング活性化の受信（RecvTPUEmbeddingActivations）
3. エンベディング勾配の送信（SendTPUEmbeddingGradients）
4. エンベディング活性化の自動微分対応（TPUEmbeddingActivations）
5. 各種オプティマイザのパラメータLoad（LoadTPUEmbedding*Parameters）
6. 各種オプティマイザのパラメータRetrieve（RetrieveTPUEmbedding*Parameters）

**関連システム・外部連携**：TPUエンベディングエンジン、TPUEmbeddingConfiguration プロトコルバッファ、各種オプティマイザ（Adagrad、Adam、FTRL、Momentum、RMSProp等）。

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

## 関連画面

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

## 機能種別

計算処理（エンベディングルックアップ・パラメータ管理）

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| batch (IntegerBatch) | N * int32 | Yes | エンベディングインデックスバッチ | N >= 1 |
| sample_indices | N * T1 | Yes | サンプルインデックス（Sparse/RaggedBatch時） | int32/int64 |
| embedding_indices | N * T2 | Yes | エンベディングインデックス | int32/int64 |
| aggregation_weights | N * T3 | Yes | 集約重み | float32/float64 |
| mode_override | string | Yes | 動作モードオーバーライド | - |
| config | string | Yes | TPUEmbeddingConfigurationプロト（シリアライズ） | 有効なプロトバッファ |
| inputs (Gradients) | N * float32 | Yes | エンベディング勾配テンソル | N >= 1 |
| learning_rates | NN * float32 | No | 動的学習率（NN >= 0、デフォルト0） | スカラー |
| parameters | float32 | Yes（Load時） | エンベディングパラメータ | - |
| combiners | list(string) | No | 集約関数（sum/mean/sqrtn） | 空または要素数Nと一致 |
| table_id | int | No | テーブルID | >= 0（Load/Retrieve時、-1=名前指定） |
| table_name | string | No | テーブル名 | - |
| num_shards | int | Yes（Load/Retrieve時） | シャード数 | 正の整数 |
| shard_id | int | Yes（Load/Retrieve時） | シャードID | 0以上 |
| device_ordinal | int | No | 対象TPUデバイス番号（デフォルト: -1） | - |

### 入力データソース

学習バッチ内のカテゴリカル特徴量（スパースまたはRaggedTensorとして表現されたインデックスと重み）。

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| outputs (RecvActivations) | num_outputs * float32 | 各テーブルのエンベディング活性化テンソル |
| output (TPUEmbeddingActivations) | float32 | 自動微分対応の活性化テンソル |
| parameters (Retrieve) | float32 | 取得したエンベディングパラメータ |
| accumulators/momenta/velocities等 | float32 | オプティマイザ状態テンソル |

### 出力先

TPUデバイスメモリ上のエンベディング活性化テンソル、ホストCPUメモリ上のパラメータテンソル（Retrieve時）。

## 処理フロー

### 処理シーケンス

```
[初期化フェーズ]
1. TPUEmbeddingConfigurationプロトを構成
2. tpu.initialize_system()でTPUエンベディングを構成
3. LoadTPUEmbedding*Parametersでパラメータをロード

[学習フェーズ（各ステップ）]
4. EnqueueTPUEmbedding*Batchでインデックスをエンキュー
5. RecvTPUEmbeddingActivationsで活性化を受信
6. TPUEmbeddingActivationsで自動微分に対応
7. SendTPUEmbeddingGradientsで勾配を送信（パラメータ更新）

[チェックポイント保存フェーズ]
8. RetrieveTPUEmbedding*Parametersでパラメータを取得
```

### フローチャート

```mermaid
flowchart TD
    A[TPUEmbeddingConfiguration構成] --> B[LoadTPUEmbedding*Parameters]
    B --> C[学習ループ開始]
    C --> D[EnqueueTPUEmbedding*Batch]
    D --> E[RecvTPUEmbeddingActivations]
    E --> F[モデルのフォワードパス]
    F --> G[SendTPUEmbeddingGradients]
    G --> H{ループ継続?}
    H -->|Yes| C
    H -->|No| I[RetrieveTPUEmbedding*Parameters]
    I --> J[チェックポイント保存]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-104-01 | configの一致 | Enqueue/Recv/SendのconfigパラメータはTPU初期化時と同一である必要がある | 全エンベディングOp |
| BR-104-02 | combinersの長さ | combinersが空でない場合、要素数はNと一致する必要がある | SparseBatch/SparseTensorBatch |
| BR-104-03 | learning_ratesのスカラー制約 | 各learning_rates要素はスカラーである必要がある | SendTPUEmbeddingGradients |
| BR-104-04 | Load/Retrieve対応 | Load Opで読み込んだオプティマイザのパラメータはRetrieve Opで同じオプティマイザの形式で取得する | チェックポイント処理 |
| BR-104-05 | シャーディング | エンベディングテーブルはnum_shardsに分割され、各シャードはshard_idで識別される | Load/Retrieve時 |

### 計算ロジック

RecvTPUEmbeddingActivationsの出力テンソル形状はTPUEmbeddingConfigurationプロトから計算される（ComputeOutputTensorShapes関数による）。集約関数（combiners）はsum/mean/sqrtnから選択可能。

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

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

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| InvalidArgument | 設定エラー | TPUEmbeddingConfigurationプロトが不正 | プロトの内容を検証 |
| InvalidArgument | サイズ不一致 | combinersの長さがNと不一致 | combinersを空またはN個に設定 |
| InvalidArgument | 出力数不一致 | num_outputsとconfigから計算された出力形状数が不一致 | configを確認 |

### リトライ仕様

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

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

エンベディングのLoad/Retrieveは明示的な操作であり、チェックポイント保存前にRetrieveを、学習開始前にLoadを実行する必要がある。全てのOpsはステートフルとして登録されている。

## パフォーマンス要件

- 大規模エンベディングテーブル（数百万～数十億エントリ）を効率的に処理
- ホストメモリに分散配置されたテーブルのシャーディングにより並列アクセスを実現
- RaggedTensorBatch版は可変長シーケンスに対応

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

- TPUデバイスへのアクセス制御が必要
- エンベディングテーブルのパラメータは機密情報を含む可能性がある

## 備考

- tpu_embedding_load_retrieve_ops.ccは自動生成ファイル（generate_tpu_embedding_load_retrieve_ops.pyによる）
- 対応オプティマイザ: Adagrad、AdagradMomentum、SGD、FTRL、ADAM、Momentum、RMSProp、CenteredRMSProp、MDLAdagradLight、Adadelta、ProximalAdagrad、ProximalYogi、FrequencyEstimator
- DynamicEnqueue系のOpsはdevice_ordinalをint32入力テンソルとして動的に指定可能

---

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

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

### 推奨読解順序

#### Step 1: データ構造を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | tpu_embedding_configuration.proto | `tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto` | TPUEmbeddingConfigurationの構造 |
| 1-2 | optimization_parameters.proto | `tensorflow/core/protobuf/tpu/optimization_parameters.proto` | 最適化パラメータの定義 |
| 1-3 | tpu_embedding_output_layout_utils.h | `tensorflow/core/tpu/tpu_embedding_output_layout_utils.h` | 出力レイアウトユーティリティ |

**読解のコツ**: TPUエンベディングは7段階のパイプラインとして設計されている（tpu_embedding_ops.ccの26-58行目のコメント参照）。

#### Step 2: エンベディングOp登録を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | tpu_embedding_ops.cc | `tensorflow/core/ops/tpu_embedding_ops.cc` | エンベディング操作Op（11 Ops） |

**主要処理フロー**:
1. **28-58行目**: 7段階パイプラインの設計コメント
2. **60-84行目**: RecvTPUEmbeddingActivations Op - config解析、ComputeOutputTensorShapesで出力形状計算
3. **86-95行目**: TPUEmbeddingActivations Op - table_id/lookup_id属性、入力形状をそのまま出力
4. **97-117行目**: SendTPUEmbeddingGradients Op - learning_ratesのスカラー検証
5. **119-125行目**: EnqueueTPUEmbeddingIntegerBatch Op - 整数インデックスバッチ
6. **127-151行目**: EnqueueTPUEmbeddingSparseBatch Op - combiners長さ検証
7. **153-168行目**: EnqueueTPUEmbeddingSparseTensorBatch Op - table_ids/max_sequence_lengths属性
8. **170-185行目**: EnqueueTPUEmbeddingRaggedTensorBatch Op - sample_splits入力
9. **204-216行目**: EnqueueTPUEmbeddingArbitraryTensorBatch Op - 汎用バッチ

#### Step 3: Load/Retrieve Op登録を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | tpu_embedding_load_retrieve_ops.cc | `tensorflow/core/ops/tpu_embedding_load_retrieve_ops.cc` | 各オプティマイザのLoad/Retrieve Op（26 Ops） |

**主要処理フロー**:
- **29-38行目**: LoadTPUEmbeddingAdagradParameters - parameters + accumulators入力
- **40-49行目**: RetrieveTPUEmbeddingAdagradParameters - parameters + accumulators出力
- **119-129行目**: LoadTPUEmbeddingADAMParameters - parameters + momenta + velocities入力
- **131-141行目**: RetrieveTPUEmbeddingADAMParameters - 同3出力

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

```
tpu.initialize_system (TPUエンベディング構成)
    │
    ├─ LoadTPUEmbedding*Parameters (初期化)
    │      ├─ LoadTPUEmbeddingAdagradParameters
    │      ├─ LoadTPUEmbeddingADAMParameters
    │      ├─ LoadTPUEmbeddingFTRLParameters
    │      └─ ... (13種類のオプティマイザ)
    │
    ├─ [学習ループ]
    │      ├─ EnqueueTPUEmbedding*Batch
    │      │      ├─ EnqueueTPUEmbeddingIntegerBatch
    │      │      ├─ EnqueueTPUEmbeddingSparseBatch
    │      │      ├─ EnqueueTPUEmbeddingSparseTensorBatch
    │      │      └─ EnqueueTPUEmbeddingRaggedTensorBatch
    │      ├─ RecvTPUEmbeddingActivations
    │      │      └─ ComputeOutputTensorShapes()
    │      ├─ TPUEmbeddingActivations (自動微分)
    │      └─ SendTPUEmbeddingGradients
    │
    └─ RetrieveTPUEmbedding*Parameters (チェックポイント)
           ├─ RetrieveTPUEmbeddingAdagradParameters
           └─ ... (13種類のオプティマイザ)
```

### データフロー図

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

インデックス+重み ──▶ EnqueueTPUEmbedding*Batch ──▶ TPUエンベディングエンジン
                                                          │
                    RecvTPUEmbeddingActivations ◀──────────┘
                           │
                    活性化テンソル ──▶ モデル計算 ──▶ 勾配テンソル
                                                          │
                    SendTPUEmbeddingGradients ◀───────────┘
                           │
                    TPUエンベディングエンジン（パラメータ更新）
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| tpu_embedding_ops.cc | `tensorflow/core/ops/tpu_embedding_ops.cc` | ソース | エンベディング操作Op登録（11 Ops） |
| tpu_embedding_load_retrieve_ops.cc | `tensorflow/core/ops/tpu_embedding_load_retrieve_ops.cc` | ソース | Load/Retrieve Op登録（26 Ops） |
| tpu_embedding_output_layout_utils.h | `tensorflow/core/tpu/tpu_embedding_output_layout_utils.h` | ヘッダ | 出力形状計算ユーティリティ |
| tpu_embedding_configuration.proto | `tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto` | プロト | エンベディング設定プロト |
| optimization_parameters.proto | `tensorflow/core/protobuf/tpu/optimization_parameters.proto` | プロト | 最適化パラメータプロト |
| tpu_embedding_optimization_parameters_utils.h | `tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h` | ヘッダ | Load/RetrieveのShapeFnユーティリティ |
