"""
=============================================================================
CS349E · 第 9 讲共享模块：loopnest.py —— 把加速器写成循环嵌套：分块、数据流与能量模型（Interstellar 的思路）
=============================================================================

【这个模块要解决什么问题】
    第 8 讲看到 TPU 用"权重驻留"的脉动阵列。还有输出驻留、行驻留……各种加速器论文各说各好。
    Interstellar（Yang et al., ASPLOS 2020）的出发点：所有这些加速器算的是**同一个循环嵌套**，
    区别只在三件事——
        (1) 分块（loop blocking）：每个循环切成几段，每一段对应存储层次的一级（寄存器 RF、片上缓冲 GB、DRAM）；
        (2) 循环顺序（reorder）：同一级里哪个循环在外、哪个在内——决定哪个张量"驻留"不动；
        (3) 空间展开（spatial unrolling）：哪些循环在 PE 阵列上并行展开——这就是"数据流"（dataflow）。
    再加上 (4) 资源分配：PE 有多少、每一级存储多大。
    给定这些，每一级存储被访问多少次就能算出来，乘以每次访问的能量（Interstellar 表 3），就得到总能量。

    本模块对**矩阵乘法** C[M, N] += A[M, K] · B[K, N]（Transformer 里的线性层；卷积可以用 im2col 写成矩阵乘法）
    实现这个分析模型，并提供：
        Mapping          ：一个完整的映射（每一维在 RF / 空间 / GB / DRAM 四层的因子 + 两级循环顺序）
        access_counts    ：解析地算出每个张量在每一级的读写次数、PE 之间的传递（hop）次数
        brute_force_fills：逐次迭代、真的数"缓冲区里的块换了几次"，用来验证解析公式
        energy           ：按 Interstellar 表 3 的能量表算总能量和各级的分解
        search           ：给定数据流（哪些维度展开到阵列的行 / 列），枚举所有合法的分块和循环顺序，找能量最低的

【模型的简化（务必知道）】
    - 和 Interstellar / Timeloop 一样，每一级为每个张量放"一个块"（另有一份用于双缓冲，只影响容量），
      块里的数据在被换掉之前反复使用；不模拟 cache 的替换细节。
    - 输出 C 的部分和：第一次放进某一级时是 0（不用读），之后每次换出都要写回上一级、再换回时要读。
    - 空间展开：一个张量和某个展开的维度无关时，几个 PE 需要同样的数据——从 GB 读一次，经过 PE 之间的传递（hop）送到其他 PE；
      输出在归约维度（K）上展开时，几个 PE 的部分和经过 hop 加起来再写回 GB。
    - 不模拟带宽和停顿：能量模型只数访问次数。阵列利用率 = 用到的 PE / 全部 PE。

【运行】python3 loopnest.py   # 自检：解析公式 = 暴力计数
【依赖】numpy；标准库 itertools。
"""

import itertools
import math
from dataclasses import dataclass, field

import numpy as np

DIMS = ("M", "N", "K")
TENSORS = ("A", "B", "C")
REL = {"A": ("M", "K"), "B": ("K", "N"), "C": ("M", "N")}       # 每个张量的下标里有哪些维度

# Interstellar 表 3（28 nm、16 位数据、每次访问的能量，单位 pJ）
RF_ENERGY = {16: 0.03, 32: 0.06, 64: 0.12, 128: 0.24, 256: 0.48, 512: 0.96}
SRAM_ENERGY = {32 * 1024: 6.0, 64 * 1024: 9.0, 128 * 1024: 13.5, 256 * 1024: 20.25, 512 * 1024: 30.375}
E_MAC, E_HOP, E_DRAM = 0.075, 0.035, 200.0


def rf_energy(nbytes):
    """寄存器堆每次访问的能量：表 3 里 16 B → 0.03 pJ，大小每翻一倍能量翻一倍（线性），中间的值按线性插值。"""
    return 0.03 * nbytes / 16


def sram_energy(nbytes):
    """SRAM 每次访问的能量：表 3 里 32 KB → 6 pJ，大小每翻一倍能量 ×1.5；其他大小按这个规律外推。"""
    return 6.0 * 1.5 ** math.log2(nbytes / (32 * 1024))


