"""
=============================================================================
CS349E · 第 9 讲 · 02_energy_model.py —— 能量模型：每一级存储的访问次数 × 每次的能量
=============================================================================

【这个脚本要回答的问题】
    Interstellar 用一个简单的公式比较各种加速器：E = Σ_级 (访问次数 × 每次访问的能量) + 乘加次数 × 每次乘加的能量。
      演示 1：Interstellar 表 3 的能量表（28 nm、16 位）：寄存器堆、SRAM、DRAM 各是一次乘加的多少倍。
      演示 2：一个 Transformer 线性层（256 个 token）在 Eyeriss 式的配置上最好的映射：能量花在哪一级（论文图 11）。
      演示 3：同一层换成 64 B 的寄存器堆：能量怎样重新分配（论文的观察 2：好的设计不应该被某一级主导）。
      演示 4：批大小变小以后：DRAM 主导，而且是"躲不掉"的（观察 2 的例外）。

【运行】python3 02_energy_model.py      （约 2 秒）
【依赖】numpy；本目录的 loopnest；../lecture01 的 console。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from loopnest import E_DRAM, E_HOP, E_MAC, RF_ENERGY, SRAM_ENERGY, Arch, search_fast  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402

LAYER = {"M": 256, "N": 512, "K": 512}          # 256 个 token 的 512×512 投影（一个小 Transformer 的 prefill）
LEVELS = ("MAC", "RF", "array", "GB", "DRAM")


def breakdown_row(label, shape, arch, rows_dims=("K",), cols_dims=("N",)):
    m, e, _, _ = search_fast(shape, arch, rows_dims, cols_dims)
    macs = shape["M"] * shape["N"] * shape["K"]
    cells = [label, f"{e['total'] / macs:.2f}"]
    for k in LEVELS:
        cells.append(f"{e[k] / macs:.2f}（{e[k] / e['total']:.0%}）")
    return cells, m, e


# -----------------------------------------------------------------------------
def demo1_table():
    print_header("演示 1：Interstellar 表 3——每次 16 位访问的能量（28 nm）")
    rows = [["一次乘加（MAC）", f"{E_MAC}", "1×"], ["PE 之间传一次（hop）", f"{E_HOP}", f"{E_HOP / E_MAC:.1f}×"]]
    for b, e in RF_ENERGY.items():
        rows.append([f"寄存器堆 {b} B", f"{e}", f"{e / E_MAC:.1f}×"])
    for b, e in SRAM_ENERGY.items():
        rows.append([f"SRAM {b // 1024} KB", f"{e}", f"{e / E_MAC:.0f}×"])
    rows.append(["DRAM", f"{E_DRAM}", f"{E_DRAM / E_MAC:.0f}×"])
    print_table(["存储 / 操作", "能量 pJ", "是一次乘加的"], rows)
    print_observations([
        "一次乘加只要 0.075 pJ；读一次 512 B 的寄存器堆就是它的 13 倍，读一次 128 KB 的 SRAM 是 180 倍，读一次 DRAM 是 2,667 倍。",
        "存储越大，每次访问越贵：寄存器堆的能量和大小成正比，SRAM 大小每翻一倍能量 ×1.5。所以\"多放一点数据在近处\"不是免费的——"
        "这正是 Interstellar 要回答的问题：每一级该多大。",
        "（和第 1 讲 Horowitz 的 45 nm 表比，28 nm、而且存储做了大量分 bank，存储和运算的比例更悬殊。）",
    ])


# -----------------------------------------------------------------------------
def demo2_eyeriss():
    print_header("演示 2：Eyeriss 式的配置（16×16 PE、512 B 寄存器堆、128 KB SRAM）上，最好的映射把能量花在哪（单位 pJ/乘加）")
    cells, m, _ = breakdown_row("256×512×512，数据流 K|N", LAYER, Arch())
    print_table(["层", "总能量"] + list(LEVELS), [cells])
    print(f"\n最好的映射：RF 层因子 {m.rf}、GB 层 {m.gb}、DRAM 层 {m.dram}；GB 层循环（外→内）{m.order_gb}、DRAM 层 {m.order_dram}")
    print_observations([
        "每次乘加要从寄存器堆读 A、B、读写 C 共 4 次，512 B 的寄存器堆每次 0.96 pJ——寄存器堆一项就占了近一半的能量。",
        "DRAM 占第二大头：128 KB 的 SRAM 放不下足够大的块，一部分数据要从 DRAM 搬不止一次（01 演示 3）。",
        "运算本身（MAC）只占约 1%。和第 8 讲的结论一样：能量几乎全花在搬数据上，而且花在**离运算最近**的那一级——"
        "因为它被访问的次数最多。",
    ])


# -----------------------------------------------------------------------------
def demo3_small_rf():
    print_header("演示 3：同一层，寄存器堆从 512 B 缩到 64 B（其他不变）")
    rows = []
    for rf in (512, 64):
        cells, m, _ = breakdown_row(f"{rf} B 寄存器堆", LAYER, Arch(rf_bytes=rf))
        rows.append(cells + [str(m.rf)])
    print_table(["配置", "总能量"] + list(LEVELS) + ["RF 层因子"], rows)
    print_observations([
        "寄存器堆小了 8 倍，每次访问便宜 8 倍；它装得下的块变小，复用少了一些，GB 的访问变多——但总能量从约 8.8 降到约 5.1 pJ/乘加，"
        "省了约 42%。Interstellar 在 AlexNet 上测到最多 2.6 倍（论文图 12），方向一致。",
        "论文的观察 2：\"一个高效系统的总能量不应该被存储层次中的某一级主导\"——某一级太大太贵，就该把它缩小、把复用推到别的级去。"
        "寄存器堆缩小以后不再是大头，可 DRAM 变成了主导（69%），而且是\"每个数只搬一次\"下限（1.56 pJ/乘加）的 2.3 倍："
        "下一步该做的是让 DRAM 的访问更少——加大 SRAM（04 脚本）。",
    ])


# -----------------------------------------------------------------------------
def demo4_small_batch():
    print_header("演示 4：批大小变小（4096×4096 的全连接层，64 B 寄存器堆、128 KB SRAM）")
    rows = []
    for M in (256, 16, 1):
        shape = {"M": M, "N": 4096 if M < 256 else 1024, "K": 4096 if M < 256 else 1024}
        cells, _, _ = breakdown_row(f"批 {M}（{shape['K']}×{shape['N']}）", shape, Arch(rf_bytes=64))
        macs = M * shape["N"] * shape["K"]
        compulsory = (M * shape["K"] + shape["K"] * shape["N"] + M * shape["N"]) * E_DRAM / macs
        rows.append(cells + [f"{compulsory:.2f}"])
    print_table(["层", "总能量"] + list(LEVELS) + ["DRAM 的下限（每个数只搬一次）"], rows)
    print_observations([
        "批 1 时每个权重只被用一次：DRAM 的能量 ≈ 200 pJ/乘加，也就是\"每个权重从 DRAM 读一次\"的下限——怎么分块、选什么数据流都省不掉。",
        "这是观察 2 的例外：DRAM 主导，但 DRAM 的访问次数已经是最少的了。这时整个系统受限于内存（Amdahl 定律：别的地方再优化也没用），"
        "片上的数据流设计基本不重要。论文专门指出 MLP、LSTM 和批 1 的推理属于这一类——和 TPU 论文（第 8 讲 §4）、"
        "LLM 的 decode（第 1 讲 §4.2）是同一个结论。",
        "批 16 时 DRAM 也接近下限（12.94 vs 12.60）。批 256 的 1024×1024 层就不一样了：DRAM 仍占 67%，却是下限的 2.8 倍——"
        "这部分不是躲不掉的，而是 128 KB 的缓冲区太小、放不下足够大的块（04 脚本加大缓冲区）。",
    ])


if __name__ == "__main__":
    demo1_table()
    demo2_eyeriss()
    demo3_small_rf()
    demo4_small_batch()