"""
=============================================================================
CS349E · 第 9 讲 · 04_memory_hierarchy.py —— 每一级存储该多大（Interstellar 的观察 2 与资源分配）
=============================================================================

【这个脚本要回答的问题】
    03 看到数据流不太重要、分块很重要。那硬件设计者能决定什么？——**每一级存储有多大**。
      演示 1：寄存器堆（每个 PE）× 全局缓冲（SRAM）的大小一起扫一遍：一个小 Transformer 层的总能量（论文图 12）。
      演示 2：最好的配置和 Eyeriss 式的配置比：能量怎样在各级之间重新分配。
      演示 3：PE 阵列变大时，最优的存储大小怎样变（论文图 13）；以及"相邻两级的容量之比约 4–16"这条经验规则。

【负载】一个 d = 512 的小 Transformer 层在 256 个 token 上的线性层：Q/K/V/O 四个 512×512、MLP 的两个 512×2048 和一个 2048×512
    （每个都按数据流 K|N、搜索最好的分块）。能量是整层加起来、除以总乘加次数。

【运行】python3 04_memory_hierarchy.py      （约 15 秒：演示 1、3 一共做了约 300 次映射搜索）
【依赖】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 Arch, search_fast  # noqa: E402
from console import fmt_bytes, print_header, print_observations, print_table  # noqa: E402

LAYER = [({"M": 256, "N": 512, "K": 512}, 4), ({"M": 256, "N": 2048, "K": 512}, 2), ({"M": 256, "N": 512, "K": 2048}, 1)]
LEVELS = ("MAC", "RF", "array", "GB", "DRAM")


def layer_energy(arch):
    """整层的能量（pJ/乘加）和各级的分解；有一个矩阵乘法找不到合法映射时返回 (None, None)。"""
    tot, macs, br = 0.0, 0, dict.fromkeys(LEVELS, 0.0)
    for shape, cnt in LAYER:
        m, e, _, _ = search_fast(shape, arch, ("K",), ("N",))
        if m is None:
            return None, None
        tot += cnt * e["total"]
        macs += cnt * shape["M"] * shape["N"] * shape["K"]
        for k in LEVELS:
            br[k] += cnt * e[k]
    return tot / macs, {k: v / macs for k, v in br.items()}


# -----------------------------------------------------------------------------
def demo1_sweep():
    print_header("演示 1：寄存器堆 × 全局缓冲的大小（16×16 PE）——整层的能量（pJ/乘加）")
    rfs = (16, 32, 64, 128, 256, 512)
    gbs = (64 * 1024, 128 * 1024, 256 * 1024, 512 * 1024)
    rows = []
    best = None
    grid = {}
    for gb in gbs:
        cells = [f"SRAM {gb // 1024} KB"]
        for rf in rfs:
            e, _ = layer_energy(Arch(rf_bytes=rf, gb_bytes=gb))
            grid[(rf, gb)] = e
            cells.append("-" if e is None else f"{e:.2f}")
            if e is not None and (best is None or e < best[0]):
                best = (e, rf, gb)
        rows.append(cells)
    print_table(["全局缓冲 \\ 寄存器堆"] + [f"{r} B" for r in rfs], rows)
    ey = grid[(512, 128 * 1024)]
    print(f"\nEyeriss 式的配置（512 B、128 KB）：{ey:.2f}；最好的配置（{best[1]} B、{best[2] // 1024} KB）：{best[0]:.2f}，"
          f"好 {ey / best[0]:.2f} 倍。只缩寄存器堆（32 B、128 KB）：{grid[(32, 128 * 1024)]:.2f}，好 {ey / grid[(32, 128 * 1024)]:.2f} 倍。")
    print_observations([
        "寄存器堆从 512 B 缩到 32–64 B，能量降约 40%；论文在 AlexNet 上测到最多 2.6 倍（图 12）。"
        "但也不能太小：16 B 时装不下足够大的块，复用不够，反而比 32 B 差——存在一个最优点。",
        "全局缓冲越大越好，但收益递减：从 256 KB 到 512 KB 只再降约 6%。论文：\"超过 256 KB 以后收益可以忽略\"，考虑到面积，不一定需要很大的缓冲。",
        "这些都是硬件设计者能决定的事；数据流（03 脚本）反而影响不大。Interstellar 的结论：资源分配（尤其是存储层次）比数据流重要。",
    ])
    return best


# -----------------------------------------------------------------------------
def demo2_breakdown(best):
    print_header("演示 2：能量在各级之间怎样重新分配（pJ/乘加，括号里是占比）")
    rows = []
    for label, rf, gb in (("Eyeriss 式：512 B、128 KB", 512, 128 * 1024), ("只缩寄存器堆：64 B、128 KB", 64, 128 * 1024),
                          (f"最好：{best[1]} B、{best[2] // 1024} KB", best[1], best[2])):
        e, br = layer_energy(Arch(rf_bytes=rf, gb_bytes=gb))
        rows.append([label, f"{e:.2f}"] + [f"{br[k]:.2f}（{br[k] / e:.0%}）" for k in LEVELS])
    print_table(["配置", "总能量"] + list(LEVELS), rows)
    print_observations([
        "Eyeriss 式的配置被寄存器堆主导（约一半）；缩小寄存器堆以后，DRAM 变成主导；再加大全局缓冲，DRAM 降下来，"
        "各级的能量更均衡——论文的观察 2：总能量不应该被某一级主导。",
        "最后剩下的 DRAM 能量里有一部分是躲不掉的（每个数至少搬一次）。这时再加大缓冲就没有用了。",
    ])


# -----------------------------------------------------------------------------
def demo3_array_size():
    print_header("演示 3：PE 阵列变大时，最优的存储配置（论文图 13）")
    rows = []
    for n in (8, 16, 32):
        best = None
        for gb in (64 * 1024, 128 * 1024, 256 * 1024, 512 * 1024, 1024 * 1024):
            for rf in (16, 32, 64, 128, 256):
                e, _ = layer_energy(Arch(rows=n, cols=n, rf_bytes=rf, gb_bytes=gb))
                if e is not None and (best is None or e < best[0]):
                    best = (e, rf, gb)
        e, rf, gb = best
        total_rf = rf * n * n
        rows.append([f"{n}×{n}（{n * n} 个 PE）", f"{rf} B", fmt_bytes(total_rf, binary=True), f"{gb // 1024} KB",
                     f"{gb / total_rf:.0f}", f"{e:.2f}"])
    print_table(["阵列", "每个 PE 的寄存器堆", "寄存器堆总量", "全局缓冲", "缓冲 / 寄存器堆总量", "能量 pJ/乘加"], rows)
    print("（全局缓冲最多扫到 1 MB。）")
    print_observations([
        "PE 多了 16 倍（8×8 → 32×32），最优的全局缓冲只大了 2 倍：存储随 PE 数**次线性**地增长——因为每一级越大、每次访问越贵，"
        "不值得为每个 PE 保留同样多的数据（论文图 13 的结论）。注意 32×32 的最优点碰到了我们扫描的上限 1 MB，真正的最优可能更大一些。"
        "每个 PE 的寄存器堆则大致保持不变（64–128 B）。",
        "能量随阵列变大而下降：更多的 PE 共享同一份数据（多播、空间归约），全局缓冲的访问被摊薄。论文说下降得\"很少\"，"
        "我们这个模型里降得多一些（约三分之一）——我们的 hop 能量很便宜，而且没有计入大阵列更长的导线。",
        "\"缓冲 / 寄存器堆总量\"在 16–64 之间。论文的经验规则：相邻两级存储的容量之比约 4–16 时最好"
        "（它加了一级 256 B 的二级寄存器堆，16 B → 256 B → 256 KB，AlexNet 再省 25%）。我们只有两级片上存储，比值偏大。",
    ])


if __name__ == "__main__":
    BEST = demo1_sweep()
    demo2_breakdown(BEST)
    demo3_array_size()