@dataclass(frozen=True)
class Arch:
    """
    加速器的资源：rows × cols 的 PE 阵列，每个 PE 一个 rf_bytes 的寄存器堆，共享一个 gb_bytes 的全局缓冲（SRAM）。

    【默认值】和 Interstellar 拿来比较的 Eyeriss 配置一样：16×16 的阵列、512 B 的 RF、128 KB 的 GB。
    """
    rows: int = 16
    cols: int = 16
    rf_bytes: int = 512
    gb_bytes: int = 128 * 1024
    elem_bytes: int = 2

    @property
    def e_rf(self):
        return rf_energy(self.rf_bytes)

    @property
    def e_gb(self):
        return sram_energy(self.gb_bytes)


@dataclass
class Mapping:
    """
    一个映射：每一维 d 的大小 = rf[d] × sp[d] × gb[d] × dram[d]。

    【字段】
        rf, gb, dram：{维度: 因子}——PE 里的时间循环、GB 层的时间循环（遍历 RF 块）、DRAM 层的时间循环（遍历 GB 块）
        sp_rows, sp_cols：展开到阵列行 / 列的维度及因子，例如 (("K", 16),) 表示 K 维的 16 段放在 16 行上；
                          一个方向上放两个维度就是"复制"（replication，Interstellar 图 2）
        order_gb, order_dram：这两层循环从外到内的顺序（RF 层内部的顺序不影响访问次数，不需要）
    【方法】sp(d)：维度 d 一共展开了几份；n_pe()：用到几个 PE
    """
    rf: dict
    gb: dict
    dram: dict
    sp_rows: tuple = ()
    sp_cols: tuple = ()
    order_gb: tuple = DIMS
    order_dram: tuple = DIMS
    label: str = field(default="")

    def sp(self, d):
        f = 1
        for dd, k in self.sp_rows + self.sp_cols:
            if dd == d:
                f *= k
        return f

    def n_pe(self):
        return int(np.prod([k for _, k in self.sp_rows + self.sp_cols])) if (self.sp_rows or self.sp_cols) else 1

    def size(self, d):
        return self.rf[d] * self.sp(d) * self.gb[d] * self.dram[d]


def tile(m, t, level):
    """张量 t 在某一级的块有多少个元素。level = "rf"（一个 PE 里）或 "gb"。"""
    n = 1
    for d in REL[t]:
        n *= m.rf[d] if level == "rf" else m.rf[d] * m.sp(d) * m.gb[d]
    return n


def versions(loops, t):
    """
    一个张量在某一级的块一共"换了几次"（被装进来几个版本）。

    【参数】loops：这一级**之上**的全部时间循环，从外到内，元素是 (维度, 上界)
    【规则（驻留 / stationarity）】从最外层开始，把上界乘起来，乘到"最内层的那个和 t 有关、上界 > 1 的循环"为止。
        它里面那些和 t 无关的循环转一圈，t 的块不变——不用重新装（这就是"t 驻留"）。
        例：C[M,N] 的块，循环从外到内 (M, N, K)：K 在最里面、和 C 无关 → C 的块装 M·N 次，在 K 上累加时一直留着（输出驻留）。
    【返回】整数
    """
    last = -1
    for i, (d, b) in enumerate(loops):
        if b > 1 and d in REL[t]:
            last = i
    v = 1
    for d, b in loops[:last + 1]:
        v *= b
    return v


