"""
=============================================================================
第 6 讲 · 代码 2：第一个 Triton 核函数 —— 向量加法与 GeLU
=============================================================================

对应讲义：lectures/lecture06_kernels_triton.md 第 4、5 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 5 讲说过：逐元素运算是内存受限的，把多个运算融合成一个核函数能大幅加速。
torch.compile 可以自动做简单的融合，但要真正掌控 GPU，我们需要自己写核函数。

**Triton** 是 OpenAI 开发的 GPU 编程语言（嵌在 Python 里）。和 CUDA 相比：
    - CUDA：你描述**每一个线程**做什么；要自己管理线程下标、共享内存、同步……
    - Triton：你描述**每一个线程块（程序）**做什么，操作的是"一块数据"；
              线程如何划分、如何合并访存、如何用共享内存，大部分交给编译器。
    Triton 更简单，对大多数工作负载也足够强大。

这个文件写两个最简单的核函数：
    1. 向量加法：讲清 Triton 的基本结构 —— program_id、arange、mask、load、store；
    2. GeLU：把 9 个逐元素运算融合成一个核函数，只读一次、写一次。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_triton_elementwise.py

- 有 NVIDIA GPU 并安装了 triton（pip install triton）：用真正的 Triton 编译、在 GPU 上运行，并测速；
- 否则：用同目录的 mini_triton.py 在 CPU 上模拟执行同一份核函数代码，验证正确性、统计读写量（不测速）。
"""

import math

import torch

# ---------------------------------------------------------------------------
# 选择"真正的 Triton"还是"模拟器"。核函数代码对两者完全相同。
# ---------------------------------------------------------------------------
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 部分：向量加法
# =============================================================================
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    """
    out[i] = x[i] + y[i]。每个程序实例（线程块）负责连续的 BLOCK_SIZE 个元素。

    【参数】
        x_ptr, y_ptr, out_ptr : 指向三个张量第一个元素的指针（传入张量时 Triton 自动转换）
        n_elements            : 元素总数
        BLOCK_SIZE            : 每个程序处理多少个元素。tl.constexpr 表示"编译期常量"：
                                Triton 会为每个不同的取值编译一个专门的版本（因此必须是 2 的幂等固定值）

    【执行过程】（以 n = 1000、BLOCK_SIZE = 256 为例，网格有 4 个程序）
        pid     = tl.program_id(0)             第几个程序：0、1、2、3
        offsets = pid * 256 + [0, 1, …, 255]   这个程序负责的元素下标
        mask    = offsets < 1000               最后一个程序的后 24 个位置超出范围，掩掉
        读 x、y 的这一块 → 相加 → 写回 out 的这一块
    """
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


def triton_add(x, y, block_size=256):
    """
    调用 add_kernel 的 Python 包装函数（"主机端"代码：在 CPU 上运行，负责分配输出、计算网格、启动核函数）。

    【参数】x, y : 形状相同的一维连续张量（真实 Triton 下需在 GPU 上）
    【返回】x + y
    """
    assert x.shape == y.shape and x.is_contiguous() and y.is_contiguous()
    out = torch.empty_like(x)
    n = x.numel()
    grid = (triton.cdiv(n, block_size),)            # 需要多少个程序：向上取整
    add_kernel[grid](x, y, out, n, BLOCK_SIZE=block_size)
    return out


# =============================================================================
# 第 2 部分：GeLU（融合）
# =============================================================================
@triton.jit
def gelu_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    """
    GeLU（tanh 近似）：y = 0.5 x (1 + tanh(a))，a = √(2/π) (x + 0.044715 x³)。

    【一个小技巧】tanh(a) = 2σ(2a) - 1，代入得 0.5 (1 + tanh(a)) = σ(2a)，所以
        y = x · σ(2a)
    用 sigmoid 代替 tanh：结果完全相同，而且 tl.sigmoid 在所有 Triton 版本里都可用。

    【融合】整个计算只从显存读一次 x、写一次 y；中间结果（x³、a、σ(2a)……）都在寄存器里。
    """
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    a = 0.7978845608028654 * (x + 0.044715 * x * x * x)
    y = x * tl.sigmoid(2.0 * a)
    tl.store(y_ptr + offsets, y, mask=mask)


