"""
simt.py —— 用 Python 写的 SIMT GPU 模拟器：按 CUDA 的方式执行内核，并把"为什么快 / 为什么慢"数出来
=====================================================================
（第 2 讲引入；第 3、4 讲大量使用；第 10、11、21 讲写融合内核时还会用）

【这个文件要回答的问题】
    没有 NVIDIA GPU，怎样学 CUDA？我们把 CUDA 的执行模型用 Python 搭一遍：
        - 内核（kernel）是一个普通的 Python 函数，第一个参数 t 提供 CUDA 的内置变量
          t.threadIdx、t.blockIdx、t.blockDim、t.gridDim（都有 .x .y .z）；
        - launch(kernel, grid, block, 参数...) 相当于 CUDA 的 kernel<<<grid, block>>>(参数...)：
          模拟器为网格里的每一个线程调用一次这个函数；
        - 显存是 Device 上分配的 DeviceArray（相当于 cudaMalloc 得到的设备指针），
          内核里对它的每一次读（A[i]）和写（A[i] = v）都会被记录下来；
        - 共享内存用 t.shared(...) 申请；__syncthreads() 写成 `yield t.syncthreads()`。
    模拟器不追求快（它比真 GPU 慢上百万倍），追求的是"能数"：
        每个线程束（warp）的每条访存指令产生了几个 32 字节的内存事务（合并访存）、
        共享内存有没有 bank 冲突、线程束里有多少线程在空转（分化）、有没有数据竞争。

【本文件提供】
    常量         : WARP_SIZE、SECTOR_BYTES、NUM_BANKS、BANK_BYTES、MAX_THREADS_PER_BLOCK 等
    Dim3 / dim3  : 三维的网格 / 线程块大小（对应 CUDA 的 dim3）
    Device       : 一块模拟的 GPU：malloc / memcpy / free（即 cudaMalloc / cudaMemcpy / cudaFree）、拷贝记账、泄漏检查
    DeviceArray  : 显存里的一维数组（相当于设备指针）：A[i] 读、A[i] = v 写、A.atomic_add(i, v)、A.atomic_max(i, v)
    launch       : 发射内核，返回 LaunchStats（线程数、访存请求、内存事务、bank 冲突、指令发射、竞争……）
    异常         : SimError 及其子类 IllegalAddress（越界）、LaunchError（配置非法）、BarrierError（栅栏用错）、MemcpyError

【模拟的规则（以及它和真 GPU 的差别）】
    1. 线程束：块内线程按 CUDA 的线性编号 x + y·Bx + z·Bx·By 排队，每 32 个一组成为一个线程束。
    2. 执行顺序：一个块一个块地执行；块内每个线程依次运行，直到结束或遇到 __syncthreads()；
       块内全部线程都到达栅栏后，再一起进入下一段（"一个栅栏周期"）。
       真 GPU 上线程是并发的、执行顺序不确定——所以"在模拟器上结果对"不能证明内核没有数据竞争，
       这就是模拟器另外做竞争检测的原因（规则 6）。
    3. 合并访存：同一个线程束里，执行"同一条访存指令"（同一行代码对同一个数组的第 k 次访问）的线程放在一起看，
       它们访问的地址落在多少个不同的 32 字节段（sector）里，这条指令就产生多少个内存事务。
       32 个线程各读相邻的 4 字节 → 128 字节 → 4 个事务（最理想）；各隔很远 → 32 个事务。
    4. bank 冲突：共享内存分成 32 个 bank，每个 bank 4 字节宽，连续的 4 字节字依次落在 bank 0、1、……、31、0……
       同一条指令里，落在同一个 bank 的不同字要排队，排得最长的那个 bank 决定这条指令要走几轮（wavefront）。
       多个线程读同一个字是广播，不算冲突。（按 4 字节字计算；8 字节类型的细节这里不模拟。）
    5. 指令计数（可选，count_instructions=True）：每执行一行 Python 代码算一条指令；
       同一个线程束里执行"同一行的第 k 次"的线程，合起来算一次"发射"（issue）。
       发射次数 × 32 和"线程实际执行的指令数"之比，就是线程束分化造成的浪费。
    6. 竞争检测：同一个栅栏周期里，一个线程写了共享内存的某个位置、另一个线程又读或写了它 → 记为一次竞争。
       全局内存的竞争检测较慢，默认关闭（check_global_races=True 打开）。原子操作不参与竞争检测。
    7. 越界一律报错（IllegalAddress）。真 GPU 上越界不一定报错（可能悄悄读到别的数据），模拟器更严格；
       共享内存还会按每一维分别检查下标。
    8. 共享内存不会被初始化：浮点数组里填 NaN、整数数组填 -999999，读到"没写过的值"一眼就能看出来。
    9. 主机代码不能直接读写 DeviceArray（真 CUDA 里这样做通常会段错误）：要用 Device.memcpy 拷来拷去。

【用法】
    dev = Device()
    dA = dev.to_device(np.arange(8, dtype=np.float32))      # cudaMalloc + cudaMemcpy(HostToDevice)
    dC = dev.malloc(8)                                       # cudaMalloc
    def scale(t, A, C, n):                                   # __global__ void scale(float* A, float* C, int n)
        i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
        if i < n:
            C[i] = A[i] * 2
    stats = launch(scale, grid=1, block=8, *(dA, dC, 8))     # 也可以写 launch(scale, 1, 8, dA, dC, 8)
    print(dev.to_host(dC), stats.summary())
"""