def access_counts(m):
    """
    解析地算出每个张量在每一级的访问次数（单位：元素，即 16 位数的个数）。

    【返回】字典 {张量: {...}}，每个张量有：
        rf_fill   ：写进各个 PE 的 RF 的总次数（所有 PE 加起来）
        gb_read   ：为了填 RF 从 GB 读出的次数（多播时只读一次）
        gb_write  ：（只对 C）RF 写回 GB 的次数
        hops      ：PE 之间传递的次数（多播分发、空间归约）
        gb_fill   ：写进 GB 的次数（从 DRAM 读来的）
        dram_read / dram_write：DRAM 的读写次数
        rf_mac    ：每次乘加对 RF 的访问（A、B 各读 1 次；C 读 1 次写 1 次）
    还有 "MACs"：乘加的总次数。
    【思路】
        RF：这一级之上的循环 = DRAM 层（外）+ GB 层（内）。版本数 v_rf = versions(...)；每个 PE 每个版本装一个 RF 块。
            n_pe 个 PE 里，和 t 无关的那些展开维度上的 PE 需要同样的数据（多播因子 mc）：GB 只读 1/mc，其余靠 hop 传。
        GB：这一级之上的循环 = DRAM 层。版本数 v_gb；每个版本从 DRAM 装一个 GB 块。
        C 的写回：每个版本最后都要写回上一级（总共 v 次），第一次装入的那 distinct 个版本不用读（从 0 开始）。
    """
    loops_gb = [(d, m.dram[d]) for d in m.order_dram]
    loops_rf = loops_gb + [(d, m.gb[d]) for d in m.order_gb]
    n_pe = m.n_pe()
    macs = 1
    for d in DIMS:
        macs *= m.size(d)
    out = {"MACs": macs}
    for t in TENSORS:
        rf_t, gb_t = tile(m, t, "rf"), tile(m, t, "gb")
        v_rf, v_gb = versions(loops_rf, t), versions(loops_gb, t)
        mc = 1
        for d in DIMS:
            if d not in REL[t]:
                mc *= m.sp(d)
        rec = {"rf_fill": v_rf * rf_t * n_pe, "gb_fill": v_gb * gb_t}
        if t != "C":
            rec["gb_read"] = rec["rf_fill"] // mc
            rec["hops"] = rec["rf_fill"] - rec["gb_read"]
            rec["dram_read"] = rec["gb_fill"]
            rec["dram_write"] = 0
            rec["gb_write"] = 0
            rec["rf_mac"] = macs
        else:
            k_sp = m.sp("K")
            distinct_rf = 1
            for d in REL["C"]:
                distinct_rf *= m.size(d) // m.rf[d]               # 不同的 RF 块（按 PE 计：每个 PE 一份）
            distinct_rf = distinct_rf * k_sp                         # K 展开时每个 PE 各有一份部分和
            writes = rec["rf_fill"]                                 # 每个版本最后都写回
            reads = max(rec["rf_fill"] - distinct_rf * rf_t, 0)     # 第一次从 0 开始，不用读
            rec["gb_write"] = writes // k_sp                        # K 方向的部分和先经过 hop 加起来，再写回 GB 一次
            rec["gb_read"] = reads // k_sp                          # 续算时只有一个 PE 需要旧的部分和，其余从 0 开始
            rec["hops"] = writes - writes // k_sp                   # 空间归约：每次写回，k_sp 个部分和要传 k_sp − 1 次
            distinct_gb = m.size("M") * m.size("N")
            rec["dram_write"] = rec["gb_fill"]
            rec["dram_read"] = max(rec["gb_fill"] - distinct_gb, 0)
            rec["rf_mac"] = 2 * macs
        out[t] = rec
    return out


def energy(m, arch, counts=None):
    """
    总能量（pJ）和分解：{"MAC", "RF", "array"（hop）, "GB", "DRAM", "total"}。

    【公式】（Interstellar §5：E = Σ 各级访问次数 × 每次访问的能量）
        RF   ：(每次乘加的 RF 访问 + 往 RF 里装数据 + C 从 RF 读出来写回) × e_rf
        array：hop 次数 × e_hop
        GB   ：A、B：(送给 RF 的读 + 从 DRAM 装进来的写)；C：(RF 续算时的读 + RF 写回 + 和 DRAM 之间的换入换出) × e_gb
        DRAM ：(读 + 写) × e_dram
    """
    c = counts or access_counts(m)
    e = {"MAC": c["MACs"] * E_MAC, "RF": 0.0, "array": 0.0, "GB": 0.0, "DRAM": 0.0}
    for t in TENSORS:
        r = c[t]
        e["RF"] += (r["rf_mac"] + r["rf_fill"] + (r["rf_fill"] if t == "C" else 0)) * arch.e_rf
        e["array"] += r["hops"] * E_HOP
        if t == "C":                                   # C：RF 读 / 写回 + 和 DRAM 之间的换入 / 换出
            gb_acc = r["gb_read"] + r["gb_write"] + r["dram_read"] + r["dram_write"]
        else:                                          # A、B：送给 RF 的读 + 从 DRAM 装进来的写
            gb_acc = r["gb_read"] + r["gb_fill"]
        e["GB"] += gb_acc * arch.e_gb
        e["DRAM"] += (r["dram_read"] + r["dram_write"]) * E_DRAM
    e["total"] = sum(e.values())
    return e


