"""
=============================================================================
CS349E · 前置讲 P1 共享模块：arch_lab.py —— 体系结构入门的小工具箱
=============================================================================

【这个模块要解决什么问题】
    前置讲 P1 要把"计算机怎样执行程序、数据怎样在存储层次里流动"讲清楚。光讲概念很难有感觉，
    所以这里提供几个可以动手玩的小模型：

      1. Cache：一个组相联（set-associative）、LRU 替换的 cache 模拟器。输入一串内存地址，
         告诉你每次访问是命中（hit）还是缺失（miss），并把缺失分成三类（3C）。
      2. Hierarchy：把几级 cache 串起来（L1 → L2 → 内存），算平均访存时间（AMAT）。
      3. 地址序列生成器：顺序、跨步、随机、按行/按列遍历矩阵、分块转置……
      4. 几个公式：流水线的周期数、多累加器的指令级并行、Little 定律、DRAM 的行缓冲命中率。
      5. 计时与读取本机 cache 信息的工具（给实测脚本用）。

    这些都是"玩具模型"：真实 CPU 的 cache 还有预取器、多个端口、非包含/包含策略、写缓冲……
    但三条最重要的规律——**局部性决定命中率、命中率决定平均延迟、在途请求数决定带宽**——
    在玩具模型里看得最清楚。

【运行】python3 arch_lab.py    # 自检：几个手算得出结果的小例子
【依赖】Python 标准库 + numpy（地址序列用 numpy 数组表示）。
"""

import os
import platform
import subprocess
import time
from collections import OrderedDict

import numpy as np


