"""
=============================================================================
第 5 讲 · 代码 6：TPU 的心脏 —— 脉动阵列（systolic array）
=============================================================================

对应讲义：lectures/lecture05_gpus_tpus.md 第 7 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
GPU 用成千上万个小核心 + 分块来做矩阵乘法；Google 的 TPU 走了另一条路：
用一个巨大的"脉动阵列"（矩阵乘法单元 MXU，例如 128×128 个乘加单元）专门做矩阵乘法。

脉动阵列的想法：数据像心脏泵血一样有节奏地"流过"一个网格状的处理单元阵列：
    - A 的每一行从左边流入，每个周期向右移动一格；
    - B 的每一列从上边流入，每个周期向下移动一格；
    - 每个处理单元（PE）在每个周期把经过它的 a 和 b 相乘，累加到自己的寄存器里。
每个数据从内存读一次，就在阵列中被相邻单元反复使用 —— 不需要反复访问内存，
这和 GPU 上"分块 + 共享内存"追求的目标（04 号文件）是一样的，只是用硬件结构直接实现了。

这个文件逐周期模拟一个小的脉动阵列，验证它算出的结果与矩阵乘法一致，
并统计周期数、利用率和内存读取次数。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 06_systolic_array.py
"""

import torch


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def systolic_matmul(A, B, verbose=False):
    """
    逐周期模拟一个 P×P 的"输出驻留"（output-stationary）脉动阵列，计算 C = A B。

    【结构】
        PE(i, j) 负责计算 C[i, j]，它有三个寄存器：
            a_reg：从左边邻居传来的 A 元素（下一个周期传给右边邻居）
            b_reg：从上边邻居传来的 B 元素（下一个周期传给下边邻居）
            acc  ：累加器，最终就是 C[i, j]
        每个周期：
            1. 所有 PE 同时把 a_reg 往右传、b_reg 往下传（最左列 / 最上行从外部读入新数据）；
            2. 每个 PE 计算 acc += a_reg × b_reg。
    【为什么要"错位"输入】
        A 的第 i 行要晚 i 个周期进入阵列，B 的第 j 列要晚 j 个周期进入，
        这样 A[i, k] 和 B[k, j] 才能恰好在同一个周期到达 PE(i, j)。
        （输入呈"平行四边形"，这是脉动阵列示意图里常见的错位形状。）

    【参数】A : (P, K)，B : (K, P)（为了简单，阵列大小 P = A 的行数 = B 的列数）
    【返回】(C, 周期数, 有效乘加次数)
    """
    P, K = A.shape
    assert B.shape == (K, P)
    a_reg = torch.zeros(P, P, dtype=A.dtype)
    b_reg = torch.zeros(P, P, dtype=A.dtype)
    acc = torch.zeros(P, P, dtype=A.dtype)
    cycles = K + 2 * (P - 1)                   # 最后一个数据 A[P-1, K-1] 在第 K-1+(P-1) 个周期进入，再走 P-1 格
    useful = 0
    for t in range(cycles):
        # 1. 数据流动：向右、向下各移一格
        a_reg = torch.roll(a_reg, shifts=1, dims=1)
        b_reg = torch.roll(b_reg, shifts=1, dims=0)
        # 最左列从外部读入：第 i 行在第 t 个周期读入 A[i, t - i]（错位 i 个周期）
        for i in range(P):
            k = t - i
            a_reg[i, 0] = A[i, k] if 0 <= k < K else 0.0
        # 最上行从外部读入：第 j 列在第 t 个周期读入 B[t - j, j]
        for j in range(P):
            k = t - j
            b_reg[0, j] = B[k, j] if 0 <= k < K else 0.0
        # 2. 所有 PE 同时乘加
        prod = a_reg * b_reg
        acc += prod
        useful += int((prod != 0).sum())       # 统计真正在做有效工作的 PE（输入为随机数时，0 表示空闲）
        if verbose and t < 4:
            print(f"    周期 {t}：正在工作的 PE 数 = {int((prod != 0).sum())}")
    return acc, cycles, useful