def triton_gelu(x, block_size=1024):
    """GeLU 的包装函数。x 可以是任意形状的连续张量（按一维处理）。"""
    assert x.is_contiguous()
    y = torch.empty_like(x)
    n = x.numel()
    gelu_kernel[(triton.cdiv(n, block_size),)](x, y, n, BLOCK_SIZE=block_size)
    return y


def manual_gelu(x):
    """用 PyTorch 基本运算手写的 GeLU：每个运算都是一个单独的核函数。"""
    return 0.5 * x * (1 + torch.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x)))


# =============================================================================
# 演示
# =============================================================================
def demo_add():
    """演示 1：向量加法 —— 程序编号与掩码。"""
    print_header("演示 1：向量加法")
    print(f"  运行模式：{'真正的 Triton（GPU）' if REAL_TRITON else 'mini_triton 模拟器（CPU）'}")
    for n in [1000, 1024, 1]:
        x = torch.randn(n, device=DEVICE)
        y = torch.randn(n, device=DEVICE)
        if not REAL_TRITON:
            triton.reset_stats()
        out = triton_add(x, y)
        ok = torch.allclose(out, x + y)
        line = f"  n = {n:>5}：网格 {triton.cdiv(n, 256)} 个程序，结果正确 = {ok}"
        if not REAL_TRITON:
            s = triton.get_stats()
            line += f"；实际读取 {s['loaded']} 个元素、写入 {s['stored']} 个（被掩掉的位置不读不写）"
        print(line)
    print("  n = 1000 时第 4 个程序的后 24 个位置被 mask 掩掉；n = 1 时一个程序里只有 1 个位置有效。")


def demo_gelu():
    """演示 2：GeLU —— 融合的效果。"""
    print_header("演示 2：GeLU 核函数（把 9 个逐元素运算融合成 1 个）")
    torch.manual_seed(0)
    n = 1_000_003                                     # 故意不是块大小的整数倍
    x = torch.randn(n, device=DEVICE)
    if not REAL_TRITON:
        triton.reset_stats()
    y = triton_gelu(x)
    ref = torch.nn.functional.gelu(x, approximate="tanh")
    print(f"  与 F.gelu(approximate='tanh') 的最大差异：{(y - ref).abs().max().item():.1e}")
    print(f"  与手写的 PyTorch 版本的最大差异：{(y - manual_gelu(x)).abs().max().item():.1e}")
    if not REAL_TRITON:
        s = triton.get_stats()
        print(f"  Triton 核函数：{s['programs']} 个程序，读 {s['loaded']:,} 个元素（= n），写 {s['stored']:,} 个（= n）")
    print("""  对比手写的 PyTorch 版本：0.5*x、x*x、*x、*0.044715、+x、*0.798、tanh、+1、最后一次乘法，共 9 个运算，
  每个运算都是一个单独的核函数，每次都要读写整个张量，总搬运量约是融合版本的 9 倍以上。""")

    if REAL_TRITON:
        import time

        def bench(fn, reps=100):
            for _ in range(5):
                fn()
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            for _ in range(reps):
                fn()
            torch.cuda.synchronize()
            return (time.perf_counter() - t0) / reps

        x = torch.randn(1 << 24, device="cuda")
        for name, fn in [("手写 PyTorch", lambda: manual_gelu(x)),
                         ("F.gelu（官方融合）", lambda: torch.nn.functional.gelu(x, approximate="tanh")),
                         ("我们的 Triton 核函数", lambda: triton_gelu(x))]:
            print(f"    {name:<20} {bench(fn) * 1e3:.3f} ms")
    else:
        print("  （模拟器逐个程序顺序执行，速度没有意义，所以不测速。在 NVIDIA GPU 上运行本文件会自动测速。）")


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


if __name__ == "__main__":
    main()