"""
=============================================================================
CS349E · 前置讲 P2 · 03_shared_memory_sync_reduction.py —— 共享内存、bank、同步与并行归约
=============================================================================

【这个脚本要回答的问题】
    共享内存（shared memory）是每个 SM 上一块程序员自己管理的 SRAM，一个线程块里的线程可以通过它交换数据。
    用好它要懂三件事：bank 冲突、同步、以及"怎样让很多线程合作完成一件事"（归约）。

      演示 1：bank 冲突——32 个线程以不同的跨度访问共享内存，各要几轮？
      演示 2：同步——用共享内存把一个块里的数组倒过来；漏写 sync 会被模拟器抓到。
      演示 3：并行归约（求和）——Mark Harris《Optimizing Parallel Reduction in CUDA》（NVIDIA, 2007）的几个版本：
              分支发散、bank 冲突、顺序寻址、加载时先加一次、warp shuffle。看每一步优化消掉了什么。

【怎么读输出】全部是模拟（真实计算 + 事件计数），每次运行完全相同。

【运行】python3 03_shared_memory_sync_reduction.py          （约 3 秒）
        python3 03_shared_memory_sync_reduction.py --quick  （不到 1 秒，数组更小）
【依赖】numpy；本目录的 simt_lab；../lecture01 的 console。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from console import print_header, print_observations, print_table, fmt_time  # noqa: E402
from simt_lab import launch, estimate_time  # noqa: E402

QUICK = "--quick" in sys.argv


# ---------------------------------------------------------------------------
# 演示 1：bank 冲突
# ---------------------------------------------------------------------------
def bank_kernel(blk, fn, size):
    """一个 warp 按 fn(tid) 给出的下标读一次共享内存（只为统计 bank 冲突）。"""
    sh = blk.shared(size)
    blk.sload(sh, fn(blk.tid) % size)


def demo1_banks():
    """
    共享内存被分成 32 个 bank，每个 bank 每个周期能服务一个 4 字节的字。第 k 个字在第 k mod 32 个 bank。
    一个 warp 的 32 个线程如果访问 32 个不同的 bank（或者同一个字），一轮完成；
    如果有 m 个线程访问同一个 bank 里的**不同**的字，就要分 m 轮——这叫 m 路 bank 冲突。
    """
    print_header("演示 1：bank 冲突——一个 warp 的 32 个线程访问共享内存，要几轮？")
    cases = [
        ("连续：tid", lambda t: t),
        ("跨 2：2·tid", lambda t: 2 * t),
        ("跨 4：4·tid", lambda t: 4 * t),
        ("跨 8：8·tid", lambda t: 8 * t),
        ("跨 16：16·tid", lambda t: 16 * t),
        ("跨 32：32·tid（按列读 32×32 的 tile）", lambda t: 32 * t),
        ("跨 33：33·tid（按列读填充成 32×33 的 tile）", lambda t: 33 * t),
        ("跨 3：3·tid（奇数跨度）", lambda t: 3 * t),
        ("广播：所有线程读同一个字", lambda t: 0 * t),
        ("两个线程一组读同一个字：tid // 2", lambda t: t // 2),
    ]
    rows = []
    for name, fn in cases:
        st = launch(bank_kernel, grid=(1,), block=(32,), args=(fn, 32 * 33))
        rows.append([name, st.smem_wavefronts])
    print_table(["访问模式", "需要的轮数"], rows)
    print_observations([
        "跨度为 2^k 时冲突 2^k 路（最多 32 路）：跨 32 时 32 个线程全挤在同一个 bank，要 32 轮——共享内存的速度降到 1/32。",
        "跨度是奇数时（3、33）没有冲突：奇数和 32 互质，32 个线程正好落到 32 个不同的 bank。这就是'填充一列'有效的原因。",
        "读同一个字不算冲突（广播）：硬件只读一次、分给所有线程。",
        "这和 P1 §2.5 的 cache 冲突缺失是同一个数学问题：地址按 2 的幂次方取模分组，2 的幂次方的跨度会让所有访问落到同一组。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：同步
# ---------------------------------------------------------------------------
def reverse_kernel(blk, d, n, do_sync):
    """
    把长度为 n 的数组倒过来（NVIDIA 博客《Using Shared Memory in CUDA C/C++》的 staticReverse 例子，单个块）：
        s[t] = d[t];  __syncthreads();  d[t] = s[n − t − 1];
    第二步每个线程读的是**另一个线程**（多半在另一个 warp）写的元素，所以中间必须同步。
    """
    s = blk.shared(n)
    t = blk.tid
    blk.sstore(s, t, blk.load(d, t))
    if do_sync:
        blk.sync()
    blk.store(d, t, blk.sload(s, n - t - 1))


def demo2_sync(n=256):
    print_header(f"演示 2：同步——用共享内存把 {n} 个数倒过来（一个块，{n} 个线程）")
    rows = []
    for do_sync in (True, False):
        d = np.arange(n, dtype=np.float32)
        st = launch(reverse_kernel, grid=(1,), block=(n,), args=(d, n, do_sync))
        rows.append(["有 sync" if do_sync else "漏写 sync", st.barriers, st.race_count,
                     "对" if np.array_equal(d, np.arange(n, dtype=np.float32)[::-1]) else "错",
                     st.races[0] if st.races else "-"])
    print_table(["版本", "sync 次数", "检测到的竞争", "模拟器算出的结果", "第一条竞争报告"], rows)
    print_observations([
        "漏写 sync 时，模拟器的结果仍然是对的——它让所有线程严格同步地执行每一条语句。真实 GPU 不是这样："
        "8 个 warp 各自推进，某个 warp 可能在别的 warp 还没把数据写进共享内存时就去读，读到旧值或垃圾。",
        "所以模拟器另外记录'谁在哪个同步阶段写过、读过'，在你读到别的 warp 本阶段写的数据时报告竞争。"
        "这类 bug 最可怕的地方是：在真实 GPU 上它可能 99 次都'碰巧'正确，第 100 次出错。（NVIDIA 的 compute-sanitizer 有 racecheck 工具。）",
        "sync 只同步**一个块内**的线程。不同块之间没有这种同步——块之间要交换数据，通常只能结束这个内核、再启动下一个。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：并行归约
# ---------------------------------------------------------------------------
def red_interleaved_divergent(blk, inp, out):
    """
    版本 1（交错寻址，分支发散）：第 s 步只有 tid 是 2s 的倍数的线程干活：sdata[tid] += sdata[tid + s]。
    【问题】if (tid % (2s) == 0) 让同一个 warp 里一半、四分之三……的线程闲着，但 warp 照样要执行这条路——分支发散。
    """
    B = blk.bdx
    sdata = blk.shared(B)
    i = blk.bx * B + blk.tid
    blk.sstore(sdata, blk.tid, blk.load(inp, i))
    blk.sync()
    s = 1
    while s < B:
        t = blk.branch(blk.tid % (2 * s) == 0)
        a = blk.sload(sdata, blk.tid, mask=t)
        b = blk.sload(sdata, blk.tid + s, mask=t)
        blk.compute(1, mask=t)
        blk.sstore(sdata, blk.tid, a + b, mask=t)
        blk.sync()
        s *= 2
    first = blk.tid == 0
    blk.store(out, blk.bx, blk.sload(sdata, 0, mask=first), mask=first)


def red_interleaved_strided(blk, inp, out):
    """
    版本 2（交错寻址，按连续的线程分配工作）：第 s 步，线程 tid 处理 index = 2·s·tid。
    【好处】干活的线程是 tid = 0, 1, 2, … 连续的一段，整 warp 要么全干、要么全不干，没有发散。
    【新问题】index = 2·s·tid 的跨度是 2s：s ≥ 1 时出现 bank 冲突（跨度 2、4、8……）。
    """
    B = blk.bdx
    sdata = blk.shared(B)
    i = blk.bx * B + blk.tid
    blk.sstore(sdata, blk.tid, blk.load(inp, i))
    blk.sync()
    s = 1
    while s < B:
        idx = 2 * s * blk.tid
        t = blk.branch(idx < B)
        a = blk.sload(sdata, np.where(t, idx, 0), mask=t)
        b = blk.sload(sdata, np.where(t, idx + s, 0), mask=t)
        blk.compute(1, mask=t)
        blk.sstore(sdata, np.where(t, idx, 0), a + b, mask=t)
        blk.sync()
        s *= 2
    first = blk.tid == 0
    blk.store(out, blk.bx, blk.sload(sdata, 0, mask=first), mask=first)


def _sequential_tail(blk, sdata, B, stop_at=1):
    """顺序寻址的循环：s = B/2, B/4, …, stop_at：前 s 个线程做 sdata[tid] += sdata[tid + s]。"""
    s = B // 2
    while s >= stop_at:
        t = blk.branch(blk.tid < s)
        a = blk.sload(sdata, blk.tid, mask=t)
        b = blk.sload(sdata, np.where(t, blk.tid + s, 0), mask=t)
        blk.compute(1, mask=t)
        blk.sstore(sdata, blk.tid, a + b, mask=t)
        blk.sync()
        s //= 2


def red_sequential(blk, inp, out):
    """
    版本 3（顺序寻址）：第 s 步前 s 个线程做 sdata[tid] += sdata[tid + s]，s 从 B/2 减半到 1。
    【好处】干活的线程连续（没有发散，直到 s < 32）；访问 sdata[tid] 和 sdata[tid + s] 都是连续的（没有 bank 冲突）。
    【剩下的浪费】第一步就有一半线程闲着：加载时它们只搬了一个数。
    """
    B = blk.bdx
    sdata = blk.shared(B)
    i = blk.bx * B + blk.tid
    blk.sstore(sdata, blk.tid, blk.load(inp, i))
    blk.sync()
    _sequential_tail(blk, sdata, B)
    first = blk.tid == 0
    blk.store(out, blk.bx, blk.sload(sdata, 0, mask=first), mask=first)


def red_first_add(blk, inp, out):
    """
    版本 4（加载时先加一次）：每个块负责 2B 个元素，每个线程读两个数先加起来再放进共享内存。
    【好处】块数减半；第一步"一半线程闲着"的浪费没有了。
    """
    B = blk.bdx
    sdata = blk.shared(B)
    i = blk.bx * (2 * B) + blk.tid
    v = blk.load(inp, i) + blk.load(inp, i + B)
    blk.compute(1)
    blk.sstore(sdata, blk.tid, v)
    blk.sync()
    _sequential_tail(blk, sdata, B)
    first = blk.tid == 0
    blk.store(out, blk.bx, blk.sload(sdata, 0, mask=first), mask=first)


def red_warp_shuffle(blk, inp, out):
    """
    版本 5（最后 32 个数用 warp shuffle）：共享内存归约到剩 64 个数后，第 0 个 warp 先把 64 个加成 32 个，
    再用 5 次 shfl_down（16、8、4、2、1）在寄存器之间直接交换，不用共享内存、不用 sync。
    【为什么能省 sync】同一个 warp 的线程一起执行 shfl 指令，交换由硬件保证完成（现代 CUDA 用 __shfl_down_sync）。
    """
    B = blk.bdx
    sdata = blk.shared(B)
    i = blk.bx * (2 * B) + blk.tid
    v = blk.load(inp, i) + blk.load(inp, i + B)
    blk.compute(1)
    blk.sstore(sdata, blk.tid, v)
    blk.sync()
    _sequential_tail(blk, sdata, B, stop_at=64)
    w0 = blk.tid < 32
    val = blk.sload(sdata, np.where(w0, blk.tid, 0), mask=w0) + blk.sload(sdata, np.where(w0, blk.tid + 32, 0), mask=w0)
    blk.compute(1, mask=w0)
    for off in (16, 8, 4, 2, 1):
        val = val + blk.shfl_down(val, off, mask=w0)
        blk.compute(1, mask=w0)
    first = blk.tid == 0
    blk.store(out, blk.bx, val, mask=first)


def demo3_reduction(n=1 << (14 if QUICK else 18), B=256):
    print_header(f"演示 3：并行归约——对 {n:,} 个 float32 求和（每块 {B} 线程，先得到每块的部分和）")
    inp = (np.arange(n) % 7).astype(np.float32)          # 小整数：float32 求和是精确的，方便核对
    expect = float(inp.astype(np.float64).sum())
    versions = [
        ("1 交错寻址（分支发散）", red_interleaved_divergent, 1),
        ("2 交错寻址（连续线程，bank 冲突）", red_interleaved_strided, 1),
        ("3 顺序寻址", red_sequential, 1),
        ("4 加载时先加一次", red_first_add, 2),
        ("5 + 最后一个 warp 用 shuffle", red_warp_shuffle, 2),
    ]
    rows = []
    base_instr = None
    t_mem = None
    for name, fn, per_thread in versions:
        n_blocks = n // (B * per_thread)
        out = np.zeros(n_blocks, dtype=np.float32)
        st = launch(fn, grid=(n_blocks,), block=(B,), args=(inp, out))
        est = estimate_time(st)
        base_instr = base_instr or st.warp_instr
        t_mem = est["t_mem"]
        ok = float(out.astype(np.float64).sum()) == expect
        rows.append([name, n_blocks, f"{st.warp_instr:,}", f"{st.warp_instr / base_instr:.2f}", f"{st.simt_efficiency() * 100:.0f}%",
                     f"{st.divergent:,}", f"{st.bank_conflict_factor():.2f}", f"{st.barriers:,}",
                     fmt_time(est["t_issue"]), fmt_time(est["t_mem"]), "对" if ok else "错"])
    print_table(["版本", "块数", "warp 指令", "相对版本 1", "SIMT 效率", "发散的 warp 次数", "bank 每指令轮数",
                 "sync 次数", "估计发射时间", "估计 HBM 时间", "结果"], rows)
    print_observations([
        "版本 1 的 SIMT 效率很低、发散极多：tid % (2s) == 0 让每个 warp 里大部分线程闲着，但 warp 照样占用发射名额。",
        "版本 2 改成'让连续的线程干活'，发散几乎消失、指令数大减；但 2·s·tid 的跨度造成了 bank 冲突（每条共享内存指令平均要好几轮）。",
        "版本 3 顺序寻址同时解决了两个问题：干活的线程连续（无发散，直到只剩不到 32 个线程），访问的地址也连续（无 bank 冲突）。",
        "版本 4 让每个线程加载时先加一次，块数减半，总指令数又少了一截——'第一步就有一半线程闲着'的浪费没有了。",
        "版本 5 把最后几步挪到 warp 内的 shuffle：省掉了这几步的共享内存读写和 sync。这正是现代 CUDA 归约的标准写法"
        "（Harris 2007 年的原版用的是'最后一个 warp 不 sync'的技巧，依赖 warp 内线程严格同步；Volta 之后线程可以在 warp 内独立推进，"
        "那个技巧不再安全，要用 __syncwarp 或 shuffle）。",
        f"归约的强度极低（每读 4 字节只做 1 次加法），理想情况下是 memory-bound：读完 {n * 4 / 1e6:.1f} MB 输入至少要 "
        f"{fmt_time(t_mem)}（估计 HBM 时间一列）。版本 1 的指令发射时间比这还长——它受限于指令，而不是带宽；"
        "后面的版本把指令和同步的开销压到读数据的时间以下，内核才真正受限于带宽。Harris 的原文也是以'接近峰值带宽'为优化目标的。",
    ])


if __name__ == "__main__":
    demo1_banks()
    demo2_sync()
    demo3_reduction()