"""
=============================================================================
CS349E · 第 10 讲 · 05_cpu_microkernel.py —— 同样的阶梯在你的 CPU 上：寄存器微内核、打包与矩阵协处理器
=============================================================================

【这个脚本要回答的问题】
    前面四个脚本的性能都是模型估计的。这个脚本在你自己的 CPU 上**真的测**：GPU 内核的那几步
    （延迟隐藏、寄存器分块、把数据搬进片上缓存并重排）在 CPU 上叫什么、各值多少。
      演示 1：Little 定律在一个 CPU 核上：同时有几条独立的 FMA 链，才能把向量单元喂满？
      演示 2：矩阵乘法的阶梯：朴素 ijk → ikj → 寄存器微内核 → 分块 + 打包（GotoBLAS / BLIS 的结构）→ numpy（厂商库）。
      演示 3：微内核的寄存器分块 MR×NR：太小喂不饱，太大寄存器不够（和 GPU 上每线程寄存器、占用率是同一个折中）。

【关于数字】这是墙钟测量，随机器、编译器和负载变化。讲义里引用的是作者机器（Apple M4 Pro，Apple clang 21）的结果；
    观察里的结论是按当次运行的数字写的。
【运行】python3 05_cpu_microkernel.py      （约 15 秒）
【依赖】numpy；系统里要有 C 编译器（没有时跳过）；../prereq4 的 cbuild（编译 C、计时）；../lecture01 的 console。
"""

import ctypes
import os
import platform
import subprocess
import sys

import numpy as np

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

from cbuild import best_time, compile_c, find_cc, ptr  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402

FLAGS = ("-O3",)

# 用 GCC/Clang 的向量扩展写的"FMA 吞吐"测试：NACC 条互不依赖的乘加链，每条链每步依赖上一步的结果。
# 写成 acc = acc + acc·a：累加器既是加数又是乘数，编译器生成一条 fmla（ARM）/ vfmadd（x86），不需要额外的寄存器搬运。
# （如果写成 acc·a + b，ARM 的 fmla 要求"加数就是目的寄存器"，编译器得先把 b 复制一份再乘加、再搬回来，
#  每条 FMA 多两条 mov，测出来的"峰值"只有真实的七成——这个脚本的第一版就是这样写错的，见 README。）
PEAK_SRC = r"""
typedef float v4 __attribute__((vector_size(16)));
float peak(long iters) {
  v4 acc[NACC];
  v4 a = {1e-7f, 1e-7f, 1e-7f, 1e-7f};
  for (int i = 0; i < NACC; i++) acc[i] = (v4){i + 1, i + 1, i + 1, i + 1};
  for (long t = 0; t < iters; t++) {
    #pragma clang loop unroll(full)
    for (int i = 0; i < NACC; i++) acc[i] = acc[i] + acc[i] * a;   /* 一条向量 FMA：4 个乘加 */
  }
  float s = 0;
  for (int i = 0; i < NACC; i++) s += acc[i][0] + acc[i][3];
  return s;                                                       /* 返回结果，防止编译器把循环删掉 */
}
"""