import inspect
import sys
from collections import defaultdict, namedtuple

import numpy as np

# =====================================================================
# 0. 硬件常量（NVIDIA 近几代 GPU 的通用值）
# =====================================================================

WARP_SIZE = 32                         # 一个线程束的线程数
SECTOR_BYTES = 32                      # 全局内存事务的粒度：一个 sector 32 字节（缓存行 128 字节 = 4 个 sector）
NUM_BANKS = 32                         # 共享内存的 bank 数
BANK_BYTES = 4                         # 每个 bank 每轮能提供 4 字节
MAX_THREADS_PER_BLOCK = 1024           # 每个线程块最多 1024 个线程
MAX_BLOCK_DIM = (1024, 1024, 64)       # 线程块每一维的上限
MAX_GRID_DIM = (2 ** 31 - 1, 65535, 65535)   # 网格每一维的上限
DEFAULT_SHARED_LIMIT = 48 * 1024       # 每个块静态共享内存的默认上限（48 KiB；更大需要显式申请，见第 4 讲）
DEVICE_ALIGNMENT = 256                 # cudaMalloc 返回的地址至少按 256 字节对齐

_THIS_FILE = __file__


# =====================================================================
# 1. dim3 与异常
# =====================================================================

class Dim3(namedtuple("Dim3", ["x", "y", "z"])):
    """
    三维大小 / 三维下标，对应 CUDA 的 dim3 / uint3。

    【Python 小知识】namedtuple 生成一个"带字段名的元组"：d.x 和 d[0] 是同一个值，元组不可修改。
        这里继承它只是为了加一个 count 属性。__slots__ = () 表示不额外占内存。
    """
    __slots__ = ()

    @property
    def count(self):
        """x·y·z：网格里的块数，或线程块里的线程数。"""
        return self.x * self.y * self.z


def dim3(x=1, y=1, z=1):
    """和 CUDA 的 dim3(x, y, z) 一样：没给的维度默认是 1。"""
    return Dim3(int(x), int(y), int(z))


def as_dim3(v):
    """
    把 int、(x,)、(x, y)、(x, y, z) 或 Dim3 统一成 Dim3。

    【例子】as_dim3(256) → Dim3(256, 1, 1)；as_dim3((16, 16)) → Dim3(16, 16, 1)
    【边界】长度超过 3 的元组报 ValueError。
    """
    if isinstance(v, Dim3):
        return v
    if isinstance(v, (int, np.integer)):
        return Dim3(int(v), 1, 1)
    v = tuple(int(a) for a in v)
    if not 1 <= len(v) <= 3:
        raise ValueError(f"dim3 只能有 1–3 维，收到 {v}")
    return Dim3(*(v + (1,) * (3 - len(v))))


class SimError(Exception):
    """模拟器检测到的错误的基类。"""


class IllegalAddress(SimError):
    """越界访问（对应 CUDA 的 cudaErrorIllegalAddress / "an illegal memory access was encountered"）。"""


class LaunchError(SimError):
    """发射配置非法（对应 cudaErrorInvalidConfiguration 等），例如一个块超过 1024 个线程、共享内存超限。"""


class BarrierError(SimError):
    """__syncthreads() 用错：同一个块里有的线程到了栅栏、有的已经结束。真 GPU 上这是未定义行为（常见后果是卡死）。"""


class MemcpyError(SimError):
    """cudaMemcpy 的参数不对：方向和指针类型不符、字节数超过数组大小等。"""


SYNC = object()          # __syncthreads() 的"信号"：内核里写 yield t.syncthreads()，模拟器收到它就让线程在栅栏处等


# =====================================================================
# 2. 显存：Device 与 DeviceArray
# =====================================================================

HostToDevice = "HostToDevice"           # 对应 cudaMemcpyHostToDevice
DeviceToHost = "DeviceToHost"           # 对应 cudaMemcpyDeviceToHost
DeviceToDevice = "DeviceToDevice"       # 对应 cudaMemcpyDeviceToDevice


