# 機能設計書 126-NCCL集約操作

## 概要

本ドキュメントは、TensorFlowにおけるNCCL（NVIDIA Collective Communications Library）集約操作の機能設計を記述する。GPU間のAllReduce、Broadcast、Reduce操作を提供し、マルチGPU環境での効率的な集約通信を実現する。

### 本機能の処理概要

**業務上の目的・背景**：データ並列分散学習において、各GPU上で計算された勾配を集約する処理はボトルネックとなりやすい。NCCLはNVIDIA GPUに最適化された集約通信ライブラリであり、GPU間の高帯域幅通信（NVLink、PCIe等）を活用して勾配集約のスループットを最大化する。

**機能の利用シーン**：マルチGPUでのデータ並列分散学習における勾配のAllReduce、パラメータサーバからの重みのBroadcast、勾配の集約Reduceなどに使用される。TensorFlowのMirroredStrategy内部で自動的に活用される。

**主要な処理内容**：
1. NcclAllReduce: 全GPU間でテンソルを集約（sum, prod, min, max）
2. NcclReduce: 複数GPUのテンソルを1つのGPUに集約
3. NcclBroadcast: 1つのGPUのテンソルを全GPUに配信
4. グラフ最適化段階でのOp分解（NcclReduce -> _NcclReduceSend + _NcclReduceRecv）

**関連システム・外部連携**：NVIDIA NCCL 2.xライブラリ、CUDA/ROCmランタイム、NcclManager（`tensorflow/core/nccl/nccl_manager.cc`）。

**権限による制御**：特になし。NVIDIA GPU搭載環境で利用可能。

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | - | - | 本機能はバックエンドGPU通信操作であり、直接関連する画面はない |

## 機能種別

データ連携（GPU間集約通信）

## 入力仕様

### 入力パラメータ

#### NcclAllReduce

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| input | T | Yes | 集約対象のテンソル | half/float/float64/int32/int64 |
| reduction | string | Yes | 集約操作の種類 | min/max/prod/sum |
| num_devices | int | Yes | 参加デバイス数 | - |
| shared_name | string | Yes | 同一集約操作の識別子 | - |

#### NcclReduce

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| input | num_devices * T | Yes | 各デバイスからの入力テンソル | half/float/float64/int32/int64 |
| reduction | string | Yes | 集約操作の種類 | min/max/prod/sum |
| num_devices | int | Yes | 参加デバイス数 | - |

#### NcclBroadcast

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| input | T | Yes | ブロードキャスト対象のテンソル | half/float/float64/int32/int64 |
| shape | shape | Yes | 出力テンソルの形状 | - |

### 入力データソース

各GPUデバイス上で計算されたテンソル（勾配、パラメータなど）。

## 出力仕様

### 出力データ

#### NcclAllReduce

| 項目名 | 型 | 説明 |
|--------|-----|------|
| data | T | 集約結果のテンソル（入力と同じ形状） |

#### NcclReduce

| 項目名 | 型 | 説明 |
|--------|-----|------|
| data | T | 集約結果のテンソル（1デバイスのみ） |

#### NcclBroadcast

| 項目名 | 型 | 説明 |
|--------|-----|------|
| output | T | ブロードキャストされたテンソル |

### 出力先

各GPUデバイスのメモリ上に集約結果が格納される。

## 処理フロー

### 処理シーケンス

```
1. 集約操作の初期化
   └─ NCCLコミュニケータの確立（shared_nameによるグループ化）
2. 各デバイスからの入力テンソル取得
   └─ GPUメモリ上のテンソルを参照
3. NCCL集約通信の実行
   └─ NVLink/PCIe経由でGPU間の直接データ転送
4. 結果の配置
   └─ AllReduce: 全GPUに結果を配置
   └─ Reduce: 指定デバイスに結果を配置
   └─ Broadcast: 全GPUにデータを配信
```

### フローチャート

