"""
=============================================================================
CS349E · 前置讲 P1 · 03_measure_memory_hierarchy.py —— 在你的电脑上"看见"存储层次
=============================================================================

【这个脚本要回答的问题】
    02 脚本的 cache 是模拟的。真实的机器上，L1、L2、内存各有多快？本脚本实测三条曲线，横轴都是
    "工作集大小"（一次反复访问的数据有多大）：

      演示 1：读出本机的 cache 大小（操作系统告诉我们的数字）。
      演示 2：**延迟**——指针追逐（pointer chasing）：每次读到的值是下一次要读的地址，前一次没读完，
              后一次就没法开始。这样测出来的就是"一次访问到底要等多久"。用一个 30 行的 C 小程序来测
              （Python 每一步的解释开销约 20 ns，会盖住 L1、L2 的差别）。需要 C 编译器（cc）；没有就跳过。
      演示 3：**随机访问的吞吐**——numpy 的 a[idx]：一次给出几百万个随机下标，访问之间没有依赖，
              处理器可以同时发出很多个请求。和演示 2 对比，就能算出"同时有多少个请求在路上"。
      演示 4：**顺序访问的带宽**——反复复制一个大小为 W 的数组。

【怎么读输出】全是墙钟测量：数字取决于你的 CPU、内存和当时的负载，每次运行都会略有不同。
    只看形状：延迟曲线是一级一级的台阶；随机访问的吞吐比延迟好得多；顺序带宽几乎是平的。
    编写时用的是 Apple M4 Pro（性能核 L1D 128 KiB、L2 16 MiB）。

【运行】python3 03_measure_memory_hierarchy.py          （约 10 秒）
        python3 03_measure_memory_hierarchy.py --quick  （约 3 秒，只测到 16 MiB）
【依赖】numpy；演示 2 需要系统里有 C 编译器 cc（macOS 装了 Xcode 命令行工具即可；Linux 一般有 gcc/cc）。
"""

import os
import shutil
import subprocess
import sys
import tempfile
import time

import numpy as np

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 best_time, cpu_cache_info, fmt_size  # noqa: E402

QUICK = "--quick" in sys.argv
KiB, MiB = 1024, 1024 * 1024
SIZES = ([16 * KiB, 64 * KiB, 256 * KiB, 1 * MiB, 4 * MiB, 16 * MiB] if QUICK else
         [16 * KiB, 32 * KiB, 64 * KiB, 128 * KiB, 256 * KiB, 512 * KiB, 1 * MiB, 2 * MiB, 4 * MiB,
          8 * MiB, 16 * MiB, 32 * MiB, 64 * MiB, 128 * MiB, 256 * MiB])

# ---------------------------------------------------------------------------
# 指针追逐的 C 程序
# ---------------------------------------------------------------------------
# 【思路】在一块 bytes 大小的内存里，每 128 字节放一个"指针"，指针连成一个**随机顺序**的环。
#   然后从环上一点出发，反复执行 p = *p（读出的值就是下一个地址）。每一步都必须等上一步的读完成，
#   所以"总时间 / 步数"就是一次访问的平均延迟。
#   - 随机顺序：避免硬件预取器猜中下一个地址（顺序环会被预取，测出来就不是真实延迟了）。
#   - 间隔 128 字节：每次访问落在不同的 cache line 上（Apple M 系列的 cache line 是 128 字节，大多数 x86 是 64 字节）。
#   - 先完整走一圈再计时：把数据"预热"进 cache（如果放得下的话）。
# 【C 小知识】void* 是"指向任意东西的指针"；*(void**)q 把地址 q 处存的 8 个字节当作一个指针读出来。
CHASE_C = r"""
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <time.h>
static double now(void){struct timespec t; clock_gettime(CLOCK_MONOTONIC,&t); return t.tv_sec+t.tv_nsec*1e-9;}
int main(int argc,char**argv){
  size_t bytes=strtoull(argv[1],0,10), steps=strtoull(argv[2],0,10), stride=128;
  size_t n=bytes/stride; if(n<2) n=2;
  char*buf=malloc(n*stride+4096);
  size_t*perm=malloc(n*sizeof(size_t));
  for(size_t i=0;i<n;i++) perm[i]=i;
  uint64_t s=12345;                                   /* 固定种子的线性同余随机数，结果可重复 */
  for(size_t i=n-1;i>0;i--){ s=s*6364136223846793005ULL+1442695040888963407ULL;
                             size_t j=(s>>33)%(i+1); size_t t=perm[i]; perm[i]=perm[j]; perm[j]=t; }
  for(size_t i=0;i<n;i++) *(void**)(buf+perm[i]*stride)=buf+perm[(i+1)%n]*stride;
  void*p=buf+perm[0]*stride;
  for(size_t i=0;i<n;i++) p=*(void**)p;               /* 预热：先走一圈 */
  double t0=now();
  for(size_t i=0;i<steps;i++) p=*(void**)p;
  double t=now()-t0;
  printf("%.4f %p\n", t/steps*1e9, p);                /* 打印 p，防止编译器把整个循环优化掉 */
  return 0;
}
"""