def demo_systolic():
    """演示 1：模拟脉动阵列，验证结果。"""
    print_header("演示 1：逐周期模拟一个 4×4 的脉动阵列")
    torch.manual_seed(0)
    P, K = 4, 6
    A = torch.randn(P, K, dtype=torch.float64)
    B = torch.randn(K, P, dtype=torch.float64)
    C, cycles, useful = systolic_matmul(A, B, verbose=True)
    print(f"  ……共 {cycles} 个周期")
    print(f"  与 A @ B 的最大差异：{(C - A @ B).abs().max().item():.1e}")
    total = P * P * cycles
    print(f"  有效乘加 {useful} 次（= P×P×K = {P * P * K}），PE 利用率 {useful / total:.0%}")
    print("""
  刚开始和快结束时，只有部分 PE 在工作（数据还在"流入"或"流出"阵列）——
  就像流水线的"填充"和"排空"。K 越大（矩阵越大），这段开销占比越小。""")


def demo_utilization_and_reuse():
    """演示 2：K 越大利用率越高；每个数据只从内存读一次。"""
    print_header("演示 2：利用率与数据复用")
    torch.manual_seed(0)
    P = 8
    print(f"  阵列大小 {P}×{P}：")
    for K in [8, 32, 128, 512]:
        A = torch.randn(P, K, dtype=torch.float64)
        B = torch.randn(K, P, dtype=torch.float64)
        C, cycles, useful = systolic_matmul(A, B)
        assert torch.allclose(C, A @ B)
        reads = A.numel() + B.numel()
        print(f"    K = {K:>4}：{cycles:>4} 个周期，利用率 {useful / (P * P * cycles):.0%}；"
              f"从内存读取 {reads:,} 个数，完成 {useful:,} 次乘加（每次乘加用到 2 个数，所以每个数被使用 {2 * useful / reads:.0f} 次）")
    print("""
  - 利用率随 K 增大趋近 100%：TPU 擅长"又大又规整"的矩阵乘法。
  - 每个数据只从内存读一次，然后在阵列里被 P 个 PE 依次使用（阵列越大，复用越多）。
    这就是 TPU 用很少的控制逻辑、很少的片上缓存，就能获得极高矩阵乘法吞吐量的秘密。""")


def demo_tpu_vs_gpu():
    """演示 3：TPU 与 GPU 的对比（文字总结）。"""
    print_header("演示 3：TPU 与 GPU 的异同")
    print("""
  相同点（CS336 课件）：核心结构都是"轻量的控制逻辑 + 很快的（大）矩阵乘法单元 + 很快的片上存储"。
  不同点：
    - GPU：很多个 SM（H100 有 132 个），每个 SM 里有若干个 Tensor Core 和大量普通计算单元，
           以 warp 为单位调度，灵活，适合各种运算。
    - TPU：每个芯片只有少数几个 TensorCore，每个 TensorCore 里有几个很大的 MXU
           （TPU v2~v5 是 128×128，v6e Trillium 是 256×256 的脉动阵列），没有 warp 的概念。
           两者的矩阵乘法峰值性能相近，但 TPU 在非矩阵乘法的运算上相对弱一些。
    - 联网方式：TPU 芯片之间用专门的高速互连（ICI）直接连成大规模的环面拓扑（pod），
                这对大规模并行训练（第 7、8 讲）非常重要。
  推荐阅读：Google 的"How to Scale Your Model"（jax-ml.github.io/scaling-book），
  CS336 课件也引用了它，其中既讲 TPU 也讲 GPU。""")


def main():
    demo_systolic()
    demo_utilization_and_reuse()
    demo_tpu_vs_gpu()
    print("\n全部演示运行完毕。")


if __name__ == "__main__":
    main()