"""
=============================================================================
第 5 讲 · 代码 4：分块（tiling）—— 矩阵乘法为什么能跑满 GPU
=============================================================================

对应讲义：lectures/lecture05_gpus_tpus.md 第 4 节（技巧 6）

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
01 号文件说矩阵乘法的算术强度约为 n/3，很高。但这有一个前提：每个输入元素只从显存读一次。
如果按最朴素的方式写 —— 每个输出元素 C[i,j] 都去显存里读 A 的第 i 行和 B 的第 j 列 ——
同一个 A[i,k] 会被读 n 次，算术强度掉回 O(1)，完全被内存卡住。

**分块**的思路：把矩阵切成 T×T 的小块，把一对小块搬进片上的共享内存（SRAM，快但小），
在里面把它们能做的乘加全部做完，再换下一对。这样每个元素从显存读的次数从 n 次降到 n/T 次。

这个文件：
    1. 模拟"朴素"和"分块"两种矩阵乘法，数一数从显存读了多少个元素，验证 2n³ vs 2n³/T；
    2. 分块大小受共享内存容量限制：H100 上最大能用多大的块？
    3. 分块的副作用：矩阵边长不能被块大小整除时的"块量化"浪费。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 04_tiling_matmul.py
"""

import math

import torch


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


# =============================================================================
# 第 1 部分：朴素 vs 分块
# =============================================================================
def naive_matmul_loads(n):
    """
    朴素矩阵乘法从显存读取的元素个数：每个输出 C[i,j] 读 A 的一行（n 个）和 B 的一列（n 个）。

    【返回】2 × n × n × n = 2n³
    """
    return 2 * n ** 3


def tiled_matmul(A, B, T):
    """
    分块矩阵乘法（模拟 GPU 上一个线程块计算一个输出块的过程），同时统计从"显存"读了多少个元素。

    【算法】对每个 T×T 的输出块 C[I, J]：
        累加器 acc = 0（在寄存器里）
        对 K = 0, T, 2T, …：
            把 A[I, K] 和 B[K, J] 两个 T×T 小块从显存搬到共享内存   ← 计数：2T² 个元素
            acc += A_tile @ B_tile                                   ← 只访问共享内存，不计数
        把 acc 写回 C[I, J]

    【参数】A、B : n×n 张量，n 必须能被 T 整除（为了简单；不能整除的情况见第 3 部分）
    【返回】(C, 读取的元素个数)
    【计数】输出块有 (n/T)² 个，每个要搬 n/T 对小块，每对 2T² 个元素：
            (n/T)² × (n/T) × 2T² = 2n³ / T
    """
    n = A.shape[0]
    if n % T != 0:
        raise ValueError("为了简单，这里要求 n 能被 T 整除")
    C = torch.zeros(n, n, dtype=A.dtype)
    loads = 0
    for I in range(0, n, T):
        for J in range(0, n, T):
            acc = torch.zeros(T, T, dtype=A.dtype)            # 寄存器里的累加器
            for K in range(0, n, T):
                a_tile = A[I:I + T, K:K + T]                  # 从显存读到共享内存
                b_tile = B[K:K + T, J:J + T]
                loads += a_tile.numel() + b_tile.numel()
                acc += a_tile @ b_tile                        # 在共享内存里完成这一块的乘加
            C[I:I + T, J:J + T] = acc                         # 写回显存
    return C, loads