# 矩阵乘法的几种写法。C[n,n] = A[n,n]·B[n,n]，行主序；微内核要求 n 是 MR、NR 的倍数。
GEMM_SRC = r"""
static inline int imin(int a, int b) { return a < b ? a : b; }

/* 1. 朴素 ijk：最内层 k，B 按列走（步长 n），每次乘加都要从内存读一个新的 B 元素 */
void ijk(const float* restrict A, const float* restrict B, float* restrict C, int n) {
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++) {
      float s = 0;
      for (int k = 0; k < n; k++) s += A[(long)i*n + k] * B[(long)k*n + j];
      C[(long)i*n + j] = s;
    }
}

/* 2. ikj：最内层 j 连续，编译器能向量化（前置讲 P4 §8.1 的最佳循环顺序）；但 C 的一行每次都要读写内存 */
void ikj(const float* restrict A, const float* restrict B, float* restrict C, int n) {
  for (long i = 0; i < (long)n*n; i++) C[i] = 0;
  for (int i = 0; i < n; i++)
    for (int k = 0; k < n; k++) {
      float a = A[(long)i*n + k];
      for (int j = 0; j < n; j++) C[(long)i*n + j] += a * B[(long)k*n + j];
    }
}

/* 3. 寄存器微内核：每次算 C 的 MR×NR 小块，累加器 acc 留在寄存器里，走完整个 k 才写回一次
      （GPU 上 k5 的"每线程 TM×TN 个累加器"）*/
void micro_nopack(const float* restrict A, const float* restrict B, float* restrict C, int n) {
  for (int i0 = 0; i0 < n; i0 += MR)
    for (int j0 = 0; j0 < n; j0 += NR) {
      float acc[MR][NR] = {{0}};
      for (int k = 0; k < n; k++) {
        const float* b = B + (long)k*n + j0;
        for (int i = 0; i < MR; i++) {
          float a = A[(long)(i0 + i)*n + k];
          for (int j = 0; j < NR; j++) acc[i][j] += a * b[j];       /* 外积：MR 个 a × NR 个 b */
        }
      }
      for (int i = 0; i < MR; i++)
        for (int j = 0; j < NR; j++) C[(long)(i0 + i)*n + j0 + j] = acc[i][j];
    }
}

/* 4. 分块 + 打包（GotoBLAS / BLIS 的结构）：
      - 外两层把 B 切成 KC×NC 的块、A 切成 MC×KC 的块，让它们分别留在 L2 / L1（GPU 上的共享内存分块）；
      - "打包"：把块拷贝成微内核要的顺序——A 每 MR 行一条、B 每 NR 列一条，k 方向连续
        （GPU 上 k6 把 As 转置存进共享内存，是同一个想法：按计算的读取顺序重新排列）；
      - 微内核只读连续内存。*/
static void micro(int kc, const float* restrict Ap, const float* restrict Bp, float* restrict C, int ldc) {
  float acc[MR][NR] = {{0}};
  for (int p = 0; p < kc; p++) {
    const float* b = Bp + p*NR;
    for (int i = 0; i < MR; i++) {
      float a = Ap[p*MR + i];
      for (int j = 0; j < NR; j++) acc[i][j] += a * b[j];
    }
  }
  for (int i = 0; i < MR; i++)
    for (int j = 0; j < NR; j++) C[(long)i*ldc + j] += acc[i][j];
}

void gemm_packed(const float* restrict A, const float* restrict B, float* restrict C, int n,
                 int KC, int MC, int NC, float* restrict Abuf, float* restrict Bbuf) {
  for (long i = 0; i < (long)n*n; i++) C[i] = 0;
  for (int jc = 0; jc < n; jc += NC) {
    int nc = imin(NC, n - jc);
    for (int pc = 0; pc < n; pc += KC) {
      int kc = imin(KC, n - pc);
      for (int jr = 0; jr < nc; jr += NR)                                   /* 打包 B 的 KC×NC 块 */
        for (int p = 0; p < kc; p++)
          for (int jj = 0; jj < NR; jj++)
            Bbuf[(long)(jr/NR)*kc*NR + p*NR + jj] = B[(long)(pc + p)*n + jc + jr + jj];
      for (int ic = 0; ic < n; ic += MC) {
        int mc = imin(MC, n - ic);
        for (int ir = 0; ir < mc; ir += MR)                                 /* 打包 A 的 MC×KC 块 */
          for (int p = 0; p < kc; p++)
            for (int ii = 0; ii < MR; ii++)
              Abuf[(long)(ir/MR)*kc*MR + p*MR + ii] = A[(long)(ic + ir + ii)*n + pc + p];
        for (int jr = 0; jr < nc; jr += NR)
          for (int ir = 0; ir < mc; ir += MR)
            micro(kc, Abuf + (long)(ir/MR)*kc*MR, Bbuf + (long)(jr/NR)*kc*NR, C + (long)(ic + ir)*n + jc + jr, n);
      }
    }
  }
}
"""

KC, MC, NC = 256, 128, 1024
ARM = platform.machine().lower() in ("arm64", "aarch64")
APPLE = ARM and sys.platform == "darwin"
N_VREGS = 32 if ARM else 16       # ARM NEON：32 个 128 位寄存器；x86 只开 -O3（SSE2）：16 个 128 位寄存器


def gflops(n, t):
    return 2 * n ** 3 / t / 1e9


def build_gemm(mr, nr):
    return compile_c(f"#define MR {mr}\n#define NR {nr}\n" + GEMM_SRC, flags=FLAGS)


def numpy_gflops(n, threads):
    """在子进程里测 numpy 的矩阵乘法（子进程才能在 numpy 导入之前设置线程数的环境变量）。"""
    env = dict(os.environ)
    if threads is not None:
        env["VECLIB_MAXIMUM_THREADS"] = env["OPENBLAS_NUM_THREADS"] = env["OMP_NUM_THREADS"] = str(threads)
    code = ("import numpy as np, time\n"
            f"n = {n}\n"
            "A = np.random.default_rng(0).random((n, n), dtype=np.float32)\n"
            "B = np.random.default_rng(1).random((n, n), dtype=np.float32)\n"
            "A @ B\n"
            "best = min((lambda t0: (A @ B, time.perf_counter() - t0)[1])(time.perf_counter()) for _ in range(5))\n"
            "print(2 * n ** 3 / best / 1e9)\n")
    r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, env=env, timeout=120)
    return float(r.stdout.strip())