def build_chaser(workdir):
    """
    把 CHASE_C 写进临时目录并编译，返回可执行文件路径；没有编译器或编译失败返回 None。

    【Python 小知识】shutil.which("cc") 在 PATH 里找名为 cc 的程序，找不到返回 None。
        subprocess.run(..., check=True) 在命令返回非 0 时抛出 CalledProcessError。
    """
    cc = shutil.which("cc") or shutil.which("gcc") or shutil.which("clang")
    if cc is None:
        return None
    src = os.path.join(workdir, "chase.c")
    exe = os.path.join(workdir, "chase")
    with open(src, "w") as f:
        f.write(CHASE_C)
    try:
        subprocess.run([cc, "-O2", "-o", exe, src], check=True, capture_output=True, timeout=60)
    except (subprocess.SubprocessError, OSError):
        return None
    return exe


def chase_latency(exe, nbytes, target_s=0.15):
    """
    测一个工作集大小下的平均访问延迟（ns）。

    【做法】先跑 20 万步估计延迟，再按"总共跑约 target_s 秒"决定正式的步数，跑 3 次取最小值。
    【边界】输出格式不对时抛出 ValueError（不应该发生）。
    """
    def run(steps):
        out = subprocess.run([exe, str(nbytes), str(steps)], check=True, capture_output=True, text=True, timeout=120)
        return float(out.stdout.split()[0])
    est = run(200_000)
    steps = int(min(5e7, max(2e5, target_s / (est * 1e-9))))
    return min(run(steps) for _ in range(3))


# ---------------------------------------------------------------------------
# 演示 1：本机 cache 信息
# ---------------------------------------------------------------------------
def demo1_info():
    print_header("演示 1：操作系统报告的 cache 信息")
    info = cpu_cache_info()
    if not info:
        print("  （这个系统上读不到 cache 信息，跳过；不影响后面的测量）")
        return info
    rows = []
    for k, v in info.items():
        rows.append([k, fmt_size(v) if isinstance(v, int) and v >= 1024 else v])
    print_table(["项目", "数值"], rows)
    return info


# ---------------------------------------------------------------------------
# 演示 2：延迟（指针追逐）
# ---------------------------------------------------------------------------
def demo2_latency():
    print_header("演示 2：延迟——指针追逐（每一步都要等上一步读完）")
    with tempfile.TemporaryDirectory() as d:
        exe = build_chaser(d)
        if exe is None:
            print("  没有找到可用的 C 编译器（cc / gcc / clang），跳过这个演示。")
            return {}
        lat = {}
        for w in SIZES:
            lat[w] = chase_latency(exe, w)
    rows = [[fmt_size(w), f"{ns:.2f} ns", f"{ns / lat[SIZES[0]]:.1f}×"] for w, ns in lat.items()]
    print_table(["工作集", "每次访问的延迟", "相对最小工作集"], rows)
    small, large = lat[SIZES[0]], lat[SIZES[-1]]
    print_observations([
        f"曲线是一级一级的台阶：工作集放得进 L1 时每次约 {small:.2f} ns（只有几个时钟周期），"
        f"放不进任何一级 cache 时约 {large:.0f} ns——慢了约 {large / small:.0f} 倍。",
        "每个台阶的拐点大致就是某一级 cache 的容量（对照演示 1 的数字）。拐点不是一个尖角，"
        "因为 cache 是组相联的、还有别的程序在用 cache，工作集接近容量时会逐渐出现缺失（02 脚本演示 2、3）。",
        "最大的几个工作集还会碰到 TLB（页表缓存）缺失：随机地跨越几百 MB，连'虚拟地址 → 物理地址'的翻译也要去内存里查，"
        "所以最右边的延迟还会继续缓慢上升。",
        "这就是第 1 讲 Horowitz 能量表的'时间版'：离计算单元越远，越慢（也越费电）。",
    ])
    return lat