# ---------------------------------------------------------------------------
# 1. Cache 模拟器
# ---------------------------------------------------------------------------
class Cache:
    """
    一个组相联、LRU 替换、只读（不模拟写）的 cache。

    【它是怎样组织的】
        cache 按 **cache line**（缓存行，例如 64 字节）为单位存数据：从内存搬数据时一次搬一整行。
        总共 size / line 行，分成 n_sets 组（set），每组 ways 行（"路"）。
        一个地址 addr 属于哪一行：line_id = addr // line；属于哪一组：set_id = line_id % n_sets。
        同一组里最多同时放 ways 行；满了再来新行，就按 LRU（最近最少使用）踢掉一行。
        - ways = 1：直接映射（direct-mapped），每行只有一个位置可放；
        - n_sets = 1：全相联（fully associative），任何行可以放在任何位置。

    【三种缺失（3C）】为了解释"为什么缺失"，我们同时维护一个同样容量的**全相联** LRU cache 做对照：
        - compulsory（强制缺失）：这一行是第一次被访问，任何 cache 都躲不掉；
        - capacity（容量缺失）：全相联 cache 也缺失了——工作集比 cache 大；
        - conflict（冲突缺失）：全相联 cache 会命中，但本 cache 因为"组"太挤而缺失。

    【参数】size：总容量（字节）；line：行大小（字节）；ways：每组几路；name：名字（打印用）；
        policy：替换策略，"lru"（踢掉最久没用的）或 "random"（随机踢一个，用固定种子保证可重复）
    【边界】size 必须能被 line × ways 整除；否则抛出 ValueError。policy 不认识时也抛出 ValueError。
        3C 分类用的对照 cache 始终是全相联 LRU，所以 policy="random" 时的分类只是近似。
    【Python 小知识】OrderedDict 记住元素插入的顺序；move_to_end(k) 把 k 挪到"最新"一端，
        popitem(last=False) 弹出"最旧"的元素——正好实现 LRU。
    """

    def __init__(self, size, line=64, ways=8, name="cache", policy="lru", seed=0):
        if size % (line * ways):
            raise ValueError("size 必须能被 line × ways 整除")
        if policy not in ("lru", "random"):
            raise ValueError("policy 只能是 'lru' 或 'random'")
        self.size, self.line, self.ways, self.name = size, line, ways, name
        self.policy = policy
        self.rng = np.random.default_rng(seed)
        self.n_sets = size // (line * ways)
        self.sets = [OrderedDict() for _ in range(self.n_sets)]
        self.full = OrderedDict()          # 同容量的全相联 LRU，用来区分冲突缺失和容量缺失
        self.full_capacity = size // line
        self.seen = set()                  # 见过的行，用来识别强制缺失
        self.reset_stats()

    def reset_stats(self):
        """把命中、缺失计数清零（cache 内容保留）。"""
        self.hits = 0
        self.misses = {"compulsory": 0, "capacity": 0, "conflict": 0}

    def _touch_full(self, line_id):
        """在对照用的全相联 cache 里访问一次，返回是否命中。"""
        hit = line_id in self.full
        if hit:
            self.full.move_to_end(line_id)
        else:
            self.full[line_id] = True
            if len(self.full) > self.full_capacity:
                self.full.popitem(last=False)
        return hit

    def access(self, addr):
        """
        访问一个字节地址，返回 True（命中）或 False（缺失），并更新统计。

        【做法】算出行号和组号；在组里查找；命中则把它标为最近使用；缺失则装入（必要时踢掉 LRU 行），
            并借助全相联对照判断这是哪一类缺失。
        """
        line_id = int(addr) // self.line
        s = self.sets[line_id % self.n_sets]
        full_hit = self._touch_full(line_id)
        if line_id in s:
            s.move_to_end(line_id)
            self.hits += 1
            return True
        if len(s) >= self.ways:
            if self.policy == "lru":
                s.popitem(last=False)
            else:
                victim = list(s.keys())[int(self.rng.integers(len(s)))]
                del s[victim]
        s[line_id] = True
        if line_id not in self.seen:
            self.misses["compulsory"] += 1
            self.seen.add(line_id)
        elif full_hit:
            self.misses["conflict"] += 1
        else:
            self.misses["capacity"] += 1
        return False

    def run(self, addrs):
        """依次访问一串地址，返回命中率。【边界】空序列返回 0。"""
        for a in addrs:
            self.access(a)
        return self.hit_rate()

    @property
    def n_access(self):
        return self.hits + sum(self.misses.values())

    def hit_rate(self):
        """命中次数 / 访问次数；没有访问时返回 0。"""
        n = self.n_access
        return self.hits / n if n else 0.0

    def describe(self):
        """一行文字描述这个 cache 的结构。"""
        kind = "直接映射" if self.ways == 1 else ("全相联" if self.n_sets == 1 else f"{self.ways} 路组相联")
        return f"{self.name}: {self.size // 1024} KiB, {self.line} B 行, {kind}（{self.n_sets} 组）"


class Hierarchy:
    """
    多级 cache：先查 L1，缺失了查 L2，再缺失去内存。

    【参数】levels：[(Cache, 命中延迟 ns), ...]，从近到远；mem_latency_ns：内存（DRAM）的延迟
    【做法】每次访问从 L1 开始往下找，找到为止；把一路上经过的延迟加起来作为这次访问的延迟。
        （真实硬件里各级的查找可以部分重叠，这里按"串行累加"的简化模型。）
        包含式（inclusive）：在某一级命中后，上面各级也都装入这一行（已经在 access 里装入了）。
    """

    def __init__(self, levels, mem_latency_ns):
        self.levels = levels
        self.mem_latency_ns = mem_latency_ns
        self.total_ns = 0.0
        self.count = 0

    def access(self, addr):
        """访问一个地址，返回这次访问的延迟（ns）。"""
        t = 0.0
        for cache, lat in self.levels:
            t += lat
            if cache.access(addr):
                break
        else:
            t += self.mem_latency_ns
        self.total_ns += t
        self.count += 1
        return t

    def run(self, addrs):
        """依次访问一串地址，返回平均访存时间 AMAT（ns）。"""
        for a in addrs:
            self.access(a)
        return self.amat()

    def amat(self):
        """平均访存时间 = 总延迟 / 访问次数；没有访问时返回 0。"""
        return self.total_ns / self.count if self.count else 0.0


