"""
=============================================================================
CS349E · 前置讲 P1 · 02_cache_simulator.py —— 用模拟器看懂 cache：局部性、3C 缺失、分块
=============================================================================

【这个脚本要回答的问题】
    处理器访问 DRAM 要一百多纳秒（03 脚本实测），访问 L1 cache 不到 1 纳秒。程序快不快，
    很大程度上取决于"要用的数据是不是恰好在 cache 里"。本脚本用 arch_lab.Cache 模拟器，
    一个现象一个演示：

      演示 1：空间局部性——顺序访问 vs 跨步访问 vs 随机访问；cache line 的大小。
      演示 2：时间局部性——工作集比 cache 小一点和大一点，命中率天差地别；LRU 的"悬崖"。
      演示 3：冲突缺失——2 的幂次方的跨度会让所有数据挤进同一组；提高相联度或"填充"（padding）来解决。
      演示 4：矩阵——按行 vs 按列遍历；朴素转置 vs 分块转置（tiling）。
      演示 5：平均访存时间（AMAT）——把命中率换算成"每次访问平均要等多久"。

【怎么读输出】全部是模拟，每次运行完全相同。模拟的 cache 参数是教学用的典型值（32 KiB、64 字节行、8 路），
    不对应某一款具体 CPU。

【运行】python3 02_cache_simulator.py        （约 2–3 秒：模拟器用 Python 逐个处理地址）
        python3 02_cache_simulator.py --quick （约 1 秒，规模更小、跳过 1024×1024 的例子）
【依赖】numpy；从 ../lecture01 导入 console。
"""

import os
import sys

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

from console import print_header, print_observations, print_table  # noqa: E402
from arch_lab import (Cache, Hierarchy, amat_formula, seq_trace, strided_trace, random_trace,  # noqa: E402
                      repeat_scan_trace, matrix_trace, transpose_trace, fmt_size)

QUICK = "--quick" in sys.argv
KiB = 1024


def miss_cells(c):
    """把一个 Cache 的统计整理成表格的几列：命中率、三类缺失各多少次。"""
    m = c.misses
    return [f"{c.hit_rate() * 100:.1f}%", m["compulsory"], m["capacity"], m["conflict"]]


