# 機能設計書 71-コレクティブ操作（Collective Operations）

## 概要

本ドキュメントは、TensorFlowにおけるコレクティブ操作（Collective Operations）機能の設計について記述する。コレクティブ操作は、複数デバイス間でのMPI互換の集約通信操作（AllReduce、AllGather、Broadcast、AllToAll）を提供する低レベルAPIである。

### 本機能の処理概要

**業務上の目的・背景**：分散機械学習において、複数のデバイス（GPU、TPU等）間でテンソルデータを効率的に集約・同期する必要がある。コレクティブ操作は、高レベルの分散ストラテジー（MirroredStrategy、MultiWorkerMirroredStrategy等）が内部で利用する基盤通信プリミティブを提供する。これにより、勾配の集約、パラメータの同期、データのブロードキャストなどが実現される。

**機能の利用シーン**：分散学習の勾配集約時にAllReduceが使用される。また、モデルパラメータの初期同期時にBroadcast、分散評価時のデータ収集にAllGather、デバイス間のデータ交換にAllToAllが使用される。高レベルAPIが内部的に使用するほか、カスタム分散学習ループの構築にも利用可能である。

**主要な処理内容**：
1. AllReduce: 全デバイスのテンソルを指定された演算（Add、Min、Max、Mul）で集約し、結果を全デバイスに配布する
2. AllGather: 全デバイスのテンソルを第一次元方向に結合し、全デバイスに結果を配布する
3. Broadcast（Send/Recv）: 一つのデバイスからグループ内の全デバイスにテンソルを送信する
4. AllToAll: 各デバイスのテンソルをグループ内の全デバイスと相互交換する
5. Communicator初期化: 通信グループの初期化とメンバーシップの確立

**関連システム・外部連携**：NCCL（NVIDIA Collective Communications Library）、gRPCベースのリング通信、デバイス間通信インフラストラクチャ

**権限による制御**：特段の権限制御はないが、group_key/instance_keyによる通信グループの論理的な分離が行われる。

## 関連画面

本機能はバックエンドの通信基盤であるため、直接的な関連画面は存在しない。

## 機能種別

計算処理 / デバイス間通信

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| input (t) | Tensor | Yes | 集約対象のテンソル | bfloat16, float, float16, float64, int32, int64のいずれか |
| group_size | int32 | Yes | グループ内のデバイス数 | 1以上（AllReduceの場合）、2以上（Broadcastの場合） |
| group_key | int32 | Yes | 通信グループの識別子 | 整数 |
| instance_key | int32 | Yes | 操作インスタンスの識別子 | 整数 |
| merge_op | string | No | 集約時の二項演算子（AllReduceのみ） | 'Min', 'Max', 'Mul', 'Add' |
| final_op | string | No | 最終的な単項演算子（AllReduceのみ） | 'Id', 'Div' |
| communication_hint | string | No | 通信方式のヒント | 'auto', 'ring', 'nccl' |
| timeout_seconds | float | No | タイムアウト（秒） | 0以上（0はタイムアウトなし） |
| ordering_token | resource | No | 実行順序制御トークン（v2のみ） | リソーステンソル |
| max_subdivs_per_device | int | No | テンソル分割数の上限 | -1（無効）、0（システムデフォルト） |

### 入力データソース

TensorFlowの計算グラフ内で生成されるテンソル（勾配テンソル、モデルパラメータ等）

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| data | Tensor | 集約結果のテンソル（入力と同じ型） |
| communicator | resource | 通信リソース（initialize_communicatorの場合） |
| group_size | int32 | グループサイズ（assign_group_v2の場合） |
| group_key | int32 | グループキー（assign_group_v2の場合） |

### 出力先

計算グラフ内の後続オペレーションへの入力として使用される

## 処理フロー

### 処理シーケンス

```
1. 通信グループの初期化
   └─ group_key、rank、group_sizeを指定してcommunicatorリソースを作成
2. テンソルの集約操作
   └─ AllReduce: 各デバイスのテンソルをmerge_opで集約、final_opを適用して全デバイスに配布
   └─ AllGather: 各デバイスのテンソルを第一次元方向に結合して全デバイスに配布
   └─ Broadcast: 送信側デバイスからグループ内全デバイスにテンソルを送信
   └─ AllToAll: 各デバイスのテンソルスライスを対応するデバイスに送信
3. タイムアウト監視
   └─ timeout_secondsが非零の場合、指定時間内に完了しなければDeadlineExceededErrorを発生
4. 結果テンソルの返却
   └─ 集約されたテンソルを後続の計算に渡す
```