class DeviceArray:
    """
    显存里的一段一维数组，相当于 CUDA 里 cudaMalloc 得到的设备指针（加上它的长度和类型）。

    【属性】name（名字，用于统计和报错）、dtype、n（元素个数）、base（起始字节地址）、nbytes
    【在内核里怎样用】
        x = A[i]           读第 i 个元素（i 必须是整数；越界抛 IllegalAddress）
        A[i] = v           写
        A.atomic_add(i, v) 原子地加上 v，返回旧值（对应 atomicAdd）
        A.atomic_max(i, v) 原子地取最大值，返回旧值（对应 atomicMax）
    【边界】在内核之外（主机代码里）读写会抛 SimError：主机只能用 Device.memcpy 访问显存。
        释放之后再访问也会报错（对应"使用已释放的指针"）。
    """

    __slots__ = ("device", "name", "dtype", "n", "base", "itemsize", "nbytes", "_data", "freed")

    def __init__(self, device, n, dtype, name, base):
        self.device = device
        self.name = name
        self.dtype = np.dtype(dtype)
        self.n = int(n)
        self.base = base
        self.itemsize = self.dtype.itemsize
        self.nbytes = self.n * self.itemsize
        # 显存刚分配时内容是"垃圾"；这里填一个容易认出来的值
        self._data = np.full(self.n, np.nan if self.dtype.kind == "f" else -999999, dtype=self.dtype)
        self.freed = False

    def __len__(self):
        return self.n

    def __repr__(self):
        return f"<DeviceArray {self.name}: {self.n} × {self.dtype}, 地址 0x{self.base:x}>"

    def _check(self, i):
        """检查下标：必须是整数、在 [0, n) 里、数组没有被释放、而且是在内核里访问。返回 int 下标。"""
        run = _RUN
        if run is None:
            raise SimError(f"主机代码不能直接读写显存数组 {self.name}：请用 cudaMemcpy（Device.memcpy / to_host）")
        if self.freed:
            raise IllegalAddress(f"访问已经 cudaFree 的数组 {self.name}")
        if not isinstance(i, (int, np.integer)):
            raise TypeError(f"{self.name}[...] 的下标必须是整数（C 指针只能按整数偏移），收到 {type(i).__name__}")
        i = int(i)
        if not 0 <= i < self.n:
            raise IllegalAddress(f"线程 {run.where()} 越界访问 {self.name}[{i}]（合法范围 0..{self.n - 1}）")
        return i

    def __getitem__(self, i):
        i = self._check(i)
        _RUN.access("global", self, "ld", i, self.base + i * self.itemsize, self.itemsize)
        return self._data[i]

    def __setitem__(self, i, v):
        i = self._check(i)
        _RUN.access("global", self, "st", i, self.base + i * self.itemsize, self.itemsize)
        self._data[i] = v

    def atomic_add(self, i, v):
        """原子加：old = A[i]; A[i] = old + v; 返回 old。模拟器里线程依次执行，天然是原子的；这里主要是计数。"""
        i = self._check(i)
        _RUN.access("global", self, "atom", i, self.base + i * self.itemsize, self.itemsize)
        old = self._data[i]
        self._data[i] = old + v
        return old

    def atomic_max(self, i, v):
        """原子最大值：old = A[i]; A[i] = max(old, v); 返回 old。"""
        i = self._check(i)
        _RUN.access("global", self, "atom", i, self.base + i * self.itemsize, self.itemsize)
        old = self._data[i]
        if v > old:
            self._data[i] = v
        return old