def fits(m, arch, double_buffer=True):
    """容量检查：RF 放得下三个张量的 RF 块；GB 放得下三个张量的 GB 块（双缓冲时 ×2）。"""
    rf_need = sum(tile(m, t, "rf") for t in TENSORS) * arch.elem_bytes
    gb_need = sum(tile(m, t, "gb") for t in TENSORS) * arch.elem_bytes * (2 if double_buffer else 1)
    return rf_need <= arch.rf_bytes and gb_need <= arch.gb_bytes


def brute_force_fills(m):
    """
    逐次迭代整个循环嵌套（DRAM 层 → GB 层，从外到内），真的数每个张量在 RF 和 GB 两级"块换了几次"。

    【用途】验证 versions() 的驻留规则。只适合小问题（迭代次数 = 所有 DRAM、GB 层因子的乘积）。
    【返回】{张量: (RF 的版本数, GB 的版本数)}
    【Python 小知识】itertools.product(*ranges) 按"最右边变得最快"的顺序生成所有下标组合，正好是嵌套循环的顺序。
    """
    loops = [("dram", d, m.dram[d]) for d in m.order_dram] + [("gb", d, m.gb[d]) for d in m.order_gb]
    res = {}
    for t in TENSORS:
        last_rf = last_gb = None
        n_rf = n_gb = 0
        for idx in itertools.product(*[range(b) for _, _, b in loops]):
            key_gb = tuple(i for (lvl, d, _), i in zip(loops, idx) if lvl == "dram" and d in REL[t])
            key_rf = tuple(i for (lvl, d, _), i in zip(loops, idx) if d in REL[t])
            if key_gb != last_gb:
                n_gb += 1
                last_gb = key_gb
            if key_rf != last_rf:
                n_rf += 1
                last_rf = key_rf
        res[t] = (n_rf, n_gb)
    return res


# -----------------------------------------------------------------------------
# 搜索：给定数据流，枚举分块和循环顺序
# -----------------------------------------------------------------------------
def divisors(n):
    return [k for k in range(1, n + 1) if n % k == 0]


def dataflow_spatial(shape, rows_dims, cols_dims, arch):
    """
    按数据流把维度展开到阵列上：rows_dims / cols_dims 是维度名的元组（一个方向放两个维度 = 复制）。
    每个维度取"不超过剩余的行（列）数、又能整除这一维"的最大因子；返回 (sp_rows, sp_cols)。
    """
    def assign(dims, cap):
        out = []
        left = cap
        for d in dims:
            f = max(k for k in divisors(shape[d]) if k <= left)
            out.append((d, f))
            left //= f
        return tuple(out)
    return assign(rows_dims, arch.rows), assign(cols_dims, arch.cols)