# ---------------------------------------------------------------------------
# 演示 3：随机访问的吞吐 vs 延迟
# ---------------------------------------------------------------------------
def demo3_random_gather(lat):
    """
    numpy 的 a[idx]：一次读 m 个随机位置。访问之间没有依赖，CPU 可以同时发出多个缺失请求。

    【在途请求数的估计】如果每次访问要等 L 纳秒（演示 2 测的延迟），而平均每 t 纳秒就完成一次访问，
        那么平均同时有 L / t 个访问"在路上"——这正是 Little 定律（P1 §4）。
    【边界】演示 2 被跳过时（没有编译器），只打印吞吐，不估计在途请求数。
    """
    print_header("演示 3：随机访问的吞吐——访问之间没有依赖时会怎样？")
    rng = np.random.default_rng(0)
    m = 1 << (20 if QUICK else 22)
    rows = []
    mlp_large = None
    for w in SIZES:
        n = w // 4
        a = np.ones(n, dtype=np.float32)
        idx = rng.integers(0, n, m)
        t = best_time(lambda: a[idx], repeats=3 if QUICK else 5) / m
        cells = [fmt_size(w), f"{t * 1e9:.2f} ns"]
        if lat:
            mlp = lat[w] / (t * 1e9)
            cells += [f"{lat[w]:.2f} ns", f"{mlp:.1f}"]
            mlp_large = mlp
        rows.append(cells)
    headers = ["工作集", "每次随机访问（吞吐）"] + (["每次访问的延迟（演示 2）", "≈ 同时在途的访问数"] if lat else [])
    print_table(headers, rows)
    obs = ["随机访问的'吞吐时间'也随工作集变大而变慢，但比演示 2 的延迟小得多："
           "处理器不需要等上一个访问完成，就可以发出下一个。"]
    if mlp_large is not None:
        obs.append(f"工作集最大时，延迟约是吞吐时间的 {mlp_large:.0f} 倍，也就是平均约有 {mlp_large:.0f} 个访问同时在路上。"
                   "这叫**内存级并行**（memory-level parallelism, MLP）：延迟没变，但很多请求重叠起来，吞吐就上去了。")
    obs.append("所以'慢'有两种：一种是**依赖**造成的（下一步要等上一步，像 LLM 的 decode），只能靠缩短延迟；"
               "另一种只是**单个请求慢**，可以靠同时发很多请求来掩盖。GPU 把第二种做到了极致（P2 §4、§7）。")
    obs.append("小工作集时每次访问不到 1 ns，主要是 numpy 本身的开销（生成结果数组、处理下标），已经不是 cache 的速度了。")
    print_observations(obs)


# ---------------------------------------------------------------------------
# 演示 4：顺序带宽
# ---------------------------------------------------------------------------
def demo4_sequential():
    """
    反复 np.copyto(b, a)，a、b 各 W/2 字节（总工作集 W），按"读 + 写 = 2 × 数组大小"记流量。

    【为什么几乎是平的】顺序访问的下一个地址完全可预测：硬件预取器会提前很多行把数据搬进来，
        许多请求同时在路上，延迟被藏起来了。剩下的限制只是"每秒能搬多少字节"。
    """
    print_header("演示 4：顺序访问的带宽——反复复制一个数组（单线程）")
    total = (1 << 26) if QUICK else (1 << 28)
    rows = []
    bw = {}
    for w in SIZES:
        n = max(1, w // 8)                    # 两个 float32 数组，各 w/2 字节
        a = np.ones(n, dtype=np.float32)
        b = np.empty_like(a)
        rep = max(1, total // (2 * a.nbytes))

        def run():
            for _ in range(rep):
                np.copyto(b, a)
        t = best_time(run, repeats=3)
        bw[w] = 2 * a.nbytes * rep / t
        rows.append([fmt_size(w), f"{bw[w] / 1e9:.1f} GB/s"])
    print_table(["工作集（两个数组合计）", "复制带宽"], rows)
    peak = max(bw.values())
    print_observations([
        f"顺序复制的带宽从最高约 {peak / 1e9:.0f} GB/s 降到工作集最大时的约 {bw[SIZES[-1]] / 1e9:.0f} GB/s，"
        f"只差 {peak / bw[SIZES[-1]]:.1f} 倍——而演示 2 的延迟差了上百倍。",
        "差别来自可预测性：顺序访问被预取器'提前搬好'，延迟藏在了带宽后面；随机且有依赖的访问没法提前，延迟全部暴露。",
        "最小的工作集反而偏慢：每次 np.copyto 调用有约零点几微秒的固定开销，数组太小时这部分占了大头"
        "（和 01 脚本'解释器开销'是同一件事）。",
        "这一组对比是整门课的缩影：**能变成大块、顺序、可并行的访问，就受带宽限制；变不成的，就受延迟限制**。"
        "GPU、HBM、分块、批处理，都是在把问题往前一种情况推。",
    ])


if __name__ == "__main__":
    t0 = time.perf_counter()
    demo1_info()
    lat = demo2_latency()
    demo3_random_gather(lat)
    demo4_sequential()
    print(f"\n⏱ 总耗时 {time.perf_counter() - t0:.1f} 秒")