def amat_formula(hit_times, miss_rates, mem_time):
    """
    平均访存时间的公式版本（教科书写法）：
        AMAT = t1 + m1 · (t2 + m2 · (t3 + ... + m_k · t_mem))
    【参数】hit_times：各级命中时间；miss_rates：各级的**局部**缺失率（到达这一级的访问里有多少缺失）
    【例子】L1 1 ns、缺失率 5%，L2 8 ns、局部缺失率 20%，内存 120 ns：
        AMAT = 1 + 0.05 × (8 + 0.2 × 120) = 2.6 ns
    """
    t = mem_time
    for h, m in zip(reversed(hit_times), reversed(miss_rates)):
        t = h + m * t
    return t


# ---------------------------------------------------------------------------
# 2. 地址序列生成器（都返回 numpy 整数数组，单位：字节地址）
# ---------------------------------------------------------------------------
def seq_trace(n_elems, elem_bytes=4, base=0):
    """顺序访问 n_elems 个元素：base, base+4, base+8, ..."""
    return base + np.arange(n_elems, dtype=np.int64) * elem_bytes


def strided_trace(n_elems, stride_elems, elem_bytes=4, base=0):
    """跨步访问：每次跳过 stride_elems 个元素（stride = 1 就是顺序访问）。"""
    return base + np.arange(n_elems, dtype=np.int64) * stride_elems * elem_bytes