```mermaid
flowchart TD
    A[GPU0: input] --> D[NcclAllReduce]
    B[GPU1: input] --> D
    C[GPU2: input] --> D
    D --> E[GPU0: result]
    D --> F[GPU1: result]
    D --> G[GPU2: result]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-126-1 | 型制限 | half, float, float64, int32, int64のみサポート | 全Op |
| BR-126-2 | デバイス一致 | 同一shared_nameの全参加デバイスが揃うまでブロック | AllReduce/Reduce |
| BR-126-3 | Op分解 | NcclReduceはグラフ最適化時に_NcclReduceSend + _NcclReduceRecvに分解 | NcclReduce |
| BR-126-4 | Op分解 | NcclBroadcastはグラフ最適化時に_NcclBroadcastSend + _NcclBroadcastRecvに分解 | NcclBroadcast |
| BR-126-5 | 形状不変 | AllReduceの出力形状は入力形状と同一 | NcclAllReduce |

### 計算ロジック

- **sum**: result = input_0 + input_1 + ... + input_n
- **prod**: result = input_0 * input_1 * ... * input_n
- **min**: result = element-wise min(input_0, input_1, ..., input_n)
- **max**: result = element-wise max(input_0, input_1, ..., input_n)

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

### 操作別データベース影響一覧

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

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| Internal | NCCLエラー | NCCL操作に失敗 | NCCL_DEBUG=WARNで詳細確認 |
| Internal | CUDAエラー | GPUデバイスエラー | デバイスの状態を確認 |
| Deadlock | デッドロック | 参加デバイス数が揃わない | num_devicesと実際のデバイス数を一致させる |

### リトライ仕様

NCCLエラーの場合、NCCLコミュニケータの再初期化が必要となる場合がある。

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

トランザクション管理なし。SetIsStateful()により副作用Opとしてマーク。

## パフォーマンス要件

NCCLはNVLink対応環境で最大300GB/sの帯域幅を提供。PCIe環境でもリングアルゴリズムにより効率的なAllReduceを実現。

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

GPU間の直接メモリ転送を行うため、同一マシン上のGPU間でのデータ露出に注意。

## 備考

- NcclReduceとNcclBroadcastはカーネル実装を持たず、グラフ最適化段階でSend/Recv Opに分解される
- _NcclReduceSend/_NcclReduceRecv、_NcclBroadcastSend/_NcclBroadcastRecvはプライベートOp（先頭アンダースコア）
- ROCm環境でも利用可能（hipifyによるCUDA→ROCm変換）

---

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

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

### 推奨読解順序

#### Step 1: Op定義を理解する

公開Op（NcclAllReduce, NcclReduce, NcclBroadcast）と内部Op（_NcclReduceSend等）の関係を把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | nccl_ops.cc | `tensorflow/core/ops/nccl_ops.cc` | 8つのOp定義と公開/内部Opの区別 |

**読解のコツ**: 先頭アンダースコア（`_NcclReduceSend`等）はTensorFlow内部Opを示す。これらはユーザが直接使用するのではなく、グラフ最適化段階でNcclReduce/NcclBroadcastから自動的に生成される。

#### Step 2: 公開Opを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | nccl_ops.cc | `tensorflow/core/ops/nccl_ops.cc` | NcclAllReduce, NcclReduce, NcclBroadcast |

**主要処理フロー**:
1. **24-32行目**: NcclAllReduce - 全GPU間集約。shared_nameで同一操作を識別
2. **36-43行目**: NcclReduce - カーネル未実装。num_devices * Tのリスト入力
3. **94-100行目**: NcclBroadcast - カーネル未実装。shape属性で出力形状を指定

#### Step 3: 内部Opを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | nccl_ops.cc | `tensorflow/core/ops/nccl_ops.cc` | _NcclReduceSend/Recv, _NcclBroadcastSend/Recv |

**主要処理フロー**:
- **45-65行目**: _NcclReduceSend - num_devices-1台が実行。出力なし（NoOutputs）
- **67-90行目**: _NcclReduceRecv - 1台が実行。集約結果を出力
- **102-121行目**: _NcclBroadcastSend - 1台が実行。出力なし
- **123-149行目**: _NcclBroadcastRecv - num_devices-1台が実行。ブロードキャストデータを受信

#### Step 4: NcclManagerの実装を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 4-1 | nccl_manager.cc | `tensorflow/core/nccl/nccl_manager.cc` | NCCLライブラリとの統合 |

**主要処理フロー**:
- **19行目**: GOOGLE_CUDAまたはTENSORFLOW_USE_ROCMの条件コンパイル
- **48-55行目**: NCCL_RETURN_IF_ERRORマクロ - NCCLエラーのハンドリング
- **37-46行目**: ROCm環境でのhipify（CUDA→ROCm型変換）

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

```
nccl_ops.cc (Op定義)
    |
    +-- REGISTER_OP("NcclAllReduce")        ... 公開Op
    |
    +-- REGISTER_OP("NcclReduce")           ... 公開Op（カーネルなし）
    |     +-- グラフ最適化で分解 →
    |         +-- _NcclReduceSend (num_devices-1台)
    |         +-- _NcclReduceRecv (1台)
    |
    +-- REGISTER_OP("NcclBroadcast")        ... 公開Op（カーネルなし）
          +-- グラフ最適化で分解 →
              +-- _NcclBroadcastSend (1台)
              +-- _NcclBroadcastRecv (num_devices-1台)

nccl_manager.cc (NCCL統合)
    +-- NCCLコミュニケータ管理
    +-- ncclAllReduce / ncclReduce / ncclBroadcast 呼び出し

nccl_rewrite.cc (グラフ書き換え)
    +-- NcclReduce → _NcclReduceSend + _NcclReduceRecv
    +-- NcclBroadcast → _NcclBroadcastSend + _NcclBroadcastRecv
```

### データフロー図

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

GPU0: tensor ───┐
GPU1: tensor ───┼──▶ NcclAllReduce ──────────▶ GPU0: reduced_tensor
GPU2: tensor ───┘     (NCCL集約通信)             GPU1: reduced_tensor
                                                  GPU2: reduced_tensor

GPU0: tensor ───▶ NcclBroadcast ──────────────▶ GPU0: tensor (同一)
                   (NCCLブロードキャスト)          GPU1: tensor (コピー)
                                                  GPU2: tensor (コピー)
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| nccl_ops.cc | `tensorflow/core/ops/nccl_ops.cc` | ソース | Op定義（8つのOp） |
| nccl_manager.cc | `tensorflow/core/nccl/nccl_manager.cc` | ソース | NCCLライブラリ統合マネージャ |
| nccl_rewrite.cc | `tensorflow/core/nccl/nccl_rewrite.cc` | ソース | グラフ最適化時のOp書き換え |
| collective_communicator.cc | `tensorflow/core/nccl/collective_communicator.cc` | ソース | 集約通信インターフェース |
| common_shape_fns.h | `tensorflow/core/framework/common_shape_fns.h` | ヘッダ | 形状推論ユーティリティ |