### フローチャート

```mermaid
flowchart TD
    A[開始] --> B{操作タイプ}
    B -->|AllReduce| C[各デバイスのテンソルをmerge_opで集約]
    B -->|AllGather| D[各デバイスのテンソルを第一次元方向に結合]
    B -->|Broadcast| E[送信側からグループ全体にテンソルを配布]
    B -->|AllToAll| F[各デバイス間でテンソルスライスを交換]
    C --> G{final_op}
    G -->|Id| H[結果をそのまま返却]
    G -->|Div| I[結果をgroup_sizeで除算]
    D --> H
    E --> H
    F --> H
    H --> J[終了]
    I --> J
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-71-01 | グループサイズ検証 | group_sizeはAllReduceで1以上、Broadcastで2以上であること | 操作実行時 |
| BR-71-02 | テンソル型制約 | 入出力テンソルはサポートされた数値型であること | 全操作 |
| BR-71-03 | デバイス一意性 | グループ内の各テンソルは異なるデバイスに配置されること | 全操作 |
| BR-71-04 | 通信一貫性 | 同一group_key/instance_keyを持つ操作は同じグループとして扱われる | 全操作 |
| BR-71-05 | Communicator一回性 | 同一group_keyのcommunicatorは一度のみ初期化可能 | initialize_communicator |

### 計算ロジック

- AllReduce (merge_op='Add', final_op='Div'): 平均値を計算（各デバイスの値を加算後、group_sizeで除算）
- AllGather: 出力の第一次元 = 入力の第一次元 x group_size
- ReduceScatter: AllReduceの結果を第一次元方向にgroup_size分割し、各デバイスに一つずつ配布

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

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

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| ValueError | パラメータ不正 | group_sizeが最低値未満 | group_sizeを適切な値に修正 |
| ValueError | 型不一致 | broadcast_sendでテンソルの形状・型が宣言と不一致 | 正しい形状・型を指定 |
| DeadlineExceededError | タイムアウト | timeout_seconds内に操作が完了しない | タイムアウト値の調整またはデバイスの状態確認 |
| InvalidArgumentError | 通信グループ重複 | 既存のgroup_keyでcommunicatorを再初期化 | 異なるgroup_keyを使用 |

### リトライ仕様

タイムアウト発生時の自動リトライ機能は提供されていない。上位レイヤーでのリトライ実装が必要。

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

コレクティブ操作はアトミックに実行される。全参加デバイスが操作を完了するまでブロックし、一部のデバイスが失敗した場合は全体がエラーとなる。

## パフォーマンス要件

- communication_hint='nccl'を指定することで、NVIDIA GPU間通信においてNCCLの高速通信を利用可能
- communication_hint='ring'はCPU環境でのリング型AllReduceを実行
- max_subdivs_per_deviceにより、テンソルの分割並列処理が可能

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

- group_keyによる通信グループの論理的分離により、異なるグループ間のデータ混入を防止
- gRPC通信層での暗号化設定は別途必要

## 備考

- V1 APIとV2 APIが存在し、V2はテンソル型のgroup_size/group_key/instance_keyをサポート
- V3 APIはcommunicator初期化を使用するより効率的なインターフェースを提供
- ordering_tokenによる自動制御依存関係の管理が可能

---

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

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

### 推奨読解順序

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

C++側のOp登録でサポートされるテンソル型や属性を把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | collective_ops.cc | `tensorflow/core/ops/collective_ops.cc` | REGISTER_OPマクロで登録される各操作の入出力型、属性を理解する |

**読解のコツ**: REGISTER_OPマクロのチェーンメソッド（.Input(), .Output(), .Attr()）でOp定義を読み取る。SetIsStateful()はステートフルOp、SetIsDistributedCommunication()は分散通信Opを示す。

#### Step 2: エントリーポイントを理解する

Python側のAPIエントリーポイントを確認する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | collective_ops.py | `tensorflow/python/ops/collective_ops.py` | Python APIの関数シグネチャ、バリデーション、gen_collective_opsへの委譲を理解する |

**主要処理フロー**:
1. **19-68行目**: all_reduce関数 - group_sizeの検証後、gen_collective_ops.collective_reduceを呼び出す
2. **95-157行目**: all_reduce_v2関数 - ordering_tokenのリスト変換後、collective_reduce_v2を呼び出す
3. **160-197行目**: all_gather関数 - group_sizeの検証後、collective_gatherを呼び出す
4. **250-311行目**: broadcast_send関数 - group_size、shape、dtypeの検証後、collective_bcast_sendを呼び出す
5. **429-465行目**: initialize_communicator関数 - communicatorリソースを生成する
6. **468-498行目**: all_reduce_v3関数 - communicatorベースの新APIでreduceを実行する
7. **501-551行目**: all_to_all_v2関数 - デバイス間のテンソル相互交換
8. **554-578行目**: all_to_all_v3関数 - communicatorベースの新APIでall-to-allを実行する

#### Step 3: C++側のOp定義を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | collective_ops.cc | `tensorflow/core/ops/collective_ops.cc` | 各Opの形状推論ロジックを確認する |

**主要処理フロー**:
- **22-37行目**: CollectiveReduce Op - 入力と同じ形状を出力（UnchangedShape）
- **39-82行目**: CollectiveGather Op - 第一次元がgroup_size倍になる形状推論
- **84-96行目**: CollectiveBcastSend Op - ExplicitShape（明示的な形状指定）
- **127-144行目**: CollectiveReduceV2 Op - V2版、ordering_tokenサポート
- **146-174行目**: CollectiveReduceScatterV2 Op - 第一次元が不明（group_size依存）
- **236-246行目**: CollectiveInitializeCommunicator Op - スカラーリソースを出力
- **248-258行目**: CollectiveReduceV3 Op - communicatorベース
- **260-274行目**: CollectiveAllToAllV2 Op - 入力と同じ形状を出力

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

```
collective_ops.py
    │
    ├─ all_reduce() / all_reduce_v2() / all_reduce_v3()
    │      └─ gen_collective_ops.collective_reduce / collective_reduce_v2 / collective_reduce_v3
    │             └─ CollectiveReduce / CollectiveReduceV2 / CollectiveReduceV3 (C++ Op)
    │
    ├─ all_gather() / all_gather_v2()
    │      └─ gen_collective_ops.collective_gather / collective_gather_v2
    │             └─ CollectiveGather / CollectiveGatherV2 (C++ Op)
    │
    ├─ broadcast_send() / broadcast_send_v2()
    │      └─ gen_collective_ops.collective_bcast_send / collective_bcast_send_v2
    │             └─ CollectiveBcastSend / CollectiveBcastSendV2 (C++ Op)
    │
    ├─ broadcast_recv() / broadcast_recv_v2()
    │      └─ gen_collective_ops.collective_bcast_recv / collective_bcast_recv_v2
    │             └─ CollectiveBcastRecv / CollectiveBcastRecvV2 (C++ Op)
    │
    ├─ all_to_all_v2() / all_to_all_v3()
    │      └─ gen_collective_ops.collective_all_to_all_v2 / collective_all_to_all_v3
    │             └─ CollectiveAllToAllV2 / CollectiveAllToAllV3 (C++ Op)
    │
    └─ initialize_communicator()
           └─ gen_collective_ops.collective_initialize_communicator
                  └─ CollectiveInitializeCommunicator (C++ Op)
```

### データフロー図

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

各デバイスの            AllReduce: merge_op + final_op      集約済みテンソル
テンソル ───────▶       AllGather: 第一次元結合          ───▶ （全デバイスに配布）
                        Broadcast: 1対多送信
                        AllToAll: 相互交換
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| collective_ops.py | `tensorflow/python/ops/collective_ops.py` | ソース | Python API定義（全操作のエントリーポイント） |
| collective_ops.cc | `tensorflow/core/ops/collective_ops.cc` | ソース | C++ Op登録・形状推論 |
| gen_collective_ops | 自動生成 | ソース | Python-C++ブリッジ（自動生成コード） |
