"""
=============================================================================
第 6 讲 · 代码 4：分块矩阵乘法 + ReLU 融合（Triton）
=============================================================================

对应讲义：lectures/lecture06_kernels_triton.md 第 8 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 5 讲 04 号文件用"数读写次数"的方式说明了分块为什么能让矩阵乘法跑满 GPU。
这里把它写成真正的 Triton 核函数：C = ReLU(A @ B)
    - 网格是二维的：程序 (pid_m, pid_n) 负责输出矩阵的一个 BLOCK_M × BLOCK_N 的块；
    - 程序内部沿 K 方向循环：每次读入 A 的一个 BLOCK_M × BLOCK_K 块和 B 的一个 BLOCK_K × BLOCK_N 块，
      用 tl.dot 相乘（在 GPU 上由 Tensor Core 完成），累加到寄存器里的累加器；
    - 循环结束后，**在写回之前**直接对累加器做 ReLU —— 这就是"融合"：
      不需要先把 A@B 写回显存、再启动一个 ReLU 核函数读出来。
这个文件还用模拟器统计实际读取了多少元素，验证第 5 讲的公式。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 04_triton_matmul.py
（有 NVIDIA GPU + triton 时用真正的 Triton；否则用 mini_triton 模拟器。）
"""

import torch

REAL_TRITON = False
try:
    import triton
    import triton.language as tl
    REAL_TRITON = torch.cuda.is_available()
except ImportError:
    pass
if not REAL_TRITON:
    import mini_triton as triton
    tl = triton.language
DEVICE = "cuda" if REAL_TRITON else "cpu"


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


@triton.jit
def matmul_relu_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                       stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                       BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    """
    C = ReLU(A @ B)，A: (M, K)，B: (K, N)，C: (M, N)。

    【步长】矩阵在内存里是一维长带（第 2 讲），元素 A[i, k] 的地址 = a_ptr + i*stride_am + k*stride_ak。
    把步长作为参数传进来，核函数就能处理转置过的（不连续的）矩阵。

    【指针块】
        offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        是一个 BLOCK_M × BLOCK_K 的"地址矩阵"（广播），一次 tl.load 读入整个 A 块。

    【边界】M、N、K 不一定是块大小的整数倍：
        读 A 块时，行号 >= M 或列号 >= K 的位置掩掉并填 0（填 0 不影响乘积之和）；
        写 C 块时，行号 >= M 或列号 >= N 的位置不写。
    【真实 Triton 的要求】tl.dot 要求块的每一维至少 16。
    """
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak    # (BLOCK_M, BLOCK_K)
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn    # (BLOCK_K, BLOCK_N)

    acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)                           # 累加器在寄存器里
    for k in range(0, K, BLOCK_K):
        a_mask = (offs_m[:, None] < M) & (offs_k[None, :] + k < K)
        b_mask = (offs_k[:, None] + k < K) & (offs_n[None, :] < N)
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak            # 沿 K 方向前进到下一块
        b_ptrs += BLOCK_K * stride_bk

    acc = tl.maximum(acc, 0.0)                   # 融合的 ReLU：在寄存器里做，不经过显存

    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc, mask=c_mask)


def triton_matmul_relu(a, b, block_m=64, block_n=64, block_k=32):
    """
    计算 ReLU(a @ b) 的包装函数。

    【参数】a : (M, K)，b : (K, N)，float32；块大小需是 2 的幂且 ≥ 16
    【返回】(M, N)
    """
    M, K = a.shape
    K2, N = b.shape
    assert K == K2, "内维不匹配"
    c = torch.empty(M, N, device=a.device, dtype=torch.float32)
    grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n))      # 二维网格
    matmul_relu_kernel[grid](a, b, c, M, N, K,
                             a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1),
                             BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k)
    return c


def demo_correctness():
    """演示 1：各种形状下的正确性。"""
    print_header("演示 1：分块矩阵乘法 + ReLU —— 正确性")
    print(f"  运行模式：{'真正的 Triton（GPU）' if REAL_TRITON else 'mini_triton 模拟器（CPU）'}")
    torch.manual_seed(0)
    for M, K, N in [(128, 64, 128), (100, 70, 90), (1, 33, 17), (256, 256, 256)]:
        a = torch.randn(M, K, device=DEVICE)
        b = torch.randn(K, N, device=DEVICE)
        c = triton_matmul_relu(a, b)
        ref = torch.relu(a @ b)
        print(f"  ({M}×{K}) @ ({K}×{N})：最大差异 {(c - ref).abs().max().item():.1e}")
    # 转置输入（不连续）：靠步长参数处理
    a = torch.randn(64, 96, device=DEVICE).t()                       # (96, 64)，步长 (1, 96)
    b = torch.randn(64, 80, device=DEVICE)
    c = triton_matmul_relu(a, b)
    print(f"  转置过的输入 (96×64) @ (64×80)：最大差异 {(c - torch.relu(a @ b)).abs().max().item():.1e}")


def demo_io():
    """演示 2：统计显存读取量，与第 5 讲的分块公式对照。"""
    print_header("演示 2：块越大，从显存读的越少")
    if REAL_TRITON:
        print("  （读写统计只在模拟器中可用。）")
        return
    torch.manual_seed(0)
    n = 256
    a, b = torch.randn(n, n), torch.randn(n, n)
    print(f"  n = {n}：")
    for bm, bn, bk in [(16, 16, 16), (32, 32, 32), (64, 64, 32), (128, 128, 32)]:
        triton.reset_stats()
        triton_matmul_relu(a, b, bm, bn, bk)
        s = triton.get_stats()
        # 每个输出块读 K/BLOCK_K 对小块：A 块 bm*bk、B 块 bk*bn；输出块共 (n/bm)*(n/bn) 个
        formula = (n // bm) * (n // bn) * (n // bk) * (bm * bk + bk * bn)
        print(f"    块 {bm:>3}×{bn:<3}（K 方向 {bk}）：{s['programs']:>3} 个程序，读 {s['loaded']:>9,} 个元素"
              f"（公式 {formula:,}），写 {s['stored']:,}")
    print("""
  输出块越大，A、B 中的每个元素被重复使用的次数越多，从显存读取的总量越少（约为 n³(1/BLOCK_N + 1/BLOCK_M)）。
  写入量始终是 n² —— ReLU 融合在核函数内部，没有额外的读写。
  如果不融合：先写 C（n²），ReLU 核函数再读 C（n²）、写回（n²），多出 2n² 的搬运和一次核函数启动。""")


def main():
    demo_correctness()
    demo_io()
    print("\n全部演示运行完毕。")


if __name__ == "__main__":
    main()