def demo1_littles_law():
    print_header("演示 1：一个 CPU 核要几条独立的 FMA 链才能喂满？（每条链 = 一个 4 宽的向量累加器）")
    rows, res = [], {}
    for nacc in (1, 2, 4, 8, 12, 16, 20):
        lib = compile_c(PEAK_SRC.replace("NACC", str(nacc)), flags=("-O3", "-ffp-contract=fast"))
        f = lib.peak
        f.restype = ctypes.c_float
        f.argtypes = [ctypes.c_long]
        iters = 20_000_000
        t = best_time(lambda: f(iters), 3)
        g = iters * nacc * 8 / t / 1e9
        res[nacc] = g
        rows.append([nacc, nacc * 4, f"{g:.1f}", f"{g / res[1]:.1f}"])
    print_table(["独立的向量累加器", "同时在算的乘加", "GFLOP/s", "相对 1 条链"], rows)
    lat_cycles_ratio = max(res.values()) / res[1]
    print_observations([
        f"只有 1 条链时每条 FMA 都要等上一条的结果：{res[1]:.1f} GFLOP/s，这就是\"延迟受限\"。"
        f"链越多越快，直到约 {max(res, key=res.get)} 条时到顶（{max(res.values()):.0f} GFLOP/s，约 1 条链的 {lat_cycles_ratio:.0f} 倍）。",
        "Little 定律：需要同时在飞的独立运算 = 吞吐 × 延迟。这个核每周期能发几条向量 FMA、每条要几个周期出结果，"
        "两者相乘就是\"至少要几条独立链\"——GPU 上同一件事靠很多 warp（占用率）和每个线程的多个累加器（ILP）来满足。",
        "寄存器分块的微内核天然有这么多独立链：4×16 的累加器块 = 16 个向量寄存器 = 16 条互不依赖的 FMA 链（演示 2、3）。",
    ])
    return max(res.values())


def demo2_ladder(peak):
    print_header("演示 2：矩阵乘法的阶梯（单核 FP32；numpy 那两行用的是厂商库）")
    rng = np.random.default_rng(0)
    lib = build_gemm(4, 16)
    rows, g = [], {}
    for n in (1024, 2048):
        A = rng.random((n, n), dtype=np.float32)
        B = rng.random((n, n), dtype=np.float32)
        C = np.zeros((n, n), np.float32)
        ref = A.astype(np.float64) @ B.astype(np.float64)
        Ab = np.zeros(MC * KC, np.float32)
        Bb = np.zeros(KC * NC, np.float32)
        cands = [("ikj（向量化的最内层）", lambda: lib.ikj(ptr(A), ptr(B), ptr(C), n)),
                 ("寄存器微内核 4×16（不打包）", lambda: lib.micro_nopack(ptr(A), ptr(B), ptr(C), n)),
                 (f"分块 + 打包 + 微内核（KC={KC}, MC={MC}, NC={NC}）",
                  lambda: lib.gemm_packed(ptr(A), ptr(B), ptr(C), n, KC, MC, NC, ptr(Ab), ptr(Bb)))]
        if n == 1024:
            cands.insert(0, ("朴素 ijk", lambda: lib.ijk(ptr(A), ptr(B), ptr(C), n)))
        for name, fn in cands:
            t = best_time(fn, 1 if name == "朴素 ijk" else 3)
            err = float(np.abs(C - ref).max() / np.abs(ref).max())
            g[(n, name.split("（")[0])] = gflops(n, t)
            rows.append([n, name, f"{t * 1e3:.1f} ms", f"{gflops(n, t):.1f}", f"{gflops(n, t) / peak:.0%}", f"{err:.1e}"])
        for thr, label in ((1, "numpy，限 1 个线程"), (None, "numpy，默认线程数")):
            gn = numpy_gflops(n, thr)
            g[(n, label)] = gn
            rows.append([n, label, "", f"{gn:.0f}", f"{gn / peak:.0%}", ""])
    print_table(["n", "写法", "时间", "GFLOP/s", "相对演示 1 的单核向量峰值", "相对误差"], rows)
    mk, pk = "寄存器微内核 4×16", "分块 + 打包 + 微内核"
    print_observations([
        f"朴素 ijk → ikj：只换循环顺序，最内层从\"按列跳着读 B\"变成连续读、能向量化，"
        f"{g[(1024, '朴素 ijk')]:.1f} → {g[(1024, 'ikj')]:.0f} GFLOP/s（前置讲 P4 §8.1）。",
        f"ikj → 寄存器微内核：累加器从\"每次读写内存里的 C\"变成\"留在寄存器里走完整个 k\"，n = 1024 时再快 "
        f"{g[(1024, mk)] / g[(1024, 'ikj')]:.1f} 倍，到了演示 1 测出的单核峰值的 {g[(1024, mk)] / peak:.0%}。这一步就是 GPU 上的 k4 / k5。",
        f"打包：n = 1024 时只多了 {g[(1024, pk)] / g[(1024, mk)] - 1:.0%}；n = 2048 时不打包的微内核掉到 {g[(2048, mk)]:.0f} GFLOP/s，"
        f"打包后仍是 {g[(2048, pk)]:.0f}。不打包时微内核读的 B 是一条 16 列宽、n 行高的竖条，相邻两行隔 n·4 字节——"
        "矩阵一大，这些散开的 cache line 和页就装不下了；打包把它们拷成一段连续内存。这就是 GPU 上\"先搬进共享内存、按计算的顺序重排\"的 CPU 版本。",
        f"numpy 调用的厂商库单线程就有 {g[(2048, 'numpy，限 1 个线程')]:.0f} GFLOP/s，"
        f"是我们最好的向量微内核的 {g[(2048, 'numpy，限 1 个线程')] / g[(2048, pk)]:.0f} 倍。"
        + ("据公开的分析，Apple 的 Accelerate 在 M 系列芯片上使用专门的矩阵协处理器（AMX / SME）——"
           "和 GPU 上 Tensor Core 对 CUDA core 的关系一样：同样的算法，换成专用的矩阵单元（第 8 讲）。" if APPLE else
           "OpenBLAS / MKL 用手写汇编的微内核和更宽的向量指令（AVX2 / AVX-512；我们的 C 代码只用 -O3 编译，没开 -march=native）。"),
    ])