def demo_tiling():
    """演示 1：分块把显存读取量降低 T 倍。"""
    print_header("演示 1：朴素 vs 分块 —— 数一数从显存读了多少")
    torch.manual_seed(0)
    n = 256
    A, B = torch.randn(n, n, dtype=torch.float64), torch.randn(n, n, dtype=torch.float64)
    ref = A @ B
    flops = 2 * n ** 3
    print(f"  n = {n}，FLOPs = 2n³ = {flops:,}")
    print(f"  {'方式':<14}{'从显存读取的元素':>18}{'公式':>16}{'算术强度(FP32)':>16}")
    print(f"  {'朴素':<16}{naive_matmul_loads(n):>18,}{'2n³':>16}{flops / (naive_matmul_loads(n) * 4):>16.2f}")
    for T in [4, 16, 64]:
        C, loads = tiled_matmul(A, B, T)
        assert torch.allclose(C, ref), "分块结果应当与普通矩阵乘法一致"
        assert loads == 2 * n ** 3 // T
        print(f"  {'分块 T=' + str(T):<16}{loads:>18,}{'2n³/' + str(T):>16}{flops / (loads * 4):>16.2f}")
    print("""
  ✓ 分块的结果与普通矩阵乘法完全一致，显存读取量正好是朴素方式的 1/T。
  朴素方式的算术强度只有 0.25 FLOP/字节（每读一个 4 字节的数，只做 1 次浮点运算），无论矩阵多大都是内存受限；
  分块后强度提高 T 倍 —— 块越大，每个搬进来的元素被重复使用得越多。
  另一个好处：搬一个小块时，同一个 warp 的线程读的是连续的一段内存，自然就是"合并访存"（02 号文件）。""")


# =============================================================================
# 第 2 部分：块能有多大？
# =============================================================================
def demo_tile_size_limit():
    """演示 2：共享内存容量决定了块的上限。"""
    print_header("演示 2：块能有多大？—— 受共享内存容量限制")
    smem = 228 * 1024                               # H100 每个 SM 的共享内存
    for name, b in [("FP32", 4), ("BF16", 2)]:
        T = int(math.sqrt(smem / (2 * b)))          # 需要同时放下 A、B 两个 T×T 的块
        print(f"  {name}：两个 T×T 的块要放进 228 KB → T 最大约 {T}")
    print("""
  实际的矩阵乘法核函数（cuBLAS、CUTLASS）通常用 128×128、128×256 这样的矩形输出块，
  沿 K 方向每次只搬 32 或 64 列（这样 A、B 的小块更小），累加器放在寄存器里，
  并且一边计算当前块、一边异步地预取下一块（"双缓冲"），把搬运和计算重叠起来。
  块也不能太大：块越大，块的总数越少，可能喂不饱所有 SM（02 号文件的波次量化）；
  每个线程占用的寄存器越多，占用率越低。块大小是一个需要权衡、常常靠实测来选的参数。""")


# =============================================================================
# 第 3 部分：块量化
# =============================================================================
def tile_efficiency(n, T):
    """
    矩阵边长 n 不能被块大小 T 整除时，最后一行/列的块只有一部分是有用的。

    【返回】有用计算 / 实际计算 = (n / (⌈n/T⌉ × T))²
    """
    padded = math.ceil(n / T) * T
    return (n / padded) ** 2


def demo_tile_quantization():
    """演示 3：块量化带来的浪费。"""
    print_header("演示 3：块量化 —— 矩阵边长最好是块大小的整数倍")
    T = 128
    print(f"  块大小 T = {T}：")
    for n in [1024, 1025, 1088, 1152, 1153, 4096, 4097]:
        blocks = math.ceil(n / T)
        print(f"    n = {n:>5}：每维 {blocks:>2} 块，有效计算占比 {tile_efficiency(n, T):.1%}")
    print("""
  n = 1025 时，每一维都多出一个几乎全空的块，约 21% 的计算是在算"填充的 0"。
  矩阵越大，这个浪费的比例越小（4097 时只有 6%）。
  另外还有"对齐"问题：如果一行的字节数不是 16/32/128 字节的倍数，块的起始地址就对不齐内存事务，
  无法合并访存。所以：**让模型的各个维度取 64 或 128 的倍数**，这是最便宜的性能优化之一。""")


def main():
    demo_tiling()
    demo_tile_size_limit()
    demo_tile_quantization()
    print("\n全部演示运行完毕。")


if __name__ == "__main__":
    main()