def search(shape, arch, rows_dims=("K",), cols_dims=("N",), keep_all=False, max_rf=64):
    """
    给定矩阵乘法的形状 {"M","N","K"}、硬件和数据流，枚举所有合法的（RF、GB、DRAM 三层的因子）×（GB、DRAM 两层的循环顺序），
    返回能量最低的映射。

    【参数】rows_dims / cols_dims：展开到行 / 列的维度（Interstellar 的 U|V 记号：rows|cols）
        keep_all：True 时额外返回所有合法映射的总能量（看分布用）
        max_rf：每一维在 RF 层的因子上限（RF 很小，不需要枚举大的）
    【返回】(最好的 Mapping, 它的能量分解, 合法映射的个数[, 所有能量的数组])
    【边界】没有合法映射（RF 太小放不下哪怕 1×1×1 的块）时返回 (None, None, 0)。
    """
    sp_rows, sp_cols = dataflow_spatial(shape, rows_dims, cols_dims, arch)
    base = Mapping({d: 1 for d in DIMS}, {d: 1 for d in DIMS}, {d: 1 for d in DIMS}, sp_rows, sp_cols)
    rest = {d: shape[d] // base.sp(d) for d in DIMS}
    opts = {}
    for d in DIMS:
        o = []
        for r in divisors(rest[d]):
            if r > max_rf:
                continue
            for g in divisors(rest[d] // r):
                o.append((r, g, rest[d] // r // g))
        opts[d] = o
    best = (None, None)
    n_ok = 0
    all_e = []
    orders = list(itertools.permutations(DIMS))
    for fm, fn, fk in itertools.product(opts["M"], opts["N"], opts["K"]):
        rf = {"M": fm[0], "N": fn[0], "K": fk[0]}
        gb = {"M": fm[1], "N": fn[1], "K": fk[1]}
        dr = {"M": fm[2], "N": fn[2], "K": fk[2]}
        m0 = Mapping(rf, gb, dr, sp_rows, sp_cols)
        if not fits(m0, arch):
            continue
        for og in orders:
            for od in orders:
                m = Mapping(rf, gb, dr, sp_rows, sp_cols, og, od)
                e = energy(m, arch)
                n_ok += 1
                if keep_all:
                    all_e.append(e["total"])
                if best[1] is None or e["total"] < best[1]["total"]:
                    best = (m, e)
    if keep_all:
        return best[0], best[1], n_ok, np.array(all_e)
    return best[0], best[1], n_ok


def _versions_vec(bounds, dims_seq, t):
    """versions() 的向量化版本：bounds 是 [映射个数, 循环个数] 的上界，dims_seq 是这些循环对应的维度名。"""
    n, L = bounds.shape
    rel = np.array([d in REL[t] for d in dims_seq])
    pos = np.where(rel[None, :] & (bounds > 1), np.arange(L)[None, :], -1).max(axis=1)
    cp = np.cumprod(bounds, axis=1)
    return np.where(pos >= 0, cp[np.arange(n), np.maximum(pos, 0)], 1)


def search_fast(shape, arch, rows_dims=("K",), cols_dims=("N",), max_rf=64):
    """
    search() 的向量化版本：同样的搜索空间、同样的公式，用 numpy 一次算一批映射。

    【为什么需要】真实尺寸的矩阵乘法（几百到几千）有几十万个合法映射，逐个用 Python 算要几十秒；向量化后不到一秒。
    【返回】(最好的 Mapping, 它的能量分解（用标量版的 energy() 重新算，保证一致）, 合法映射的个数, 所有总能量的数组)
    """
    sp_rows, sp_cols = dataflow_spatial(shape, rows_dims, cols_dims, arch)
    base = Mapping({d: 1 for d in DIMS}, {d: 1 for d in DIMS}, {d: 1 for d in DIMS}, sp_rows, sp_cols)
    sp = {d: base.sp(d) for d in DIMS}
    n_pe = base.n_pe()
    rest = {d: shape[d] // sp[d] for d in DIMS}
    opts = {}
    for d in DIMS:
        opts[d] = np.array([(r, g, rest[d] // r // g) for r in divisors(rest[d]) if r <= max_rf
                            for g in divisors(rest[d] // r)])
    grid = np.array(list(itertools.product(range(len(opts["M"])), range(len(opts["N"])), range(len(opts["K"])))))
    F = {d: opts[d][grid[:, i]] for i, d in enumerate(DIMS)}          # 每一维：[n, 3] = (rf, gb, dram)
    rf = {d: F[d][:, 0] for d in DIMS}
    gb = {d: F[d][:, 1] for d in DIMS}
    dr = {d: F[d][:, 2] for d in DIMS}
    size = {d: shape[d] for d in DIMS}

    def tile_v(t, level):
        n = np.ones(len(grid), dtype=np.int64)
        for d in REL[t]:
            n = n * (rf[d] if level == "rf" else rf[d] * sp[d] * gb[d])
        return n
    rf_t = {t: tile_v(t, "rf") for t in TENSORS}
    gb_t = {t: tile_v(t, "gb") for t in TENSORS}
    ok = (sum(rf_t.values()) * arch.elem_bytes <= arch.rf_bytes) & \
         (sum(gb_t.values()) * arch.elem_bytes * 2 <= arch.gb_bytes)
    idx = np.nonzero(ok)[0]
    macs = shape["M"] * shape["N"] * shape["K"]
    k_sp = sp["K"]
    best = (None, np.inf)
    all_e = []
    for og in itertools.permutations(DIMS):
        for od in itertools.permutations(DIMS):
            b_gbl = np.stack([dr[d][idx] for d in od], axis=1)
            b_rfl = np.concatenate([b_gbl, np.stack([gb[d][idx] for d in og], axis=1)], axis=1)
            seq_gb, seq_rf = list(od), list(od) + list(og)
            e_tot = np.full(len(idx), macs * E_MAC, dtype=np.float64)
            for t in TENSORS:
                v_rf = _versions_vec(b_rfl, seq_rf, t)
                v_gb = _versions_vec(b_gbl, seq_gb, t)
                rf_fill = v_rf * rf_t[t][idx] * n_pe
                gb_fill = v_gb * gb_t[t][idx]
                if t != "C":
                    mc = 1
                    for d in DIMS:
                        if d not in REL[t]:
                            mc *= sp[d]
                    gb_read = rf_fill // mc
                    hops = rf_fill - gb_read
                    e_tot += (macs + rf_fill) * arch.e_rf + hops * E_HOP + (gb_read + gb_fill) * arch.e_gb \
                        + gb_fill * E_DRAM
                else:
                    distinct_rf = (size["M"] // rf["M"][idx]) * (size["N"] // rf["N"][idx]) * k_sp
                    reads = np.maximum(rf_fill - distinct_rf * rf_t[t][idx], 0)
                    gb_write = rf_fill // k_sp
                    gb_read = reads // k_sp
                    hops = rf_fill - gb_write
                    dram_read = np.maximum(gb_fill - size["M"] * size["N"], 0)
                    e_tot += (2 * macs + 2 * rf_fill) * arch.e_rf + hops * E_HOP \
                        + (gb_read + gb_write + dram_read + gb_fill) * arch.e_gb + (dram_read + gb_fill) * E_DRAM
            all_e.append(e_tot)
            j = int(np.argmin(e_tot))
            if e_tot[j] < best[1]:
                i = idx[j]
                m = Mapping({d: int(rf[d][i]) for d in DIMS}, {d: int(gb[d][i]) for d in DIMS},
                            {d: int(dr[d][i]) for d in DIMS}, sp_rows, sp_cols, og, od)
                best = (m, e_tot[j])
    all_e = np.concatenate(all_e) if all_e else np.array([])
    if best[0] is None:
        return None, None, 0, all_e
    return best[0], energy(best[0], arch), len(all_e), all_e


def utilization(shape, rows_dims, cols_dims, arch):
    """数据流在阵列上的利用率 = 用到的 PE / 全部 PE（维度太小或不能整除时，一部分 PE 空着）。"""
    sr, sc = dataflow_spatial(shape, rows_dims, cols_dims, arch)
    used = int(np.prod([k for _, k in sr + sc]))
    return used / (arch.rows * arch.cols)


def _selftest():
    rng = np.random.default_rng(0)
    for _ in range(20):
        dims = {d: int(rng.choice([1, 2, 3, 4])) for d in DIMS}
        gb = {d: int(rng.choice([1, 2, 3])) for d in DIMS}
        rf = {d: int(rng.choice([1, 2])) for d in DIMS}
        og = tuple(rng.permutation(DIMS))
        od = tuple(rng.permutation(DIMS))
        m = Mapping(rf, gb, dims, (("K", 2),), (("N", 2),), og, od)
        bf = brute_force_fills(m)
        loops_gb = [(d, m.dram[d]) for d in m.order_dram]
        loops_rf = loops_gb + [(d, m.gb[d]) for d in m.order_gb]
        for t in TENSORS:
            assert bf[t] == (versions(loops_rf, t), versions(loops_gb, t)), (t, m, bf[t])
    arch = Arch()
    best, e, n = search({"M": 64, "N": 64, "K": 64}, arch)
    assert best is not None and fits(best, arch)
    _, e2, n2, _ = search_fast({"M": 64, "N": 64, "K": 64}, arch)
    assert n2 == n and abs(e2["total"] - e["total"]) < 1e-6 * e["total"], (e["total"], e2["total"])
    print(f"loopnest 自检通过：20 个随机映射的驻留公式 = 暴力计数；逐个计算的搜索 = 向量化的搜索；64³ 的矩阵乘法在 16×16、512 B RF 上有 {n:,} 个合法映射，"
          f"最低 {e['total'] / best.size('M') / best.size('N') / best.size('K'):.2f} pJ/MAC")


if __name__ == "__main__":
    _selftest()