"""
=============================================================================
第 6 讲 · 代码 5：融合的 RMSNorm —— 前向、反向、接入自动求导（Triton）
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 3 讲说过：RMSNorm 的 FLOPs 很少，却因为内存受限而占用可观的时间。
用 PyTorch 基本运算写的 RMSNorm（x.pow(2).mean().rsqrt() * x * g）会启动好几个核函数，
每个都要读写整个激活张量；反向传播还要更多。
把它写成一个融合的 Triton 核函数，就能做到"前向只读 x 一次、写 y 一次"。
（2026 年版 CS336 的作业 2 就包含"用 Triton 实现融合 RMSNorm 核函数"。）

这个文件包括：
    1. 前向核函数：每个程序处理一行（与 03 号文件的 softmax 同一个模式）；
    2. 反向核函数：推导梯度公式，每个程序处理一行；增益 g 的梯度要对所有行求和，
       为了避免多个程序同时写同一个位置（数据竞争），先让每行写一份"部分梯度"，再在外面求和；
    3. 用 torch.autograd.Function 把两者包装起来，像普通 PyTorch 模块一样使用 .backward()；
    4. 与 PyTorch 自动求导的结果逐一对比。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 05_triton_rmsnorm.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)


# =============================================================================
# 第 1 部分：参照实现
# =============================================================================
def rmsnorm_reference(x, g, eps=1e-5):
    """RMSNorm 的 PyTorch 参照实现：y = x / √(mean(x²) + eps) ⊙ g。自动求导用它来对照。"""
    return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * g


# =============================================================================
# 第 2 部分：前向核函数
# =============================================================================
@triton.jit
def rmsnorm_fwd_kernel(x_ptr, g_ptr, y_ptr, r_ptr, n_cols, row_stride, eps, BLOCK_SIZE: tl.constexpr):
    """
    每个程序处理一行：
        r = 1 / √(mean(x²) + eps)
        y = x · r · g
    额外把每行的 r 存下来（每行 1 个数，很便宜），反向时直接用，不必重算。

    【为什么在 FP32 中计算】见第 3 讲 3.2 节：平方求和是累加，低精度会出错。
    （真实场景中 x 可能是 BF16，这里 .to(tl.float32) 把它转成 FP32 再算。）
    """
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(x_ptr + row * row_stride + cols, mask=mask, other=0.0).to(tl.float32)
    g = tl.load(g_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    ms = tl.sum(x * x, axis=0) / n_cols          # 被掩掉的位置填 0，不影响平方和
    r = tl.rsqrt(ms + eps)
    tl.store(y_ptr + row * row_stride + cols, x * r * g, mask=mask)
    tl.store(r_ptr + row, r)


# =============================================================================
# 第 3 部分：反向核函数
# =============================================================================
@triton.jit
def rmsnorm_bwd_kernel(x_ptr, g_ptr, dy_ptr, r_ptr, dx_ptr, dg_partial_ptr, n_cols, row_stride,
                       BLOCK_SIZE: tl.constexpr):
    """
    RMSNorm 的反向传播（每个程序处理一行）。

    【推导】对一行：y_i = g_i · x_i · r，r = (mean(x²) + eps)^(-1/2)
        ∂r/∂x_j = -r³ · x_j / N
        dx_j = Σ_i dy_i · ∂y_i/∂x_j = dy_j g_j r  -  (r³ x_j / N) · Σ_i dy_i g_i x_i
        dg_i = Σ_{所有行} dy_i · x_i · r
    所以每行需要：x、g、dy 和前向存下的 r；输出 dx 的这一行，以及"这一行对 dg 的贡献"。

    【数据竞争】dg 是对所有行求和。如果每个程序都直接往同一个 dg 数组里加，
    不同程序（在 GPU 上并行运行）会同时读-改-写同一个位置，结果不确定。
    解决办法之一：每行写到部分梯度缓冲区的一行（dg_partial[row, :]），最后在外面 sum(dim=0)。
    （另一种办法是原子加 tl.atomic_add，但它慢而且结果不确定。）
    """
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(x_ptr + row * row_stride + cols, mask=mask, other=0.0).to(tl.float32)
    dy = tl.load(dy_ptr + row * row_stride + cols, mask=mask, other=0.0).to(tl.float32)
    g = tl.load(g_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    r = tl.load(r_ptr + row)
    dyg = dy * g
    dot = tl.sum(dyg * x, axis=0)                        # Σ_i dy_i g_i x_i
    dx = dyg * r - x * (r * r * r * dot / n_cols)
    tl.store(dx_ptr + row * row_stride + cols, dx, mask=mask)
    tl.store(dg_partial_ptr + row * n_cols + cols, dy * x * r, mask=mask)


# =============================================================================
# 第 4 部分：接入自动求导
# =============================================================================
class TritonRMSNorm(torch.autograd.Function):
    """
    用 Triton 核函数实现的 RMSNorm，支持 .backward()。

    【形状】x : (..., N)，最后一维连续；g : (N,)
    【保存给反向的东西】x、g、每行的 r（N 很大时，r 只占 1/N，几乎不花显存）
    """

    @staticmethod
    def forward(ctx, x, g, eps=1e-5):
        shape = x.shape
        x2 = x.reshape(-1, shape[-1]).contiguous()
        n_rows, n_cols = x2.shape
        y = torch.empty_like(x2)
        r = torch.empty(n_rows, device=x.device, dtype=torch.float32)
        block = triton.next_power_of_2(n_cols)
        rmsnorm_fwd_kernel[(n_rows,)](x2, g, y, r, n_cols, x2.stride(0), eps, BLOCK_SIZE=block)
        ctx.save_for_backward(x2, g, r)
        ctx.shape = shape
        return y.reshape(shape)

    @staticmethod
    def backward(ctx, dy):
        x2, g, r = ctx.saved_tensors
        n_rows, n_cols = x2.shape
        dy2 = dy.reshape(-1, n_cols).contiguous()
        dx = torch.empty_like(x2)
        dg_partial = torch.empty(n_rows, n_cols, device=x2.device, dtype=torch.float32)
        block = triton.next_power_of_2(n_cols)
        rmsnorm_bwd_kernel[(n_rows,)](x2, g, dy2, r, dx, dg_partial, n_cols, x2.stride(0), BLOCK_SIZE=block)
        return dx.reshape(ctx.shape), dg_partial.sum(dim=0), None      # eps 不需要梯度


def demo_forward_backward():
    """演示 1：与 PyTorch 自动求导逐一对比。"""
    print_header("演示 1：融合 RMSNorm 的前向与反向")
    print(f"  运行模式：{'真正的 Triton（GPU）' if REAL_TRITON else 'mini_triton 模拟器（CPU）'}")
    torch.manual_seed(0)
    for shape in [(4, 16, 64), (37, 100), (2, 3, 513)]:
        n = shape[-1]
        x = torch.randn(*shape, device=DEVICE, requires_grad=True)
        g = torch.randn(n, device=DEVICE, requires_grad=True)
        x_ref = x.detach().clone().requires_grad_(True)
        g_ref = g.detach().clone().requires_grad_(True)

        y = TritonRMSNorm.apply(x, g)
        y_ref = rmsnorm_reference(x_ref, g_ref)
        dy = torch.randn_like(y)
        y.backward(dy)
        y_ref.backward(dy)
        print(f"  x 形状 {shape}：y 差异 {(y - y_ref).abs().max().item():.1e}，"
              f"dx 差异 {(x.grad - x_ref.grad).abs().max().item():.1e}，"
              f"dg 差异 {(g.grad - g_ref.grad).abs().max().item():.1e}")


def demo_io():
    """演示 2：读写量对比。"""
    print_header("演示 2：融合前后的显存读写")
    if REAL_TRITON:
        print("  （读写统计只在模拟器中可用。）")
        return
    torch.manual_seed(0)
    M, N = 256, 1024
    x = torch.randn(M, N, requires_grad=True)
    g = torch.randn(N, requires_grad=True)
    triton.reset_stats()
    y = TritonRMSNorm.apply(x, g)
    fwd = triton.get_stats()
    triton.reset_stats()
    y.backward(torch.randn_like(y))
    bwd = triton.get_stats()
    MN = M * N
    print(f"  M = {M} 行，N = {N} 列（MN = {MN:,}）")
    print(f"  前向：读 {fwd['loaded']:,} ≈ MN + M·N（每行都读一遍 g），写 {fwd['stored']:,} = MN + M")
    print(f"  反向：读 {bwd['loaded']:,} ≈ 2MN + M·N + M，写 {bwd['stored']:,} = 2MN（dx 和部分 dg）")
    print("""
  注意 g 被每一行都读了一遍（统计里的 M·N）。在真实 GPU 上，g 只有 N 个数，
  读第一遍之后就会待在 L2 缓存里，后面的读取几乎不花显存带宽。
  对比 PyTorch 基本运算写的版本：前向就有平方、求均值、加 eps、rsqrt、乘 x、乘 g 等多个核函数，
  每个都要读写 MN 量级的数据，反向更多。这就是融合核函数能快好几倍的原因。
  【可以改进的地方】部分梯度缓冲区有 MN 个数，本身也要写和读一遍。更好的做法是让每个程序
  处理多行、先在寄存器里累加这些行的 dg，再写出 —— 缓冲区就只有 (M / 每程序行数) × N 了。""")


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


if __name__ == "__main__":
    main()