class Device:
    """
    一块模拟的 GPU：负责显存的分配与释放、主机 ↔ 设备的拷贝，以及拷贝时间的记账。

    【参数】
        name           : 名字
        pcie_bandwidth : 主机 ↔ 设备拷贝的带宽（字节/秒），默认 32e9（PCIe 4.0 x16 的标称值，原课 slides 同）
        pcie_latency   : 每次拷贝的固定开销（秒），默认 10 微秒（量级）
    【记账】copies（次数）、h2d_bytes、d2h_bytes、copy_seconds（按 α + 字节数 / 带宽 估算的拷贝时间）
    【边界】只模拟"同步拷贝"（cudaMemcpy 返回时拷贝已完成）；异步拷贝、页锁定内存见第 3 讲讲义。
    """

    def __init__(self, name="模拟 GPU", pcie_bandwidth=32e9, pcie_latency=10e-6):
        self.name = name
        self.pcie_bandwidth = pcie_bandwidth
        self.pcie_latency = pcie_latency
        self._next = 0x7000_0000_0000        # 第一个设备地址（只是一个"看起来像地址"的整数）
        self.live = {}                        # 基地址 → DeviceArray（还没释放的）
        self.allocated_bytes = 0
        self.peak_bytes = 0
        self.mallocs = 0
        self.copies = 0
        self.h2d_bytes = 0
        self.d2h_bytes = 0
        self.d2d_bytes = 0
        self.copy_seconds = 0.0
        self._names = defaultdict(int)

    # ---------------- 分配与释放 ----------------
    def malloc(self, n, dtype=np.float32, name=None):
        """
        cudaMalloc：在显存上申请 n 个 dtype 类型的元素，返回 DeviceArray（真 CUDA 通过 void** 参数把地址写回）。

        【边界】n 必须 ≥ 1；地址按 256 字节对齐（cudaMalloc 的保证），所以不同数组的起点总是对齐的。
        """
        if n < 1:
            raise ValueError("cudaMalloc 的大小必须至少是 1 个元素")
        if name is None:
            self._names["d"] += 1
            name = f"d{self._names['d']}"
        arr = DeviceArray(self, n, dtype, name, self._next)
        self._next += -(-arr.nbytes // DEVICE_ALIGNMENT) * DEVICE_ALIGNMENT    # 向上取整到 256 的倍数
        self.live[arr.base] = arr
        self.mallocs += 1
        self.allocated_bytes += arr.nbytes
        self.peak_bytes = max(self.peak_bytes, self.allocated_bytes)
        return arr

    def free(self, arr):
        """cudaFree。【边界】重复释放、释放不是本设备分配的数组 → SimError。"""
        if arr.freed or self.live.get(arr.base) is not arr:
            raise SimError(f"cudaFree 了一个无效的指针（{arr.name}：重复释放或不属于这块 GPU）")
        arr.freed = True
        del self.live[arr.base]
        self.allocated_bytes -= arr.nbytes

    def leaks(self):
        """还没有 cudaFree 的数组列表（程序结束时它们就是"显存泄漏"）。"""
        return list(self.live.values())

    # ---------------- 拷贝 ----------------
    def memcpy(self, dst, src, nbytes, kind):
        """
        cudaMemcpy(dst, src, nbytes, kind)：按字节数拷贝。

        【参数】
            dst, src : 一个是 NumPy 数组（主机内存），一个是 DeviceArray（显存），或两个都是 DeviceArray
            nbytes   : 拷贝的字节数（必须是元素大小的整数倍）
            kind     : HostToDevice / DeviceToHost / DeviceToDevice
        【边界】方向和参数类型对不上、字节数超出任一方的大小、两边元素类型不同 → MemcpyError。
            真 CUDA 里主机指针和设备指针在 C 的类型系统里都是 float*，编译器分不出来，这类错误要到运行时才暴露。
        """
        want = {HostToDevice: (DeviceArray, np.ndarray), DeviceToHost: (np.ndarray, DeviceArray),
                DeviceToDevice: (DeviceArray, DeviceArray)}
        if kind not in want:
            raise MemcpyError(f"未知的拷贝方向 {kind}")
        tdst, tsrc = want[kind]
        if not isinstance(dst, tdst) or not isinstance(src, tsrc):
            raise MemcpyError(f"cudaMemcpy 方向 {kind} 要求 dst 是 {tdst.__name__}、src 是 {tsrc.__name__}，"
                              f"实际是 {type(dst).__name__}、{type(src).__name__}")
        for a in (dst, src):
            if isinstance(a, DeviceArray) and a.freed:
                raise MemcpyError(f"拷贝已释放的数组 {a.name}")
        if dst.dtype != src.dtype:
            raise MemcpyError(f"两边的元素类型不同（{dst.dtype} vs {src.dtype}）；cudaMemcpy 只按字节搬，不做类型转换")
        item = dst.dtype.itemsize
        if nbytes % item:
            raise MemcpyError(f"字节数 {nbytes} 不是元素大小 {item} 的整数倍")
        count = nbytes // item
        for a, role in ((dst, "dst"), (src, "src")):
            size = a.n if isinstance(a, DeviceArray) else a.size
            if count > size:
                raise MemcpyError(f"拷贝 {count} 个元素，超出了 {role} 的大小 {size}")
        if isinstance(src, np.ndarray) and not src.flags["C_CONTIGUOUS"]:
            raise MemcpyError("主机数组必须是连续存放的（cudaMemcpy 只认一段连续的字节）")
        sdata = src._data if isinstance(src, DeviceArray) else src.reshape(-1)
        if isinstance(dst, DeviceArray):
            dst._data[:count] = sdata[:count]
        else:
            if not dst.flags["C_CONTIGUOUS"]:
                raise MemcpyError("主机数组必须是连续存放的")
            dst.reshape(-1)[:count] = sdata[:count]
        self.copies += 1
        if kind == HostToDevice:
            self.h2d_bytes += nbytes
        elif kind == DeviceToHost:
            self.d2h_bytes += nbytes
        else:
            self.d2d_bytes += nbytes
        if kind != DeviceToDevice:             # 设备内部的拷贝走显存带宽，不经过 PCIe，这里不计入 PCIe 时间
            self.copy_seconds += self.pcie_latency + nbytes / self.pcie_bandwidth

    def to_device(self, host, name=None):
        """便捷函数：malloc + memcpy(HostToDevice)。host 会先被展平成一维（按行优先）。"""
        host = np.ascontiguousarray(host)
        arr = self.malloc(host.size, host.dtype, name)
        self.memcpy(arr, host.reshape(-1), host.nbytes, HostToDevice)
        return arr

    def to_host(self, arr, shape=None):
        """便捷函数：新建一个主机数组 + memcpy(DeviceToHost)；shape 给出时把结果 reshape 成这个形状。"""
        out = np.empty(arr.n, dtype=arr.dtype)
        self.memcpy(out, arr, arr.nbytes, DeviceToHost)
        return out.reshape(shape) if shape is not None else out

    def copy_report(self):
        """拷贝记账的汇总字典。"""
        return {"copies": self.copies, "h2d_bytes": self.h2d_bytes, "d2h_bytes": self.d2h_bytes,
                "d2d_bytes": self.d2d_bytes, "copy_seconds": self.copy_seconds}


# =====================================================================
# 3. 共享内存
# =====================================================================

class SharedArray:
    """
    一个线程块的共享内存数组（对应 __shared__ float As[T][T];）。同一个块的所有线程看到的是同一个对象。

    【在内核里怎样用】As = t.shared("As", (T, T))；As[ty, tx] = v；x = As[k]（一维时）
    【边界】下标必须是整数或整数元组，个数等于维数，每一维都在范围内（比 C 更严格：C 只要不出整块内存就不报错）。
    """

    __slots__ = ("name", "shape", "dtype", "itemsize", "offset", "_data", "_strides", "n")

    def __init__(self, name, shape, dtype, offset):
        self.name = name
        self.shape = tuple(int(s) for s in shape)
        self.dtype = np.dtype(dtype)
        self.itemsize = self.dtype.itemsize
        self.offset = offset                       # 在本块共享内存里的起始字节偏移
        self.n = int(np.prod(self.shape))
        self._data = np.full(self.n, np.nan if self.dtype.kind == "f" else -999999, dtype=self.dtype)
        strides, acc = [], 1
        for s in reversed(self.shape):
            strides.append(acc)
            acc *= s
        self._strides = tuple(reversed(strides))

    @property
    def nbytes(self):
        return self.n * self.itemsize

    def _flat(self, idx):
        """多维下标 → 一维位置（行优先），同时检查每一维的范围。"""
        if not isinstance(idx, tuple):
            idx = (idx,)
        if len(idx) != len(self.shape):
            raise IllegalAddress(f"共享数组 {self.name} 是 {len(self.shape)} 维的，下标却有 {len(idx)} 个")
        flat = 0
        for i, s, st in zip(idx, self.shape, self._strides):
            if not isinstance(i, (int, np.integer)):
                raise TypeError(f"{self.name}[...] 的下标必须是整数")
            if not 0 <= i < s:
                raise IllegalAddress(f"线程 {_RUN.where()} 越界访问共享数组 {self.name}{list(idx)}（形状 {self.shape}）")
            flat += int(i) * st
        return flat

    def __getitem__(self, idx):
        f = self._flat(idx)
        _RUN.access("shared", self, "ld", f, self.offset + f * self.itemsize, self.itemsize)
        return self._data[f]

    def __setitem__(self, idx, v):
        f = self._flat(idx)
        _RUN.access("shared", self, "st", f, self.offset + f * self.itemsize, self.itemsize)
        self._data[f] = v

    def atomic_add(self, idx, v):
        """共享内存上的原子加，返回旧值。"""
        f = self._flat(idx)
        _RUN.access("shared", self, "atom", f, self.offset + f * self.itemsize, self.itemsize)
        old = self._data[f]
        self._data[f] = old + v
        return old


# =====================================================================
# 4. 线程上下文与统计
# =====================================================================

class ThreadCtx:
    """
    内核的第一个参数 t：一个线程看得到的内置变量和工具函数。

    【属性】threadIdx、blockIdx、blockDim、gridDim（Dim3，和 CUDA 同名）；
        warp（块内第几个线程束）、lane（线程束里的第几个线程，0..31）
    【方法】shared(name, shape, dtype)：申请 / 取得本块的共享数组；syncthreads()：配合 yield 使用的栅栏
    """

    __slots__ = ("threadIdx", "blockIdx", "blockDim", "gridDim", "linear", "_run")

    def __init__(self, run, thread_idx, block_idx, linear):
        self.threadIdx = thread_idx
        self.blockIdx = block_idx
        self.blockDim = run.block
        self.gridDim = run.grid
        self.linear = linear
        self._run = run

    @property
    def warp(self):
        return self.linear // WARP_SIZE

    @property
    def lane(self):
        return self.linear % WARP_SIZE

    def shared(self, name, shape, dtype=np.float32):
        """__shared__ 声明：第一个调用的线程负责"分配"，同一块里的其他线程拿到同一个数组。"""
        return self._run.get_shared(name, shape, dtype)

    @staticmethod
    def syncthreads():
        """__syncthreads()：必须写成 `yield t.syncthreads()`，否则不起作用。"""
        return SYNC


class LaunchStats:
    """
    一次内核发射的统计结果。

    【全局内存】（"请求"= 一个线程束的一条访存指令；"事务"= 一个 32 字节 sector）
        ld_requests / ld_sectors、st_requests / st_sectors、atom_requests / atom_sectors
        thread_loads / thread_stores / thread_atomics ：按线程数的访问次数
        bytes_requested ：线程们要求读写的字节数之和（每个线程每次访问 itemsize 字节）
        unique_sectors  ：整个内核里不同 sector 的个数——如果缓存无限大，显存只需要传这么多（"最少流量"）
    【共享内存】sh_requests（线程束级访问指令数）、sh_wavefronts（实际轮数；有 bank 冲突时 > 请求数）、thread_shared
    【执行】threads、blocks、warps（总线程束数）、idle_lanes（不满的线程束里空着的线程位置数）、barriers（块级栅栏次数）
        warp_issues、thread_instructions（只有 count_instructions=True 时才统计）
    【竞争】races：检测到的数据竞争（最多保存前 20 条描述），race_count：总数
    【分数组】by_array：数组名 → 各项计数
    """

    def __init__(self, grid, block):
        self.grid, self.block = grid, block
        self.blocks = grid.count
        tpb = block.count
        self.threads = self.blocks * tpb
        wpb = -(-tpb // WARP_SIZE)
        self.warps = self.blocks * wpb
        self.idle_lanes = self.blocks * (wpb * WARP_SIZE - tpb)
        self.ld_requests = self.ld_sectors = 0
        self.st_requests = self.st_sectors = 0
        self.atom_requests = self.atom_sectors = 0
        self.thread_loads = self.thread_stores = self.thread_atomics = 0
        self.bytes_requested = 0
        self.unique_sectors = 0
        self.sh_requests = self.sh_wavefronts = self.thread_shared = 0
        self.barriers = 0
        self.warp_issues = self.thread_instructions = 0
        self.counted_instructions = False
        self.races = []
        self.race_count = 0
        self.by_array = defaultdict(lambda: defaultdict(int))

    # 派生指标 ---------------------------------------------------------
    @property
    def global_requests(self):
        return self.ld_requests + self.st_requests + self.atom_requests

    @property
    def global_sectors(self):
        return self.ld_sectors + self.st_sectors + self.atom_sectors

    @property
    def sector_bytes(self):
        """按 32 字节事务计算的全局内存流量（L1 / L2 和 SM 之间实际搬运的字节数，不考虑缓存命中）。"""
        return self.global_sectors * SECTOR_BYTES

    @property
    def sectors_per_request(self):
        """平均每条线程束访存指令产生几个事务；4 字节元素完全合并时是 4，完全不合并时是 32。"""
        return self.global_sectors / self.global_requests if self.global_requests else 0.0

    @property
    def ld_sectors_per_request(self):
        """只看读：平均每条线程束读指令产生几个事务。"""
        return self.ld_sectors / self.ld_requests if self.ld_requests else 0.0

    @property
    def bank_conflict_factor(self):
        """共享内存每条指令平均要几轮：1 表示没有 bank 冲突，32 表示最坏（32 路冲突）。"""
        return self.sh_wavefronts / self.sh_requests if self.sh_requests else 0.0

    @property
    def simt_efficiency(self):
        """线程实际执行的指令数 ÷ (发射次数 × 32)。1 表示每次发射 32 个线程都在干活。"""
        if not self.warp_issues:
            return float("nan")
        return self.thread_instructions / (self.warp_issues * WARP_SIZE)

    def summary(self):
        """常用指标的字典（打印用）。"""
        d = {"threads": self.threads, "blocks": self.blocks, "warps": self.warps,
             "thread_loads": self.thread_loads, "thread_stores": self.thread_stores,
             "ld_requests": self.ld_requests, "ld_sectors": self.ld_sectors,
             "st_requests": self.st_requests, "st_sectors": self.st_sectors,
             "sectors_per_request": round(self.sectors_per_request, 3),
             "unique_sectors": self.unique_sectors, "races": self.race_count}
        if self.sh_requests:
            d.update(sh_requests=self.sh_requests, sh_wavefronts=self.sh_wavefronts,
                     bank_conflict_factor=round(self.bank_conflict_factor, 3))
        if self.counted_instructions:
            d.update(warp_issues=self.warp_issues, thread_instructions=self.thread_instructions,
                     simt_efficiency=round(self.simt_efficiency, 4))
        return d


# =====================================================================
# 5. 发射内核
# =====================================================================

_RUN = None              # 当前正在执行的发射（模块级变量：DeviceArray 通过它知道"现在是哪个线程在访问"）


class _Run:
    """一次 launch 的运行时状态（模块内部使用）。"""

    def __init__(self, grid, block, count_instructions, check_global_races, shared_limit):
        self.grid, self.block = grid, block
        self.stats = LaunchStats(grid, block)
        self.count_instructions = count_instructions
        self.check_global_races = check_global_races
        self.shared_limit = shared_limit
        self.arrays = {}                       # id(数组) → 数组对象（统计时要用名字和元素大小）
        self.unique = set()                    # 全部访问过的全局 sector
        self.global_access = {}                # 全局竞争检测：(数组 id, 下标) → [(全局线程号, 块号, 栅栏周期, 读/写)]
        self.block_linear = 0
        self.in_thread = False                 # 只有在执行某个线程的代码时才计指令

    # ---------------- 块的开始与结束 ----------------
    def start_block(self, block_idx, block_linear):
        self.block_idx = block_idx
        self.block_linear = block_linear
        self.shared = {}                       # 名字 → SharedArray
        self.shared_bytes = 0
        self.groups = defaultdict(list)        # (线程束, 访问点, 第 k 次) → [(字节地址, 字节数)]
        self.line_groups = defaultdict(int)    # (线程束, 代码行, 第 k 次) → 执行这一行的线程数
        self.occ_by_thread = {}                # 线程 → {访问点: 已经执行了几次}
        self.line_occ_by_thread = {}
        self.epoch = 0                         # 第几个栅栏周期
        self.sh_readers = defaultdict(set)     # 本栅栏周期内：(数组 id, 下标) → 读过它的线程
        self.sh_writers = defaultdict(set)

    def set_thread(self, ctx):
        """切换"当前线程"：之后的每一次访存都记在这个线程名下。"""
        self.ctx = ctx
        self.tid = ctx.linear
        self.warp = ctx.linear // WARP_SIZE
        self.gtid = self.block_linear * self.block.count + ctx.linear
        self.occ = self.occ_by_thread.setdefault(ctx.linear, {})
        self.line_occ = self.line_occ_by_thread.setdefault(ctx.linear, {})

    def where(self):
        """报错时描述当前线程的位置。"""
        c = self.ctx
        return f"blockIdx={tuple(c.blockIdx)} threadIdx={tuple(c.threadIdx)}"

    def barrier(self):
        """全块到达栅栏：进入下一个栅栏周期，清空共享内存的竞争记录。"""
        self.epoch += 1
        self.stats.barriers += 1
        self.sh_readers.clear()
        self.sh_writers.clear()

    def get_shared(self, name, shape, dtype):
        arr = self.shared.get(name)
        shape = tuple(int(s) for s in (shape if isinstance(shape, (tuple, list)) else (shape,)))
        if arr is None:
            offset = -(-self.shared_bytes // 16) * 16          # 按 16 字节对齐
            arr = SharedArray(name, shape, dtype, offset)
            if offset + arr.nbytes > self.shared_limit:
                raise LaunchError(f"共享内存超限：本块需要 {offset + arr.nbytes} 字节 > 上限 {self.shared_limit} 字节")
            self.shared_bytes = offset + arr.nbytes
            self.shared[name] = arr
        elif arr.shape != shape or arr.dtype != np.dtype(dtype):
            raise SimError(f"共享数组 {name} 被以不同的形状 / 类型声明了两次")
        return arr

    # ---------------- 访存记录（热路径） ----------------
    def access(self, space, arr, kind, index, addr, nbytes):
        """
        记录一次访问。由 DeviceArray / SharedArray 调用：
            sys._getframe(2) 是"调用 A[i] 的那一行代码"所在的栈帧（0 是本函数，1 是 __getitem__ 等）。
        【Python 小知识】sys._getframe(k) 返回调用栈上往上数第 k 层的帧对象；f_code 是那段代码，
            f_lineno 是当前行号。两者合起来唯一确定"哪条指令"。
        """
        f = sys._getframe(2)
        key = (space, id(arr), kind, f.f_code, f.f_lineno)
        occ = self.occ
        k = occ.get(key, 0)
        occ[key] = k + 1
        self.groups[(self.warp, key, k)].append((addr, nbytes))
        self.arrays[id(arr)] = arr
        if space == "global":
            s0, s1 = addr // SECTOR_BYTES, (addr + nbytes - 1) // SECTOR_BYTES
            self.unique.add(s0)
            if s1 != s0:
                self.unique.update(range(s0 + 1, s1 + 1))
            if self.check_global_races and kind != "atom":
                self._global_race(arr, index, kind, f.f_lineno)
        elif kind != "atom":
            self._shared_race(arr, index, kind, f.f_lineno)

    def _report_race(self, msg):
        self.stats.race_count += 1
        if len(self.stats.races) < 20:
            self.stats.races.append(msg)

    def _shared_race(self, arr, index, kind, lineno):
        k = (id(arr), index)
        tid = self.tid
        if kind == "st":
            others = (self.sh_writers[k] | self.sh_readers[k]) - {tid}
            self.sh_writers[k].add(tid)
        else:
            others = self.sh_writers[k] - {tid}
            self.sh_readers[k].add(tid)
        if others:
            self._report_race(f"共享 {arr.name}[{index}]：线程 {tid} 在第 {lineno} 行{'写' if kind == 'st' else '读'}，"
                              f"同一栅栏周期里线程 {sorted(others)[:3]} 也访问过它（块 {tuple(self.block_idx)}）")

    def _global_race(self, arr, index, kind, lineno):
        k = (id(arr), index)
        me = (self.gtid, self.block_linear, self.epoch, kind)
        hist = self.global_access.setdefault(k, [])
        for g, b, e, kd in hist:
            if g != self.gtid and (kind == "st" or kd == "st") and (b != self.block_linear or e == self.epoch):
                self._report_race(f"全局 {arr.name}[{index}]：线程 {self.gtid} 在第 {lineno} 行"
                                  f"{'写' if kind == 'st' else '读'}，与线程 {g} 的{'写' if kd == 'st' else '读'}冲突")
                break
        hist.append(me)

    def count_line(self, code, lineno):
        """指令计数：当前线程执行了 code 的第 lineno 行。"""
        key = (code, lineno)
        occ = self.line_occ
        k = occ.get(key, 0)
        occ[key] = k + 1
        self.line_groups[(self.warp, key, k)] += 1

    def end_block(self):
        """块执行完：把记录下的访问按"线程束 × 指令"汇总成请求、事务、bank 冲突。"""
        st = self.stats
        for (warp, key, k), lst in self.groups.items():
            space, aid, kind = key[0], key[1], key[2]
            arr = self.arrays[aid]
            rec = st.by_array[arr.name]
            if space == "global":
                secs = set()
                for addr, nb in lst:
                    s0, s1 = addr // SECTOR_BYTES, (addr + nb - 1) // SECTOR_BYTES
                    secs.add(s0)
                    if s1 != s0:
                        secs.update(range(s0 + 1, s1 + 1))
                ns = len(secs)
                st.bytes_requested += sum(nb for _, nb in lst)
                if kind == "ld":
                    st.ld_requests += 1
                    st.ld_sectors += ns
                    st.thread_loads += len(lst)
                elif kind == "st":
                    st.st_requests += 1
                    st.st_sectors += ns
                    st.thread_stores += len(lst)
                else:
                    st.atom_requests += 1
                    st.atom_sectors += ns
                    st.thread_atomics += len(lst)
                rec[kind + "_requests"] += 1
                rec[kind + "_sectors"] += ns
                rec["thread_" + kind] += len(lst)
            else:
                banks = defaultdict(set)
                for off, nb in lst:
                    for w in range(off // BANK_BYTES, (off + nb - 1) // BANK_BYTES + 1):
                        banks[w % NUM_BANKS].add(w)
                wf = max(len(v) for v in banks.values())
                st.sh_requests += 1
                st.sh_wavefronts += wf
                st.thread_shared += len(lst)
                rec["sh_requests"] += 1
                rec["sh_wavefronts"] += wf
                rec["thread_sh"] += len(lst)
        if self.count_instructions:
            st.warp_issues += len(self.line_groups)
            st.thread_instructions += sum(self.line_groups.values())


def _global_tracer(frame, event, arg):
    """sys.settrace 的全局回调：只跟踪内核及其调用的 Python 函数（跳过模拟器自身的代码）。"""
    if frame.f_code.co_filename == _THIS_FILE:
        return None
    return _line_tracer


def _line_tracer(frame, event, arg):
    """局部回调：每执行一行，记一条指令。"""
    if event == "line" and _RUN is not None and _RUN.in_thread:
        _RUN.count_line(frame.f_code, frame.f_lineno)
    return _line_tracer


def _check_config(grid, block):
    """检查发射配置是否合法（对应 CUDA 的 invalid configuration argument）。"""
    for name, d, lim in (("block", block, MAX_BLOCK_DIM), ("grid", grid, MAX_GRID_DIM)):
        for axis, v, m in zip("xyz", d, lim):
            if v < 1:
                raise LaunchError(f"{name}.{axis} = {v}：每一维至少是 1")
            if v > m:
                raise LaunchError(f"{name}.{axis} = {v} 超过上限 {m}")
    if block.count > MAX_THREADS_PER_BLOCK:
        raise LaunchError(f"一个线程块有 {block.count} 个线程，超过上限 {MAX_THREADS_PER_BLOCK}")


def launch(kernel, grid, block, *args, count_instructions=False, check_global_races=False,
           shared_limit=DEFAULT_SHARED_LIMIT):
    """
    发射内核：相当于 CUDA 的 kernel<<<grid, block>>>(*args)，返回 LaunchStats。

    【参数】
        kernel             : 内核函数 kernel(t, *args)。要用 __syncthreads() 的内核必须是生成器函数（里面有 yield）
        grid, block        : int、元组或 Dim3
        *args              : 传给内核的参数（DeviceArray、整数、元组……；对应 CUDA 的内核参数）
        count_instructions : True 时统计指令发射与线程束分化（用 sys.settrace，会慢 3–10 倍）
        check_global_races : True 时检查全局内存的数据竞争（很慢，只在小规模演示时用）
        shared_limit       : 每个块共享内存的上限（字节）
    【思路】
        对每个块：为每个线程建一个 ThreadCtx；
            普通内核：依次调用 kernel(t, ...)；
            生成器内核：反复"让每个还活着的线程跑到下一个 yield 或结束"，直到全部结束；
                如果同一轮里有的线程停在栅栏、有的已经结束 → BarrierError。
        块结束时把访存记录汇总（见 _Run.end_block）。
    【边界】内核不能有返回值（CUDA 的 __global__ 函数返回 void），返回非 None 时报 SimError；
        不支持在内核里再发射内核（动态并行）。
    【Python 小知识】生成器函数（函数体里有 yield）被调用时并不执行，而是返回一个生成器对象；
        next(gen) 让它运行到下一个 yield 并把 yield 的值交出来；函数执行完时 next 抛出 StopIteration，
        异常对象的 .value 是函数的返回值。我们正是用它来实现"线程跑到栅栏处停下"。
    """
    global _RUN
    grid, block = as_dim3(grid), as_dim3(block)
    _check_config(grid, block)
    if _RUN is not None:
        raise SimError("模拟器不支持在内核里再发射内核")
    is_gen = inspect.isgeneratorfunction(kernel)
    run = _Run(grid, block, count_instructions, check_global_races, shared_limit)
    run.stats.counted_instructions = count_instructions
    _RUN = run
    if count_instructions:
        sys.settrace(_global_tracer)
    try:
        threads_idx = [Dim3(x, y, z) for z in range(block.z) for y in range(block.y) for x in range(block.x)]
        bl = 0
        for bz in range(grid.z):
            for by in range(grid.y):
                for bx in range(grid.x):
                    bidx = Dim3(bx, by, bz)
                    run.start_block(bidx, bl)
                    ctxs = [ThreadCtx(run, ti, bidx, lin) for lin, ti in enumerate(threads_idx)]
                    if not is_gen:
                        for c in ctxs:
                            run.set_thread(c)
                            run.in_thread = True
                            ret = kernel(c, *args)
                            run.in_thread = False
                            if ret is not None:
                                raise SimError("内核不能有返回值（__global__ 函数返回 void）：结果要写进显存数组")
                    else:
                        gens = [kernel(c, *args) for c in ctxs]
                        alive = list(range(len(ctxs)))
                        while alive:
                            waiting, finished = [], []
                            for i in alive:
                                run.set_thread(ctxs[i])
                                run.in_thread = True
                                try:
                                    tok = next(gens[i])
                                except StopIteration as stop:
                                    run.in_thread = False
                                    if stop.value is not None:
                                        raise SimError("内核不能有返回值（__global__ 函数返回 void）") from None
                                    finished.append(i)
                                    continue
                                run.in_thread = False
                                if tok is not SYNC:
                                    raise SimError("内核里的 yield 只能用来表示 __syncthreads()：请写 yield t.syncthreads()")
                                waiting.append(i)
                            if waiting and finished:
                                c = ctxs[finished[0]]
                                raise BarrierError(
                                    f"块 {tuple(bidx)} 里有 {len(waiting)} 个线程停在 __syncthreads()，"
                                    f"另有 {len(finished)} 个线程（例如 threadIdx={tuple(c.threadIdx)}）已经结束——"
                                    "真 GPU 上这是未定义行为（常见后果是卡死）。不要在 if 分支里提前 return 再用栅栏。")
                            if waiting:
                                run.barrier()
                            alive = waiting
                    run.end_block()
                    bl += 1
    finally:
        if count_instructions:
            sys.settrace(None)
        _RUN = None
    run.stats.unique_sectors = len(run.unique)
    return run.stats


def ceil_div(a, b):
    """向上取整的整数除法：CUDA 里常写成 (a + b - 1) / b，用来算"需要几个块"。"""
    return -(-a // b)