def random_trace(n_elems, footprint_bytes, elem_bytes=4, seed=0):
    """在 footprint_bytes 大小的区域里随机访问 n_elems 次（对齐到元素）。"""
    rng = np.random.default_rng(seed)
    return rng.integers(0, footprint_bytes // elem_bytes, n_elems, dtype=np.int64) * elem_bytes


def repeat_scan_trace(footprint_bytes, passes, elem_bytes=4):
    """
    把一个 footprint_bytes 大的数组从头到尾扫 passes 遍（典型的"工作集"实验）。

    【用途】看"工作集比 cache 大一点点"时 LRU 会发生什么（02 脚本演示 2）。
    """
    one = np.arange(footprint_bytes // elem_bytes, dtype=np.int64) * elem_bytes
    return np.tile(one, passes)


def matrix_trace(rows, cols, order="row", elem_bytes=4, base=0, ld=None):
    """
    遍历一个**按行存储**（row-major，C/numpy 的默认）的 rows × cols 矩阵。

    【参数】order："row" 按行遍历（先走完一行再下一行），"col" 按列遍历；
        ld（leading dimension，行距）：相邻两行首地址相隔几个元素，默认等于 cols。
        ld > cols 表示每行末尾留了 ld − cols 个元素的空隙（"填充"，padding）。
    【直觉】按行存储时，A[i][j] 的地址 = base + (i·ld + j)·elem_bytes。按行遍历地址连续；
        按列遍历每一步跳过一整行（ld·elem_bytes 字节）。
    """
    ld = cols if ld is None else ld
    i, j = np.meshgrid(np.arange(rows), np.arange(cols), indexing="ij")
    if order == "col":
        i, j = i.T, j.T
    return base + (i.ravel().astype(np.int64) * ld + j.ravel()) * elem_bytes


def transpose_trace(n, tile=None, elem_bytes=4, src_base=0, dst_base=None, ld=None):
    """
    矩阵转置 B[j][i] = A[i][j] 的访存序列（A、B 都是 n × n、按行存储），交替出现"读 A、写 B"。

    【参数】tile：None 表示朴素的两重循环；否则按 tile × tile 的小块转置（分块，tiling）
    【直觉】朴素转置时，读 A 是按行（连续），写 B 却是按列（每步跳一整行）——总有一边没有空间局部性。
        分块之后，一小块的 A 和对应的一小块 B 都能同时待在 cache 里，两边都能复用 cache line。
    【参数】ld：A、B 的行距（元素个数），默认 n；ld > n 表示每行末尾有填充。
    【边界】dst_base 默认放在 A 后面（不重叠）；tile 必须整除 n。
    """
    ld = n if ld is None else ld
    if dst_base is None:
        dst_base = src_base + n * ld * elem_bytes
    if tile is None:
        i, j = np.meshgrid(np.arange(n), np.arange(n), indexing="ij")
        i, j = i.ravel(), j.ravel()
    else:
        if n % tile:
            raise ValueError("tile 必须整除 n")
        ii, jj, ti, tj = np.meshgrid(np.arange(n // tile), np.arange(n // tile),
                                     np.arange(tile), np.arange(tile), indexing="ij")
        i = (ii * tile + ti).ravel()
        j = (jj * tile + tj).ravel()
    i = i.astype(np.int64)
    j = j.astype(np.int64)
    reads = src_base + (i * ld + j) * elem_bytes
    writes = dst_base + (j * ld + i) * elem_bytes
    out = np.empty(2 * len(reads), dtype=np.int64)
    out[0::2] = reads
    out[1::2] = writes
    return out


# ---------------------------------------------------------------------------
# 3. 几个公式
# ---------------------------------------------------------------------------
def pipeline_cycles(n_instr, stages, stall_cycles=0):
    """
    执行 n_instr 条指令要多少个时钟周期：不分级 vs 分 stages 级流水线。

    【公式】
        不用流水线：每条指令走完全部 stages 步才开始下一条 → n · stages
        理想流水线：第一条指令要 stages 个周期才出来，之后每个周期出一条 → stages + n − 1
        有停顿时再加上停顿的周期数 stall_cycles（数据依赖、分支预测失败等造成的"气泡"）
    【返回】(不用流水线的周期数, 流水线的周期数)
    """
    return n_instr * stages, stages + n_instr - 1 + stall_cycles


def reduction_throughput(n_accumulators, add_latency, adds_per_cycle):
    """
    求和循环 s += x[i] 每个周期能完成几次加法（指令级并行 ILP 的玩具模型）。

    【直觉】s += x[i] 的每一次加法都要等上一次加法的结果（循环携带依赖）。浮点加法有 add_latency
        个周期的延迟，所以只用一个累加器时，每 add_latency 个周期才能完成一次加法。
        用 k 个独立的累加器（s0 += x[0], s1 += x[1], ..., 最后再加起来），k 条依赖链可以交错进行，
        吞吐提高到 k / add_latency，直到碰到硬件每周期最多能发射的加法数 adds_per_cycle。
    【返回】每周期完成的加法数 = min(k / latency, adds_per_cycle)
    """
    return min(n_accumulators / add_latency, adds_per_cycle)


def littles_law_inflight(bandwidth, latency_s):
    """
    Little 定律：系统里平均"在路上"的量 = 到达速率 × 停留时间。

    【放到内存上】要以 bandwidth（byte/s）的速度拿到数据，而每个请求从发出到返回要 latency_s 秒，
        那么任意时刻必须有 bandwidth × latency_s 字节的请求"在途"（已发出、未返回）。
    【例子】H100：3.35 TB/s × 500 ns ≈ 1.7 MB 的请求同时在途——这就是 GPU 需要成千上万个线程的原因。
    【返回】在途字节数
    """
    return bandwidth * latency_s


def dram_row_hit_rate(addrs, row_bytes=8192, n_banks=16):
    """
    一个极简的 DRAM 行缓冲（row buffer）模型：每个 bank 记住当前"打开"的那一行。

    【背景】DRAM 读数据时先把一整行（几 KB）读到这个 bank 的行缓冲里，再从行缓冲里取需要的那几个字节。
        如果下一次访问恰好在同一行（row hit），只要从行缓冲取数，很快；
        如果是另一行（row miss），要先把旧行写回、再打开新行，慢好几倍。
    【参数】addrs：地址序列；row_bytes：每行字节数；n_banks：bank 数（地址按行轮流分到各 bank）
    【返回】行命中率
    """
    rows = np.asarray(addrs, dtype=np.int64) // row_bytes
    banks = rows % n_banks
    open_row = {}
    hits = 0
    for r, b in zip(rows.tolist(), banks.tolist()):
        if open_row.get(b) == r:
            hits += 1
        open_row[b] = r
    return hits / len(rows) if len(rows) else 0.0


# ---------------------------------------------------------------------------
# 4. 计时与本机信息
# ---------------------------------------------------------------------------
def best_time(fn, repeats=5, warmup=1):
    """
    多次运行 fn()，返回最短的一次耗时（秒）。

    【为什么取最短】测量噪声（别的进程、缓存被冲掉、降频）几乎只会让程序变慢，最短时间最接近机器本来的能力。
    【Python 小知识】time.perf_counter() 是高精度计时器，只能做差，绝对值没有意义。
    """
    for _ in range(warmup):
        fn()
    best = float("inf")
    for _ in range(repeats):
        t0 = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t0)
    return best


def cpu_cache_info():
    """
    尽量读出本机的 cache 大小。返回字典（读不到就返回空字典）。

    【做法】macOS 用 sysctl 命令（hw.perflevel0.* 是性能核的数字）；Linux 读
        /sys/devices/system/cpu/cpu0/cache/index*/ 下的文件。其他系统返回空字典。
    【边界】任何一步失败都静默跳过：这只是"打印给你看"的辅助信息，不影响实验。
    【Python 小知识】subprocess.run([...], capture_output=True, text=True) 运行一个外部命令并拿到它的输出。
    """
    info = {}
    if platform.system() == "Darwin":
        keys = {"hw.perflevel0.l1dcachesize": "L1D（性能核）", "hw.perflevel0.l2cachesize": "L2（性能核簇共享）",
                "hw.perflevel0.cpusperl2": "共享一个 L2 的核数", "hw.cachelinesize": "cache line",
                "hw.perflevel0.physicalcpu": "性能核数", "hw.perflevel1.physicalcpu": "能效核数"}
        for k, label in keys.items():
            try:
                r = subprocess.run(["sysctl", "-n", k], capture_output=True, text=True, timeout=5)
                if r.returncode == 0 and r.stdout.strip():
                    info[label] = int(r.stdout.strip())
            except (OSError, ValueError, subprocess.SubprocessError):
                pass
    elif platform.system() == "Linux":
        base = "/sys/devices/system/cpu/cpu0/cache"
        try:
            for d in sorted(os.listdir(base)):
                p = os.path.join(base, d)
                with open(os.path.join(p, "level")) as f:
                    level = f.read().strip()
                with open(os.path.join(p, "type")) as f:
                    typ = f.read().strip()
                with open(os.path.join(p, "size")) as f:
                    size = f.read().strip()
                info[f"L{level} {typ}"] = size
        except OSError:
            pass
    return info


def fmt_size(nbytes):
    """把字节数写成 KiB / MiB（cache 大小习惯用二进制单位）。"""
    if nbytes >= 1 << 20:
        return f"{nbytes / (1 << 20):g} MiB"
    if nbytes >= 1 << 10:
        return f"{nbytes / (1 << 10):g} KiB"
    return f"{nbytes} B"


# ---------------------------------------------------------------------------
# 自检
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # 1) 顺序读 4 字节元素、64 字节行：每 16 次访问只有第 1 次缺失 → 命中率 15/16
    c = Cache(32 * 1024, 64, 8)
    print(c.describe(), "顺序访问命中率", round(c.run(seq_trace(4096)), 4), "（应为 0.9375）")
    # 2) 跨步 = 行大小：每次都是新行 → 命中率 0
    c = Cache(32 * 1024, 64, 8)
    print("跨 16 个元素（64 字节）访问命中率", c.run(strided_trace(4096, 16)), "（应为 0）")
    # 3) AMAT 公式
    print("AMAT 公式例子", amat_formula([1, 8], [0.05, 0.2], 120), "ns（应为 2.6）")
    # 4) 流水线
    print("1000 条指令、5 级流水线：", pipeline_cycles(1000, 5), "（应为 (5000, 1004)）")
    # 5) Little 定律
    print("H100 在途字节", littles_law_inflight(3.35e12, 500e-9) / 1e6, "MB")
    print("本机 cache 信息：", cpu_cache_info())