"""
=============================================================================
CS349E · 第 9 讲 · 03_dataflow_vs_blocking.py —— 数据流重要，还是分块重要？（Interstellar 的观察 1）
=============================================================================

【这个脚本要回答的问题】
    加速器论文常常为自己的"数据流"（权重驻留、输出驻留、行驻留……）辩护。Interstellar 把它们放在同一个框架里公平比较：
      演示 1：同样的硬件、每种数据流都用它自己最好的分块——最低能量差多少？
      演示 2：同一种数据流，分块方案之间差多少？（论文：只有 30% 的分块方案在最低能量的 1.25 倍以内）
      演示 3：数据流影响的是利用率：维度太小填不满阵列时，"复制"（replication，一个方向上展开两个维度）把利用率补回来。

【数据流的记号】rows|cols：哪些维度展开到阵列的行、哪些到列（Interstellar 的 U|V）。矩阵乘法 C[M,N] += A[M,K]·B[K,N] 里：
    K|N：每个 PE 存一个 B[k,n]——权重驻留，TPU 式（Interstellar 的 C|K）；
    M|N：每个 PE 负责一个输出 C[m,n]——输出驻留；
    M|K：每个 PE 存一个 A[m,k]——输入驻留；以及把行列对调的三种。

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

import os
import sys

import numpy as np

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, utilization  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402

DATAFLOWS = [(("K",), ("N",), "权重驻留（TPU 式）"), (("N",), ("K",), "权重驻留（行列对调）"),
             (("M",), ("N",), "输出驻留"), (("N",), ("M",), "输出驻留（对调）"),
             (("M",), ("K",), "输入驻留"), (("K",), ("M",), "输入驻留（对调）")]


# -----------------------------------------------------------------------------
def demo1_dataflows():
    print_header("演示 1：每种数据流配上它自己最好的分块（16×16 PE、64 B 寄存器堆、128 KB SRAM）——最低能量（pJ/乘加）")
    arch = Arch(rf_bytes=64)
    shapes = [("prefill：256×512×512", {"M": 256, "N": 512, "K": 512}),
              ("批 16 的全连接：16×1024×1024", {"M": 16, "N": 1024, "K": 1024})]
    rows = []
    best_all = {}
    for rd, cd, name in DATAFLOWS:
        cells = [f"{''.join(rd)}|{''.join(cd)}", name]
        for label, shape in shapes:
            _, e, _, _ = search_fast(shape, arch, rd, cd)
            macs = shape["M"] * shape["N"] * shape["K"]
            cells.append(f"{e['total'] / macs:.2f}")
            cells.append(f"{utilization(shape, rd, cd, arch):.0%}")
            best_all.setdefault(label, []).append(e["total"] / macs)
        rows.append(cells)
    heads = ["数据流", "名字"]
    for label, _ in shapes:
        heads += [label, "利用率"]
    print_table(heads, rows)
    for label, v in best_all.items():
        print(f"{label}：最好 / 最差的数据流差 {max(v) / min(v):.2f} 倍")
    print_observations([
        "每种数据流都配上最适合它的分块以后，最低能量只差 1%–2%：Interstellar 的观察 1——\"只要分块和复制用得好，"
        "很多不同的数据流都能达到相近的、接近最优的能效\"。",
        "原因：矩阵乘法（和卷积）里每个数都有很大的复用潜力，好的分块能让绝大多数访问发生在最小、最便宜的那一级；"
        "数据在 PE 之间怎么传（数据流）只占总能量的很小一部分（02 脚本里 array 一项不到 1%）。",
        "批 16 的全连接层里 M 只有 16：把 M 展开到行上恰好填满，所以这里利用率都是 100%。维度不够大时利用率会受影响（演示 3）。",
    ])


# -----------------------------------------------------------------------------
def demo2_blocking():
    print_header("演示 2：同一种数据流（K|N），所有合法的分块方案的能量分布（256×512×512，64 B 寄存器堆、128 KB SRAM）")
    arch = Arch(rf_bytes=64)
    shape = {"M": 256, "N": 512, "K": 512}
    _, _, n, all_e = search_fast(shape, arch, ("K",), ("N",))
    macs = shape["M"] * shape["N"] * shape["K"]
    v = np.sort(all_e) / macs
    lo = v[0]
    rows = [["最低", f"{lo:.2f}", "1.00×"]]
    for q in (1, 10, 25, 50, 75, 90):
        x = float(np.percentile(v, q))
        rows.append([f"第 {q} 百分位", f"{x:.2f}", f"{x / lo:.2f}×"])
    rows.append(["最高", f"{v[-1]:.2f}", f"{v[-1] / lo:.1f}×"])
    print(f"合法的映射（分块 × 两级循环顺序）一共 {n:,} 个")
    print_table(["", "能量 pJ/乘加", "是最低的几倍"], rows)
    print(f"\n在最低能量 1.25 倍以内的：{np.mean(v <= 1.25 * lo):.1%}；2 倍以内：{np.mean(v <= 2 * lo):.1%}"
          "（Interstellar 图 10：AlexNet CONV3 上只有 30% 在 1.25 倍以内）")
    print_observations([
        "同一种数据流、同一块硬件，分块方案之间差一个数量级以上；随手挑一个（中位数）就比最好的差好几倍。"
        "在 1.25 倍以内的方案只占很小一部分（我们这个层比论文的 CONV3 更挑分块）。",
        "所以真正决定能效的是分块（以及循环顺序），不是数据流的名字——这是 Interstellar 最重要的结论之一。"
        "它也说明了为什么加速器离不开一个好的编译器 / 映射搜索（第 12、13 讲）。",
    ])


# -----------------------------------------------------------------------------
def demo3_replication():
    print_header("演示 3：数据流影响的是利用率——K 只有 3（比如卷积的第一层：3 个颜色通道）")
    arch = Arch(rf_bytes=64)
    shape = {"M": 960, "N": 64, "K": 3}
    macs = shape["M"] * shape["N"] * shape["K"]
    rows = []
    for rd, cd, name in ((("K",), ("N",), "K|N：K 放在 16 行上"), (("K", "M"), ("N",), "KM|N：K 和 M 一起放在行上（复制）"),
                         (("M",), ("N",), "M|N：输出驻留")):
        u = utilization(shape, rd, cd, arch)
        _, e, _, _ = search_fast(shape, arch, rd, cd)
        rows.append([name, f"{u:.0%}", f"{macs / (u * 256):,.0f}", f"{e['total'] / macs:.2f}"])
    print_table(["数据流", "PE 利用率", "需要的周期（不计停顿）", "最低能量 pJ/乘加"], rows)
    print_observations([
        "K = 3 放在 16 行上，只有 3 行在干活（利用率 19%），要多花约 5 倍的时间——Interstellar 图 2 的例子。",
        "\"复制\"：在同一个方向上再展开一个维度（这里 M 展开 5 份），3 × 5 = 15 行在干活，利用率回到 94%（论文：3/16 → 15/16）。"
        "这需要阵列的互连足够灵活（Eyeriss v2、CGRA、FPGA 支持）；固定互连的阵列只能挑\"天然够大\"的维度（输出驻留在这里也是满的）。",
        "能量几乎不受影响：利用率决定的是**时间**（吞吐），不是能量。论文图 9：用了复制以后各种数据流的利用率差别也大多消失了。"
        "论文最后选 C|K（这里的 K|N）是因为通道维通常最大、最容易填满阵列。",
    ])


if __name__ == "__main__":
    demo1_dataflows()
    demo2_blocking()
    demo3_replication()