# ---------------------------------------------------------------------------
# 演示 1：空间局部性
# ---------------------------------------------------------------------------
def demo1_spatial():
    """
    同样访问 N 次 4 字节的元素，只改变访问的"跨度"（stride）。

    【直觉】cache 每次从内存搬一整行（64 字节 = 16 个 float32）。顺序访问时，第 1 个元素缺失、
        把整行搬进来，后面 15 个都命中——这叫**空间局部性**：用了一个数，很快就会用它旁边的数。
        跨度越大，一行里被用到的数越少；跨度 ≥ 一行时，每次访问都是新的一行，命中率 0。
    """
    print_header("演示 1：空间局部性——同样访问 4096 次，只改变跨度（32 KiB、64 B 行、8 路）")
    n = 4096
    rows = []
    for stride in (1, 2, 4, 8, 16, 64):
        c = Cache(32 * KiB, 64, 8)
        c.run(strided_trace(n, stride))
        per_line = 64 // min(64, 4 * stride)      # 一行 64 字节里有几个元素会被访问到
        rows.append([f"跨 {stride} 个元素（{4 * stride} B）", per_line] + miss_cells(c))
    c = Cache(32 * KiB, 64, 8)
    c.run(random_trace(n, 4 * 1024 * KiB))
    rows.append(["4 MiB 范围内随机", "≈1"] + miss_cells(c))
    print_table(["访问方式", "每行被用到几个数", "命中率", "强制缺失", "容量缺失", "冲突缺失"], rows)

    print("\ncache line 的大小（顺序访问 4096 个 float32）：")
    rows = []
    for line in (16, 32, 64, 128, 256):
        c = Cache(32 * KiB, line, 8)
        c.run(seq_trace(n))
        rows.append([f"{line} B", f"{c.hit_rate() * 100:.2f}%", f"{1 - 4 / line:.4f}"])
    print_table(["行大小", "命中率", "理论值 1 − 4/行大小"], rows)
    print_observations([
        "顺序访问的命中率 = 1 − 元素大小 / 行大小：每搬一行，后面的元素都免费。行越大，顺序访问越占便宜。",
        "跨度达到一整行（16 个 float32 = 64 B）以后，每次访问都是一行新数据，命中率掉到 0：搬进来 64 字节只用了 4 字节，"
        "带宽浪费了 15/16。GPU 上同样的现象叫'没有合并的访存'（P2 §5）。",
        "随机访问几乎没有空间局部性，范围又远大于 cache，命中率接近 0——这正是 decode 时'每个序列读自己的 KV cache'、"
        "推荐系统查大嵌入表这类负载难的原因。",
        "这里所有的缺失都是'强制缺失'：每一行只被访问一轮，cache 再大也躲不掉。下一个演示看反复访问的情况。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：时间局部性与工作集
# ---------------------------------------------------------------------------
def demo2_working_set(passes=4):
    """
    把大小为 W 的数组从头到尾扫 passes 遍，每行只访问一次（元素大小取 64 B，排除空间局部性的干扰）。

    【理想情况】W ≤ cache 时，第一遍全是强制缺失，之后各遍全命中：命中率 = (passes − 1) / passes。
    【LRU 的悬崖】W 只比 cache 大一点点时，LRU 每次都恰好踢掉"马上就要再用"的那一行——
        顺序循环访问正好是 LRU 最怕的模式，命中率直接掉到 0。随机替换反而能保住一部分。
    """
    print_header(f"演示 2：时间局部性——把 W 大小的数组扫 {passes} 遍（32 KiB 全相联 cache，每行访问一次）")
    rows = []
    for w_kib in (8, 16, 24, 32, 33, 36, 48, 64):
        trace = repeat_scan_trace(w_kib * KiB, passes, elem_bytes=64)
        lru = Cache(32 * KiB, 64, 512, policy="lru")          # 512 路 × 64 B = 32 KiB → 全相联
        rnd = Cache(32 * KiB, 64, 512, policy="random", seed=0)
        lru.run(trace)
        rnd.run(trace)
        rows.append([fmt_size(w_kib * KiB), f"{w_kib / 32:.2f}", f"{lru.hit_rate() * 100:.1f}%",
                     lru.misses["capacity"], f"{rnd.hit_rate() * 100:.1f}%"])
    print_table(["工作集 W", "W / cache", "LRU 命中率", "LRU 容量缺失", "随机替换命中率"], rows)
    print_observations([
        f"工作集放得下时（W ≤ 32 KiB），除了第一遍的强制缺失全部命中：{(passes - 1) / passes * 100:.0f}%。这叫**时间局部性**：用过的数据很快又被用。",
        "工作集只比 cache 大 3%（33 KiB），LRU 命中率就直接掉到 0：每次要用的那一行恰好刚被踢掉。性能不是'慢慢变差'，而是'掉下悬崖'。",
        "随机替换在这里反而更好：它不会系统性地踢掉下一个要用的行。没有一种替换策略在所有访问模式下都最好——"
        "第 7 讲 KV cache 的前缀缓存会再遇到'该踢掉谁'的问题。",
        "对程序员的启示：让工作集**放得进** cache，比任何替换策略都重要。这就是'分块'（tiling）的意义：演示 4。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：冲突缺失
# ---------------------------------------------------------------------------
def demo3_conflict(n_arrays=12, n_iter=512):
    """
    同时顺序读 n_arrays 个数组的第 i 个元素（比如 n 个向量逐元素相加），数组首地址相隔 32 KiB。

    【问题】组号 = (地址 // 行大小) % 组数。32 KiB 正好是"组数 × 行大小"的整数倍时，所有数组的
        第 i 个元素都落在**同一组**。组里只有 ways 个位置，数组个数超过 ways 就会互相踢——即使整个
        cache 远没有用满。这就是**冲突缺失**。
    【两种办法】提高相联度（组里位置更多）；或者在数组之间**填充**（padding）一点空隙，让首地址错开一行。
        P2 §6 的共享内存 bank conflict 用的是同一个填充技巧。
    """
    print_header(f"演示 3：冲突缺失——同时读 {n_arrays} 个首地址相隔 32 KiB 的数组（32 KiB cache）")
    rows = []
    for gap, label in ((32 * KiB, "相隔 32 KiB"), (32 * KiB + 64, "相隔 32 KiB + 64 B（填充一行）")):
        trace = []
        for i in range(n_iter):
            for k in range(n_arrays):
                trace.append(k * gap + i * 4)
        for ways in (1, 4, 8, 16, 512):
            c = Cache(32 * KiB, 64, ways)
            c.run(trace)
            kind = "直接映射" if ways == 1 else ("全相联" if ways == 512 else f"{ways} 路")
            rows.append([label, kind] + miss_cells(c))
    print_table(["数组布局", "相联度", "命中率", "强制缺失", "容量缺失", "冲突缺失"], rows)
    print_observations([
        f"{n_arrays} 个数组一共只需要 {n_arrays} 行就能顺序读（{n_arrays * 64} 字节，远小于 32 KiB），"
        "但首地址相隔 32 KiB 时它们全部挤在同一组：直接映射、4 路、8 路都放不下，几乎每次都是冲突缺失。",
        f"相联度 ≥ {n_arrays}（16 路、全相联）时冲突消失，命中率恢复到顺序访问应有的 15/16。",
        "更便宜的办法是填充：每个数组错开 64 字节，它们就分散到不同的组，8 路甚至直接映射都够用了。",
        "2 的幂次方大小的数组和矩阵在 AI 里到处都是（4096、8192……），冲突问题在 CPU cache、GPU 共享内存的 bank、"
        "甚至 HBM 的通道上都会出现，填充是最常用的解法。",
    ])


# ---------------------------------------------------------------------------
# 演示 4：矩阵遍历与分块转置
# ---------------------------------------------------------------------------
def demo4_matrix(n=256):
    """
    n × n 的 float32 矩阵（n = 256 时 256 KiB）按行存储，cache 只有 32 KiB。

    【按行 vs 按列】按行遍历地址连续，享受空间局部性；按列遍历每一步跳过一整行（n × 4 字节），
        每个元素都在不同的 cache line 上。如果这一列的所有行都能留在 cache 里，下一列（同一批行）
        就能全部命中；留不住，就每次都缺失。
    【为什么留不住】n = 256 时一列只涉及 256 行 × 64 B = 16 KiB，比 32 KiB 小——容量其实够。
        问题在于行距 1 KiB 是 2 的幂：一列的 256 个元素只落在少数几组里（和演示 3 一样），互相冲突。
        所以我们同时看"行距填充 16 个元素（64 B）"的版本：ld = n + 16。
    【转置】B[j][i] = A[i][j]：读 A 按行、写 B 按列，朴素写法总有一边是"按列"。
        分块：每次只处理 t × t 的一小块，A 的 t 行和 B 的 t 行一共 2t 个 cache line，两边都能把一行用完再走。
    """
    print_header(f"演示 4：{n}×{n} float32 矩阵（{n * n * 4 // KiB} KiB）在 32 KiB、8 路 cache 上——遍历顺序、填充与分块")
    rows = []
    for ld, ld_label in ((n, "无填充"), (n + 16, "行距 +64 B")):
        for order, label in (("row", "按行求和"), ("col", "按列求和")):
            c = Cache(32 * KiB, 64, 8)
            c.run(matrix_trace(n, n, order, ld=ld))
            rows.append([label, ld_label] + miss_cells(c))
        for tile in (None, 8, 16, 32, 64):
            c = Cache(32 * KiB, 64, 8)
            c.run(transpose_trace(n, tile, ld=ld))
            label = "朴素转置" if tile is None else f"分块转置 {tile}×{tile}"
            rows.append([label, ld_label] + miss_cells(c))
    if not QUICK:
        big = 1024                      # 4 MiB 的矩阵：一列就要 1024 行 × 64 B = 64 KiB，超过 cache
        c = Cache(32 * KiB, 64, 8)
        c.run(matrix_trace(big, big, "col", ld=big + 16))
        rows.append([f"按列求和（{big}×{big}）", "行距 +64 B"] + miss_cells(c))
        for tile in (None, 32):
            c = Cache(32 * KiB, 64, 8)
            c.run(transpose_trace(big, tile, ld=big + 16))
            label = "朴素转置" if tile is None else f"分块转置 {tile}×{tile}"
            rows.append([f"{label}（{big}×{big}）", "行距 +64 B"] + miss_cells(c))
    print_table(["操作", "行距", "命中率", "强制缺失", "容量缺失", "冲突缺失"], rows)
    print_observations([
        "同一个矩阵、同样的运算，按列遍历的命中率可以接近 0，按行遍历是 15/16。numpy、C 的数组默认按行存储，"
        "所以'最内层循环沿着最后一个下标走'是最基本的性能常识。",
        f"一个意外：{n}×{n} 时按列求和和朴素转置的缺失几乎全是**冲突缺失**，不是容量缺失。一列只要 16 KiB，"
        "容量够；是 1 KiB 这个 2 的幂次方的行距让它们挤进了少数几组。给每行填充 64 字节之后，按列求和就回到了 15/16。",
        "不填充时，分块也能绕开冲突：16×16 的块一次只碰 32 行，命中率回到 15/16；块太大（64×64）又重新撞上冲突。",
        "但填充救不了真正的容量问题：1024×1024 的矩阵一列要 1024 行 × 64 B = 64 KiB，比 cache 大。填充之后，"
        "按列求和、朴素转置仍然几乎全是容量缺失；32×32 分块转置一次只碰 64 行（4 KiB），命中率回到 15/16。",
        "'把大问题切成放得进快速存储的小块'——这就是分块（tiling）。第 1 讲 §3.9 的分块矩阵乘法、第 10 讲的 GPU 内核、"
        "第 11 讲的 FlashAttention 都是它的变体。",
    ])


# ---------------------------------------------------------------------------
# 演示 5：平均访存时间
# ---------------------------------------------------------------------------
def demo5_amat(n=20_000 if QUICK else 60_000):
    """
    两级 cache + 内存：L1 32 KiB（1 ns）、L2 1 MiB（8 ns）、内存 120 ns。每次访问的延迟 = 一路上经过的各级延迟之和。

    【延迟的取值】大致参考 03 脚本在 Apple M4 Pro 上实测的数量级（L1 不到 1 ns、L2 约 7–10 ns、内存约 120 ns），
        取整以便计算；这里的 L1、L2 容量是教学用的典型值，不是那台机器的。
    """
    print_header("演示 5：平均访存时间 AMAT（L1 32 KiB / 1 ns → L2 1 MiB / 8 ns → 内存 120 ns）")
    patterns = [
        ("顺序访问", seq_trace(n)),
        ("随机，范围 16 KiB（放得进 L1）", random_trace(n, 16 * KiB, seed=1)),
        ("随机，范围 512 KiB（放得进 L2）", random_trace(n, 512 * KiB, seed=2)),
        ("随机，范围 64 MiB（都放不下）", random_trace(n, 64 * 1024 * KiB, seed=3)),
    ]
    rows = []
    for name, trace in patterns:
        l1 = Cache(32 * KiB, 64, 8, "L1")
        l2 = Cache(1024 * KiB, 64, 16, "L2")
        h = Hierarchy([(l1, 1.0), (l2, 8.0)], mem_latency_ns=120.0)
        amat = h.run(trace)
        m1 = 1 - l1.hit_rate()
        m2 = 1 - l2.hit_rate() if l2.n_access else 0.0
        rows.append([name, f"{m1 * 100:.1f}%", f"{m2 * 100:.1f}%", f"{amat:.2f} ns",
                     f"{amat_formula([1.0, 8.0], [m1, m2], 120.0):.2f} ns"])
    print_table(["访问模式", "L1 缺失率", "L2 局部缺失率", "AMAT（模拟）", "AMAT（公式）"], rows)
    print_observations([
        "AMAT = L1 时间 + L1 缺失率 × (L2 时间 + L2 缺失率 × 内存时间)。模拟和公式完全一致（模拟就是在逐次累加这个式子）。",
        "随机访问一大片内存时，AMAT 接近内存延迟本身，是 L1 的上百倍；工作集放得进 L2 时，L2 接住了大部分 L1 缺失，AMAT 降到几十分之一。",
        "顺序访问在这个模型里也要 9 ns：每 16 次访问有 1 次要一路走到内存（1 + 8 + 120 ns），被平均到 16 次上。"
        "这个模型没有'预取'。真实 CPU 的硬件预取器看到顺序访问会提前把后面的行搬进来，顺序访问的实际代价低得多——"
        "03 脚本实测：顺序拷贝的带宽在 DRAM 上只比在 L1 上低约 30%，而随机访问的延迟差上百倍。",
        "预取能成功，是因为顺序访问的'下一个地址'可以预测，而且很多请求可以同时在路上——这是 P1 §4（Little 定律）的主题。"
        "随机访问没法预测，延迟就全暴露出来。",
    ])


if __name__ == "__main__":
    demo1_spatial()
    demo2_working_set()
    demo3_conflict()
    demo4_matrix(128 if QUICK else 256)
    demo5_amat()