def demo3_register_tile():
    print_header(f"演示 3：微内核的寄存器分块 MR×NR（n = 1024，不打包；这台机器有 {N_VREGS} 个 128 位向量寄存器）")
    rng = np.random.default_rng(1)
    n = 1024
    A = rng.random((n, n), dtype=np.float32)
    B = rng.random((n, n), dtype=np.float32)
    C = np.zeros((n, n), np.float32)
    ref = A.astype(np.float64) @ B.astype(np.float64)
    rows = []
    for mr, nr in ((1, 16), (2, 16), (4, 8), (4, 16), (8, 8), (8, 16), (4, 32)):
        lib = build_gemm(mr, nr)
        t = best_time(lambda: lib.micro_nopack(ptr(A), ptr(B), ptr(C), n), 3)
        acc = mr * nr // 4
        need = acc + nr // 4 + 1
        err = float(np.abs(C - ref).max() / np.abs(ref).max())
        rows.append([f"{mr}×{nr}", acc, need, "装得下" if need <= N_VREGS else f"超过 {N_VREGS}", f"{(mr + nr) / (mr * nr):.3f}",
                     f"{gflops(n, t):.1f}", f"{err:.1e}"])
    print_table(["MR×NR", "累加器（向量寄存器）", "大致需要的寄存器（+ B 的一行 + a）", "", "每次乘加读几个数",
                 "GFLOP/s", "相对误差"], rows)
    print_observations([
        "每次乘加要读 (MR + NR)/(MR·NR) 个数：块越大越省（和 GPU 演示 4 的 4·(TM + TN)/(TM·TN) 字节是同一个式子）。",
        "但累加器要放在寄存器里：8×16 和 4×32 需要 32 个累加器寄存器，加上 B 和 a 就超过了 32 个，编译器只好把一部分放回内存（spill），"
        "速度掉一半；4×16、8×8 这类刚好装下的块最快。",
        "GPU 上同样的约束表现为\"每线程寄存器多了，占用率就降\"；CPU 上是直接溢出到内存。两边的最优块大小都由寄存器的个数决定，"
        "而且要按机器调——BLIS 为每种 CPU 手写一个微内核，cuBLAS 为每种 GPU 和形状准备很多个内核。",
    ])


if __name__ == "__main__":
    if find_cc() is None:
        print("没有找到 C 编译器（cc / clang / gcc），这个脚本的演示都需要编译 C 代码，跳过。")
        sys.exit(0)
    pk = demo1_littles_law()
    demo2_ladder(pk)
    demo3_register_tile()