(MAX を導入する意味)=
# MAX を導入する意味

## この章で学ぶこと

- なぜ microgpt.py / Mojo 版のままでは速さに限界があるのか
- MAX が解決しようとしている問題とは何か
- Graph・Module・Weight という三つの中心概念
- MAX と Mojo の役割分担

---

## microgpt.py と Mojo 版の限界

{ref}`microgpt.py を Mojo で書き直す` で実装した Mojo 版は、すべての演算を**スカラー単位でテープに追記**していきます。

```
t.add(a, b)   → テープにノードを 1 個追記
t.mul(c, d)   → テープにノードを 1 個追記
...           → 推論 1 ステップで数千〜数万ノードが積み上がる
```

これは自動微分（学習）のためには正しい設計ですが、**推論だけを速く動かしたい場合には非効率**です。

| 実装 | 特徴 | 問題 |
|------|------|------|
| Python 版 | 短く読みやすい | Python インタプリタのオーバーヘッド |
| Mojo 版 （Tape） | 型安全・自動微分可能 | スカラーループ、JIT/SIMD 最適化なし |
| **MAX** | グラフを先に構築してから実行 | → 次節で説明 |

---

## MAX とは何か

**MAX**（Modular Accelerated Execution）は、モデルの計算を**グラフとして表し、実行しやすい形へ整理する基盤**です。

microgpt.py のような実行方式と MAX の実行方式を比べると次のようになります。

:::: {container} mermaid-flow-half

```{mermaid}
flowchart LR
    subgraph "Python / Mojo Tape 方式"
        direction TB
        A1["入力"] --> B1["演算1<br/>（その場で実行）"]
        B1 --> C1["演算2<br/>（その場で実行）"]
        C1 --> D1["出力"]
    end
    subgraph "MAX 方式"
        direction TB
        A2["Graph 構築<br/>（形だけ定義）"] --> B2["コンパイル<br/>（最適化）"]
        B2 --> C2["実行<br/>（高速）"]
    end
```

::::

MAX では先に「**何をどんな順で計算するか**」というグラフを組み立て、その後でコンパイル・最適化してから実行します。このため：

- **演算の融合**（複数の小さな演算を1回のカーネルにまとめる）
- **SIMD/GPU への自動マッピング**
- **メモリレイアウトの最適化**

が可能になります。

---

## 三つの中心概念

### Graph — 計算の設計図

`Graph` は「何を入力にし、どんな演算を経て、何を出力するか」を**実行前に定義する設計図**です。

```
Graph
├── 入力ノード（TensorType で型・形状を宣言）
├── 演算ノード（ops.matmul, ops.relu, ...）
└── 出力ノード
```

グラフを定義しただけでは何も実行されません。`InferenceSession` に渡してコンパイルして初めて動きます。

### Module — 層や重みを持つまとまり

`Module` は、PyTorch の `nn.Module` に近い概念です。層・重み・計算の流れをひとまとまりで表します。

microgpt の例で言えば、次のような単位を `Module` として整理できます。
```
GPTModule
├── wte  (Weight: トークン埋め込み)
├── wpe  (Weight: 位置埋め込み)
├── AttentionModule (サブモジュール)
├── MLPModule       (サブモジュール)
└── lm_head (Weight: 出力線形層)
```

### Weight — 学習される値を明示する

`Weight` は、モデルの学習可能なパラメータを表します。入力テンソルと区別して明示することで、「どれが固定値でどれが更新される値か」がグラフ上で明確になります。

---

## MAX と Mojo の役割分担

MAX と Mojo は競合ではなく、**役割を分けて組み合わせる**設計になっています。

```{mermaid}
flowchart LR
    A["モデルの流れを整理<br/>(Graph / Module / Weight)"] -- MAX --> B["実行・最適化"]
    C["カスタム演算<br/>(独自のカーネル)"] -- Mojo カスタム ops --> B
```

| 役割 | 担当 |
|------|------|
| モデル全体の計算グラフを組む | MAX（Python API） |
| 標準にない演算を高速実装する | Mojo（custom ops） |
| 学習（自動微分・Adam） | Mojo の Tape（または別途 MAX Engine） |
| 推論の高速実行 | MAX（InferenceSession） |

典型的な進め方は「**まず標準の MAX ops でグラフを組み、ボトルネックになった部分だけを Mojo の custom ops に置き換える**」です。

---

## まとめ

- Python ループ / Mojo Tape のスカラー実行は学習には適しているが、推論の高速化には限界がある
- MAX は計算グラフを先に構築し、コンパイル・最適化してから実行する
- `Graph`（設計図）・`Module`（まとまり）・`Weight`（学習パラメータ）が三つの中心概念
- MAX と Mojo は競合せず役割分担する

次章では MAX の Python API を使った具体的なコード例を示します。

## 次に読む章

{numref}`MAX でコードを書く`（{ref}`MAX でコードを書く`）へ進みます。