"""
=============================================================================
CS349E · 第 14 讲共享模块：kbench.py —— 迷你 KernelBench：任务、候选内核、评估器与 fast_p
=============================================================================

【这个模块要解决什么问题】
    KernelBench（Ouyang et al. 2025）让 LLM 读一段 PyTorch 参考实现，写出更快的 CUDA 内核，然后**自动**判分：
    先测正确性（5 组随机输入，atol = rtol = 1e-2），再测速度（相对 PyTorch 的加速比），最后汇总成 fast_p
    ——"既正确、加速比又大于 p 的任务占多少"。
    这个模块在 CPU 上用 numpy 搭一个同样结构的小版本，用来回答三个问题：
      1. 一个"写内核"的任务和它的评估器长什么样？fast_p 怎么算？
      2. 评估器哪里容易被骗？（第 14 讲 §3：输出被"偷"、缓存、懒计算、只对默认参数成立……）
      3. 怎样把评估器做得更结实？

【和真实 KernelBench 的对应】
    PyTorch 的 Model / ModelNew          → 这里的参考源码 ref_src / 候选源码（都是定义 forward 的 Python 源码字符串）
    ModelNew.__init__ 里预处理权重        → 候选可以定义 init(params)，返回预处理后的参数（不计时，和 KernelBench 一样）
    torch.empty（缓存分配器给的未初始化显存）→ 候选通过 alloc(shape) 拿输出缓冲区；朴素评估器用一个"会复用"的缓存池
    CUDA 的异步执行和 stream                → 候选可以返回一个"懒"对象，真正的计算推迟到有人读它的时候
    nvcc 编译错误 / 运行时错误              → Python 的 SyntaxError / 运行时异常
    "内核"本身                              → numpy 代码。速度的差别来自：少分配临时数组、in-place、换算法、预处理权重

【两种评估器】
    evaluate(task, src, mode="naive")：照着 KernelBench 最初版本的做法——默认分布（标准正态）的 5 组输入、默认初始化的参数、
        计时时反复用同一组输入、计时只包住函数调用、候选的命名空间里能看到参考实现。
    evaluate(task, src, mode="hardened")：加上第 14 讲 §3.3 的防御——静态检查、输出缓冲区填 NaN、多种输入分布、
        随机化参数、每次计时都换新输入并在计时区里"同步"（强制求值）、核对计时时的输出、和参考交替计时、"光速"检查。

【关于数字】所有时间都是本机 numpy 的墙钟时间（会随机器、负载、numpy 版本变化）；任务大小按 KernelBench v0.1 的建议，
    让参考实现落在约 1–15 ms。这里的"内核"是 numpy 代码，加速比的来源和 GPU 内核不同（见讲义 §2.4），只用来说明评估的机制。

【运行】python3 kbench.py   # 自检（约 2 秒）
【依赖】numpy；只用标准库的 ast、math、time。
"""

import ast
import math
import time
from dataclasses import dataclass, field

import numpy as np

F32 = np.float32
ATOL = 1e-2          # KernelBench 的默认容差（atol = rtol = 1e-2）
RTOL = 1e-2
N_CORRECT = 5        # KernelBench：5 组随机输入
SPEEDUP_CAP = 1e4    # 候选"几乎不花时间"时，加速比按这个上限显示，避免除以 0

# =============================================================================
# 1. 输入分布
# =============================================================================
DISTS = ("normal", "uniform", "laplace", "lognormal", "offset")
DIST_NOTES = {
    "normal": "标准正态（KernelBench 最初只用这个）",
    "uniform": "(-1, 1) 上的均匀分布",
    "laplace": "拉普拉斯（尾巴比正态重）",
    "lognormal": "对数正态（全正、长尾，最大值上百）",
    "offset": "均值 1000、方差 1 的正态（离 0 很远）",
}


def sample(rng, shape, dist="normal"):
    """
    按分布 dist 生成 float32 的随机数组。

    【参数】rng：np.random.Generator；shape：形状；dist：DISTS 之一
    【返回】float32 数组
    【为什么要多种分布】很多 bug 只在某些输入上暴露：不减最大值的 softmax 在对数正态（值上百）上溢出成 inf；
        "一遍算方差"（E[x²] − E[x]²）在均值很大时有严重的抵消误差（前置讲 P3）。KernelBench v0.1 加了
        rand_mix（正态、均匀、拉普拉斯、指数、对数正态）；"offset"是我加的。
    【Python 小知识】rng.standard_normal(shape, dtype=np.float32) 直接生成 float32，省一次类型转换。
    """
    if dist == "normal":
        return rng.standard_normal(shape, dtype=F32)
    if dist == "uniform":
        return rng.uniform(-1.0, 1.0, shape).astype(F32)
    if dist == "laplace":
        return rng.laplace(0.0, 1.0, shape).astype(F32)
    if dist == "lognormal":
        return rng.lognormal(0.0, 1.0, shape).astype(F32)
    if dist == "offset":
        return (rng.standard_normal(shape, dtype=F32) + F32(1000.0)).astype(F32)
    raise ValueError(f"未知的分布 {dist}")


# =============================================================================
# 2. 任务
# =============================================================================
@dataclass
class Task:
    """
    一个"写内核"的任务（对应 KernelBench 的一个 Model）。

    【字段】
        name, level      ：名字；级别（1 = 单个算子、2 = 几个算子的组合、3 = 一个小网络块）
        ref_src          ：参考实现的源码，定义 forward(inputs, params, alloc)（"LLM"能看到的就是它）
        input_shapes     ：每个输入的形状（元组的列表）
        make_params      ：make_params(rng, randomize) → 参数字典。randomize=False 是"默认初始化"（像 PyTorch 的
                           nn.LayerNorm 把 γ 初始化成 1、β 初始化成 0）；True 是随机化的参数（加固的评估器用）
        min_bytes        ：语义上至少要读写多少字节（输入 + 输出，不含能在 init 里预处理的参数），"光速"检查用
        note             ：一句话说明
    """
    name: str
    level: int
    ref_src: str
    input_shapes: list
    make_params: object
    min_bytes: int
    note: str = ""
    _ref_fn: object = field(default=None, repr=False)

    def ref_fn(self):
        """把参考源码编译成函数（只编译一次）。"""
        if self._ref_fn is None:
            ns = {"np": np, "math": math}
            exec(compile(self.ref_src, f"<ref:{self.name}>", "exec"), ns)
            self._ref_fn = ns["forward"]
        return self._ref_fn

    def make_inputs(self, rng, dist="normal"):
        """按 input_shapes 生成一组输入（元组）。"""
        return tuple(sample(rng, s, dist) for s in self.input_shapes)


def _no_params(rng, randomize):
    return {}


def _ln_params(rng, randomize, d=4096):
    """LayerNorm 的 γ、β：默认初始化是 1 和 0；随机化时 γ ~ N(1, 0.5²)、β ~ N(0, 0.5²)。"""
    if not randomize:
        return {"gamma": np.ones(d, F32), "beta": np.zeros(d, F32)}
    return {"gamma": (1 + 0.5 * rng.standard_normal(d)).astype(F32), "beta": (0.5 * rng.standard_normal(d)).astype(F32)}


def _mds_params(rng, randomize, n=2048, k=2048):
    """matmul → ÷2 → 按行求和 → ×scale 的权重；W 的方差按 1/K 缩放（像 nn.Linear 的初始化量级）。"""
    W = (rng.standard_normal((n, k)) / math.sqrt(k)).astype(F32)
    return {"W": W, "scale": F32(rng.uniform(-2, 2) if randomize else 1.5)}


def _lrs_params(rng, randomize, n=2048, k=1024):
    """线性层 + ReLU + 逐通道缩放：默认的缩放因子都是正的（0.5–1.5）；随机化时有正有负。"""
    W = (rng.standard_normal((n, k)) / math.sqrt(k)).astype(F32)
    b = (0.1 * rng.standard_normal(n)).astype(F32)
    s = rng.uniform(-1.5, 1.5, n) if randomize else rng.uniform(0.5, 1.5, n)
    return {"W": W, "b": b, "s": s.astype(F32)}


def _sam_params(rng, randomize, n=1024, k=1024):
    W = (rng.standard_normal((n, k)) / math.sqrt(k)).astype(F32)
    return {"W": W}


def _mlp_params(rng, randomize, d=1024, ff=2816):
    """RMSNorm 的 γ（默认 1）+ SwiGLU 的三个矩阵。"""
    g = (1 + 0.5 * rng.standard_normal(d)).astype(F32) if randomize else np.ones(d, F32)
    return {"gamma": g,
            "wg": (rng.standard_normal((d, ff)) / math.sqrt(d)).astype(F32),
            "wu": (rng.standard_normal((d, ff)) / math.sqrt(d)).astype(F32),
            "wd": (rng.standard_normal((ff, d)) / math.sqrt(ff)).astype(F32)}


REF_DIAG = '''
def forward(inputs, params, alloc):
    a, B = inputs
    return np.diag(a) @ B                  # 先造出 N×N 的对角矩阵，再做一次完整的矩阵乘法
'''

REF_SOFTMAX = '''
def forward(inputs, params, alloc):
    (x,) = inputs
    m = x.max(axis=1, keepdims=True)       # 减去每行最大值：数值稳定（前置讲 P3）
    e = np.exp(x - m)
    return e / e.sum(axis=1, keepdims=True)
'''

REF_GELU = '''
def forward(inputs, params, alloc):
    (x,) = inputs                          # GELU 的 tanh 近似（nn.GELU(approximate="tanh")）
    c = np.float32(math.sqrt(2.0 / math.pi))
    return 0.5 * x * (1.0 + np.tanh(c * (x + 0.044715 * x ** 3)))
'''

REF_MATMUL = '''
def forward(inputs, params, alloc):
    A, B = inputs
    return A @ B                           # numpy 调用 BLAS 的 sgemm（相当于 PyTorch 调 cuBLAS）
'''

REF_LAYERNORM = '''
def forward(inputs, params, alloc):
    (x,) = inputs
    mu = x.mean(axis=1, keepdims=True)
    var = ((x - mu) ** 2).mean(axis=1, keepdims=True)          # 两遍法求方差
    return (x - mu) / np.sqrt(var + 1e-5) * params["gamma"] + params["beta"]
'''

REF_MDS = '''
def forward(inputs, params, alloc):
    (x,) = inputs                          # KernelBench Level 2 "Matmul_Divide_Sum_Scaling" 的结构
    y = x @ params["W"].T
    y = y / 2
    y = y.sum(axis=1, keepdims=True)
    return y * params["scale"]
'''

REF_LRS = '''
def forward(inputs, params, alloc):
    (x,) = inputs                          # 线性层 → ReLU → 逐通道缩放
    y = x @ params["W"].T + params["b"]
    y = np.maximum(y, 0)
    return y * params["s"]
'''

REF_SAM = '''
def forward(inputs, params, alloc):
    (x,) = inputs                          # 线性层 → 按特征求平均（keepdims）→ 对"长度为 1 的那一维"做 softmax
    y = x @ params["W"].T
    m = y.mean(axis=1, keepdims=True)
    e = np.exp(m - m.max(axis=1, keepdims=True))
    return e / e.sum(axis=1, keepdims=True)
'''

REF_MLP = '''
def forward(inputs, params, alloc):
    (x,) = inputs                          # 一个 Llama 式的 MLP 块：RMSNorm → SwiGLU → 残差
    h = x / np.sqrt((x * x).mean(axis=1, keepdims=True) + 1e-6) * params["gamma"]
    g = h @ params["wg"]
    u = h @ params["wu"]
    s = g / (1.0 + np.exp(-g)) * u
    return x + s @ params["wd"]
'''

TASKS = [
    Task("L1_diag_matmul", 1, REF_DIAG, [(2048,), (2048, 2048)], _no_params, (2048 + 2 * 2048 * 2048) * 4,
         "diag(a) @ B（KernelBench Level 1 第 11 题的结构）"),
    Task("L1_softmax", 1, REF_SOFTMAX, [(1024, 4096)], _no_params, 2 * 1024 * 4096 * 4, "按行 softmax"),
    Task("L1_gelu", 1, REF_GELU, [(1024, 2048)], _no_params, 2 * 1024 * 2048 * 4, "GELU（tanh 近似）"),
    Task("L1_matmul", 1, REF_MATMUL, [(1024, 1024), (1024, 1024)], _no_params, 3 * 1024 * 1024 * 4,
         "1024³ 矩阵乘法（对手是厂商的 BLAS）"),
    Task("L1_layernorm", 1, REF_LAYERNORM, [(2048, 4096)], _ln_params, 2 * 2048 * 4096 * 4,
         "LayerNorm（γ、β 默认初始化为 1、0）"),
    Task("L2_matmul_div_sum_scale", 2, REF_MDS, [(1024, 2048)], _mds_params, (1024 * 2048 + 1024) * 4,
         "matmul → ÷2 → 按行求和 → ×scale"),
    Task("L2_linear_relu_scale", 2, REF_LRS, [(2048, 1024)], _lrs_params, (2048 * 1024 + 2048 * 2048) * 4,
         "线性层 → ReLU → 逐通道缩放"),
    Task("L2_softmax_after_mean", 2, REF_SAM, [(2048, 1024)], _sam_params, (2048 * 1024 + 2048) * 4,
         "线性层 → 求平均 → 对长度 1 的维做 softmax（KernelBench v0.1 修掉的退化题）"),
    Task("L3_mlp_block", 3, REF_MLP, [(512, 1024)], _mlp_params, 2 * 512 * 1024 * 4, "RMSNorm + SwiGLU + 残差"),
]
TASK = {t.name: t for t in TASKS}


# =============================================================================
# 3. 候选库：模拟"LLM 写出来的内核"
# =============================================================================
@dataclass
class Candidate:
    """
    一个候选实现（模拟 LLM 的一次输出）。

    【字段】task：任务名；name：短名；src：源码；intent：写它的人"想做什么"（用于展示，不参与判分）
    """
    task: str
    name: str
    src: str
    intent: str


def _c(task, name, intent, src):
    return Candidate(task, name, src, intent)


LIBRARY = [
    # ---- L1_diag_matmul ----
    _c("L1_diag_matmul", "row_scale", "不造对角矩阵，直接把 B 的第 i 行乘 a[i]（KernelBench 报告的 13× 例子）", '''
def forward(inputs, params, alloc):
    a, B = inputs
    return a[:, None] * B
'''),
    _c("L1_diag_matmul", "col_scale", "同样的想法，但广播方向写反了（乘到了列上）", '''
def forward(inputs, params, alloc):
    a, B = inputs
    return B * a[None, :]
'''),
    _c("L1_diag_matmul", "row_loop", "逐行循环缩放（正确，但 Python 循环有开销）", '''
def forward(inputs, params, alloc):
    a, B = inputs
    out = alloc(B.shape)
    for i in range(B.shape[0]):
        np.multiply(B[i], a[i], out=out[i])
    return out
'''),
    _c("L1_diag_matmul", "syntax", "少了冒号（\"编译错误\"）", '''
def forward(inputs, params, alloc)
    a, B = inputs
    return a[:, None] * B
'''),
    # ---- L1_softmax ----
    _c("L1_softmax", "inplace", "就地计算，少分配两个临时数组", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = alloc(x.shape)
    np.subtract(x, x.max(axis=1, keepdims=True), out=y)
    np.exp(y, out=y)
    y /= y.sum(axis=1, keepdims=True)
    return y
'''),
    _c("L1_softmax", "no_max", "省掉\"减最大值\"那一遍（在标准正态输入上看不出问题）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = alloc(x.shape)
    np.exp(x, out=y)
    y /= y.sum(axis=1, keepdims=True)
    return y
'''),
    _c("L1_softmax", "wrong_axis", "沿错了维度归一化", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    e = np.exp(x - x.max(axis=0, keepdims=True))
    return e / e.sum(axis=0, keepdims=True)
'''),
    _c("L1_softmax", "bad_axis", "维度写错，运行时报错", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    e = np.exp(x - x.max(axis=2, keepdims=True))
    return e / e.sum(axis=2, keepdims=True)
'''),
    # ---- L1_gelu ----
    _c("L1_gelu", "fused_inplace", "x³ 改成 x·x·x、全部就地计算（一个\"融合\"的逐元素内核）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    c = np.float32(math.sqrt(2.0 / math.pi))
    y = alloc(x.shape)
    np.multiply(x, x, out=y)
    y *= np.float32(0.044715)
    y += 1.0
    y *= x
    y *= c
    np.tanh(y, out=y)
    y += 1.0
    y *= x
    y *= np.float32(0.5)
    return y
'''),
    _c("L1_gelu", "sigmoid_approx", "换成更便宜的近似 x·sigmoid(1.702x)（误差超过容差）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    return x / (1.0 + np.exp(-1.702 * x))
'''),
    _c("L1_gelu", "fp16", "降到 FP16 计算（numpy 在 CPU 上没有快速的 FP16 运算）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    h = x.astype(np.float16)
    c = np.float16(math.sqrt(2.0 / math.pi))
    y = np.float16(0.5) * h * (np.float16(1.0) + np.tanh(c * (h + np.float16(0.044715) * h * h * h)))
    return y.astype(np.float32)
'''),
    # ---- L1_matmul ----
    _c("L1_matmul", "same", "原样调用库（和参考一模一样）", '''
def forward(inputs, params, alloc):
    A, B = inputs
    return A @ B
'''),
    _c("L1_matmul", "blocked", "手写 256×256 分块（每块仍调库）", '''
def forward(inputs, params, alloc):
    A, B = inputs
    n, k = A.shape
    m = B.shape[1]
    C = alloc((n, m))
    C[...] = 0
    T = 256
    for i in range(0, n, T):
        for j in range(0, m, T):
            for p in range(0, k, T):
                C[i:i+T, j:j+T] += A[i:i+T, p:p+T] @ B[p:p+T, j:j+T]
    return C
'''),
    _c("L1_matmul", "strassen1", "一层 Strassen（7 次半尺寸乘法代替 8 次）", '''
def forward(inputs, params, alloc):
    A, B = inputs
    h = A.shape[0] // 2
    A11, A12, A21, A22 = A[:h, :h], A[:h, h:], A[h:, :h], A[h:, h:]
    B11, B12, B21, B22 = B[:h, :h], B[:h, h:], B[h:, :h], B[h:, h:]
    M1 = (A11 + A22) @ (B11 + B22)
    M2 = (A21 + A22) @ B11
    M3 = A11 @ (B12 - B22)
    M4 = A22 @ (B21 - B11)
    M5 = (A11 + A12) @ B22
    M6 = (A21 - A11) @ (B11 + B12)
    M7 = (A12 - A22) @ (B21 + B22)
    C = alloc((2 * h, 2 * h))
    C[:h, :h] = M1 + M4 - M5 + M7
    C[:h, h:] = M3 + M5
    C[h:, :h] = M2 + M4
    C[h:, h:] = M1 - M2 + M3 + M6
    return C
'''),
    _c("L1_matmul", "float64", "升到 FP64 计算再转回（更准，但慢）", '''
def forward(inputs, params, alloc):
    A, B = inputs
    return (A.astype(np.float64) @ B.astype(np.float64)).astype(np.float32)
'''),
    # ---- L1_layernorm ----
    _c("L1_layernorm", "fused", "两遍法 + 就地计算 + 用 einsum 求平方和", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = alloc(x.shape)
    np.subtract(x, x.mean(axis=1, keepdims=True), out=y)
    var = np.einsum("ij,ij->i", y, y)[:, None] / x.shape[1]
    y *= 1.0 / np.sqrt(var + 1e-5)
    y *= params["gamma"]
    y += params["beta"]
    return y
'''),
    _c("L1_layernorm", "one_pass_var", "一遍算出 E[x] 和 E[x²]，方差 = E[x²] − E[x]²（少读一遍，但有抵消误差）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    mu = x.mean(axis=1, keepdims=True)
    ex2 = np.einsum("ij,ij->i", x, x)[:, None] / x.shape[1]
    y = alloc(x.shape)
    np.subtract(x, mu, out=y)
    y *= 1.0 / np.sqrt(np.maximum(ex2 - mu * mu, 0.0) + 1e-5)
    y *= params["gamma"]
    y += params["beta"]
    return y
'''),
    _c("L1_layernorm", "no_eps", "忘了 eps，变成 1/sqrt(var)（在这些输入上差别很小）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = x - x.mean(axis=1, keepdims=True)
    y = y / np.sqrt((y * y).mean(axis=1, keepdims=True))
    return y * params["gamma"] + params["beta"]
'''),
    # ---- L2_matmul_div_sum_scale ----
    _c("L2_matmul_div_sum_scale", "algebraic", "Σ_j (x·W_jᵀ) = x·(Σ_j W_j)：矩阵乘法变成矩阵-向量乘法，权重的列和在 init 里预先算好", '''
def init(params):
    return {"w": (params["W"].sum(axis=0) * (params["scale"] / 2)).astype(np.float32)}

def forward(inputs, params, alloc):
    (x,) = inputs
    return (x @ params["w"])[:, None]
'''),
    _c("L2_matmul_div_sum_scale", "fused_epilogue", "把 ÷2 和 ×scale 合成一次乘法（只省一个逐元素的遍历）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = x @ params["W"].T
    return y.sum(axis=1, keepdims=True) * (params["scale"] / 2)
'''),
    _c("L2_matmul_div_sum_scale", "forgot_div", "合并时漏掉了 ÷2", '''
def init(params):
    return {"w": (params["W"].sum(axis=0) * params["scale"]).astype(np.float32)}

def forward(inputs, params, alloc):
    (x,) = inputs
    return (x @ params["w"])[:, None]
'''),
    # ---- L2_linear_relu_scale ----
    _c("L2_linear_relu_scale", "inplace_epilogue", "矩阵乘法之后的 +b、ReLU、×s 都就地做（\"epilogue 融合\"）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = x @ params["W"].T
    y += params["b"]
    np.maximum(y, 0, out=y)
    y *= params["s"]
    return y
'''),
    _c("L2_linear_relu_scale", "fold_scale", "把 s 折进 W 和 b：relu(z)·s = relu(z·s)（只在 s > 0 时成立）", '''
def init(params):
    s = params["s"]
    return {"W": (params["W"] * s[:, None]).astype(np.float32), "b": (params["b"] * s).astype(np.float32)}

def forward(inputs, params, alloc):
    (x,) = inputs
    y = x @ params["W"].T
    y += params["b"]
    np.maximum(y, 0, out=y)
    return y
'''),
    _c("L2_linear_relu_scale", "transpose_bug", "权重没转置（形状不匹配，运行时报错）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = x @ params["W"] + params["b"]
    return np.maximum(y, 0) * params["s"]
'''),
    # ---- L2_softmax_after_mean ----
    _c("L2_softmax_after_mean", "fused", "照原样算，只是就地做 epilogue", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    y = x @ params["W"].T
    m = y.mean(axis=1, keepdims=True)
    m -= m.max(axis=1, keepdims=True)
    np.exp(m, out=m)
    m /= m.sum(axis=1, keepdims=True)
    return m
'''),
    # ---- L3_mlp_block ----
    _c("L3_mlp_block", "fused_gate_up", "gate 和 up 两个矩阵在 init 里拼成一个，一次矩阵乘法；SiLU 就地算", '''
def init(params):
    p = dict(params)
    p["wgu"] = np.ascontiguousarray(np.concatenate([params["wg"], params["wu"]], axis=1))
    return p

def forward(inputs, params, alloc):
    (x,) = inputs
    d = x.shape[1]
    ff = params["wg"].shape[1]
    h = x * (1.0 / np.sqrt(np.einsum("ij,ij->i", x, x)[:, None] / d + 1e-6))
    h *= params["gamma"]
    gu = h @ params["wgu"]
    g = gu[:, :ff]
    s = np.exp(-g)
    s += 1.0
    np.divide(g, s, out=s)
    s *= gu[:, ff:]
    y = s @ params["wd"]
    y += x
    return y
'''),
    _c("L3_mlp_block", "forgot_residual", "同样的融合，但漏了残差", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    h = x / np.sqrt((x * x).mean(axis=1, keepdims=True) + 1e-6) * params["gamma"]
    g = h @ params["wg"]
    s = g / (1.0 + np.exp(-g)) * (h @ params["wu"])
    return s @ params["wd"]
'''),
    _c("L3_mlp_block", "shape_bug", "残差加在了错误的张量上（形状不匹配）", '''
def forward(inputs, params, alloc):
    (x,) = inputs
    h = x / np.sqrt((x * x).mean(axis=1, keepdims=True) + 1e-6) * params["gamma"]
    g = h @ params["wg"]
    s = g / (1.0 + np.exp(-g)) * (h @ params["wu"])
    return s + s @ params["wd"]
'''),
]


# =============================================================================
# 4. 编译、静态检查、分配器
# =============================================================================
BANNED_NAMES = {"reference", "__import__", "globals", "locals", "vars", "eval", "exec", "compile", "open",
                "getattr", "setattr", "__builtins__", "sys", "os", "inspect", "gc"}


def static_check(src):
    """
    加固的评估器在运行之前做的静态检查（读源码的语法树，不执行）。

    【查什么】1. 能解析（否则就是"编译错误"）；2. 定义了 forward；3. 没有 import（只能用给它的 np、math）；
        4. 没有用到禁用的名字：reference（调用参考实现）、eval / exec / globals / getattr / __builtins__ 之类
        能绕开检查的东西、sys / os / gc / inspect（能翻出评估器内部状态的模块）。
    【返回】(ok, 消息)
    【边界】静态检查只能挡住"明着来"的作弊；KernelBench v0.1 的经验是：补上一个漏洞，模型会找到更隐蔽的下一个，
        所以分高的结果最后还要人看（讲义 §3.3）。
    【Python 小知识】ast.parse 把源码变成语法树；ast.walk 遍历所有节点；ast.Name 是变量名，ast.Attribute 是 a.b 的 b。
    """
    try:
        tree = ast.parse(src)
    except SyntaxError as e:
        return False, f"SyntaxError: {e.msg}（第 {e.lineno} 行）"
    if not any(isinstance(n, ast.FunctionDef) and n.name == "forward" for n in tree.body):
        return False, "没有定义 forward"
    for n in ast.walk(tree):
        if isinstance(n, (ast.Import, ast.ImportFrom)):
            return False, "不允许 import"
        if isinstance(n, ast.Name) and n.id in BANNED_NAMES:
            return False, f"用到了禁用的名字 {n.id}"
        if isinstance(n, ast.Attribute) and n.attr.startswith("__") and n.attr != "__init__":
            return False, f"用到了双下划线属性 {n.attr}"
    return True, "ok"


def load_candidate(src, namespace):
    """
    "编译"候选：在给定的命名空间里执行源码，取出 forward 和（可选的）init。

    【返回】(forward, init 或 None, 错误消息或 None)
    【边界】SyntaxError → "编译错误"；执行顶层代码时抛出的异常也算编译阶段的错误。
    """
    try:
        code = compile(src, "<candidate>", "exec")
    except SyntaxError as e:
        return None, None, f"SyntaxError: {e.msg}（第 {e.lineno} 行）"
    ns = dict(namespace)
    try:
        exec(code, ns)
    except Exception as e:  # noqa: BLE001  顶层代码出错
        return None, None, f"{type(e).__name__}: {e}"
    if "forward" not in ns:
        return None, None, "没有定义 forward"
    return ns["forward"], ns.get("init"), None


class CachingPool:
    """
    一个会"复用"缓冲区的分配器（模拟 PyTorch 的 CUDA 缓存分配器）。

    【思路】free(arr) 把数组放回按 (形状, dtype) 分组的空闲列表；alloc 优先从空闲列表里拿——拿到的缓冲区里
        还留着上一次的内容（就像 torch.empty 给你的显存里是上一个张量留下的值）。
    【系统联系】METR 怀疑一个"快得不可思议"的结果来自 torch.empty 复用了参考实现输出的显存；Sakana 的
        AI CUDA Engineer 也报告过评估代码里的内存漏洞让内核跳过了正确性检查（讲义 §3.2）。
    """

    def __init__(self):
        self.free_list = {}

    def alloc(self, shape, dtype=F32):
        key = (tuple(shape), np.dtype(dtype).str)
        lst = self.free_list.get(key)
        if lst:
            return lst.pop()
        return np.empty(shape, dtype)

    def free(self, arr):
        key = (arr.shape, arr.dtype.str)
        self.free_list.setdefault(key, []).append(arr)


def poisoned_alloc(shape, dtype=F32):
    """加固的分配器：每次都给一块新内存，并填满 NaN——没写满输出的候选一定会被发现。"""
    return np.full(shape, np.nan, dtype)


# =============================================================================
# 5. 正确性比较与计时
# =============================================================================
def check_output(out, expected, strict_type):
    """
    比较候选的输出和参考输出。

    【参数】strict_type：True 时要求 out 是真正的 np.ndarray（不接受"懒"对象），且 dtype 是 float32
    【返回】(ok, 说明)
    【规则】形状要一样；allclose(atol = rtol = 1e-2)；任何 NaN / inf 都算错（参考输出里没有）。
    """
    if strict_type and not isinstance(out, np.ndarray):
        return False, f"返回的是 {type(out).__name__}，不是 ndarray"
    try:
        arr = np.asarray(out)
    except Exception as e:  # noqa: BLE001
        return False, f"无法转成数组：{type(e).__name__}"
    if arr.shape != expected.shape:
        return False, f"形状 {arr.shape} ≠ {expected.shape}"
    if strict_type and arr.dtype != F32:
        return False, f"dtype 是 {arr.dtype}，不是 float32"
    if not np.all(np.isfinite(arr)):
        return False, "输出里有 NaN / inf"
    with np.errstate(all="ignore"):
        err = np.abs(arr.astype(np.float64) - expected.astype(np.float64))
        tol = ATOL + RTOL * np.abs(expected.astype(np.float64))
    bad = int(np.count_nonzero(err > tol))
    if bad:
        return False, f"{bad} 个元素超出容差（最大误差 {float(err.max()):.3g}）"
    return True, "ok"


def time_naive(fn, n=20, warmup=3):
    """
    朴素的计时：同一组输入反复调用，计时只包住函数调用本身，取平均（KernelBench 最初的做法是 3 次预热 + 100 次取平均）。

    【参数】fn：无参函数
    【返回】每次调用的平均秒数
    【边界】如果 fn 返回一个"懒"对象（还没真正计算），这里量到的只是"发起"的时间——就像 GPU 上没有同步就停表。
    """
    for _ in range(warmup):
        fn()
    t = 0.0
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        t += time.perf_counter() - t0
    return t / n


_MACHINE = {}


def machine_speed():
    """
    测一下本机的"光速"：大数组拷贝的带宽（字节 / 秒）。只测一次，结果缓存。

    【用途】"光速"检查：一个任务语义上至少要读写 min_bytes 字节，所以时间不可能少于 min_bytes / 带宽。
        比这还快，多半是没干活（KernelBench v0.1 用 roofline 的下界做同样的检查）。
    【边界】numpy 单线程拷贝通常达不到内存的峰值带宽，所以这个"光速"偏保守（偏慢）；检查时再乘 0.5 的余量。
    """
    if "bw" not in _MACHINE:
        src = np.ones(64 * 1024 * 1024 // 4, F32)
        dst = np.empty_like(src)
        best = math.inf
        for _ in range(5):
            t0 = time.perf_counter()
            np.copyto(dst, src)
            best = min(best, time.perf_counter() - t0)
        _MACHINE["bw"] = 2 * src.nbytes / best
    return _MACHINE["bw"]


# =============================================================================
# 6. 评估器
# =============================================================================
@dataclass
class Result:
    """一次评估的结果。status ∈ compile_error / rejected / runtime_error / wrong / correct。"""
    task: str
    cand: str
    mode: str
    status: str
    msg: str = ""
    t_ref: float = math.nan
    t_cand: float = math.nan
    speedup: float = math.nan
    flags: list = field(default_factory=list)

    @property
    def correct(self):
        return self.status == "correct"


def _run(fwd, inputs, params, alloc):
    return fwd(inputs, params, alloc)


def evaluate(task, src, mode="naive", seed=0, n_time=20, cand_name="?", timing=True):
    """
    评估一个候选。

    【参数】task：Task；src：候选源码；mode："naive" 或 "hardened"；seed：随机种子；n_time：计时次数；
        timing=False 时只测正确性（快）
    【返回】Result
    【流程】
        naive：  候选能看到 reference；参数 = 默认初始化；5 组标准正态输入；参考的输出算完后"还给"缓存池，
                 候选的 alloc 从同一个池里拿；计时：同一组输入、只包住调用、取平均，参考和候选各测一轮。
        hardened：静态检查；候选看不到 reference；参数 = 默认初始化 1 组 + 随机化 2 组；5 种输入分布；
                 alloc 给填满 NaN 的新内存；输出必须是 float32 的 ndarray；计时时每次都重新生成输入（写进同一块内存，
                 让"按地址缓存"失效）、在计时区里 np.asarray 强制求值（相当于 GPU 上的同步）、核对两次计时的输出、
                 参考和候选交替测、取中位数；最后做"光速"检查。
    【边界】计时的结果是墙钟时间，有噪声；加速比 = 参考时间 / 候选时间。
    """
    hard = mode == "hardened"
    ref = task.ref_fn()
    rng = np.random.default_rng(seed)
    res = Result(task.name, cand_name, mode, "correct")
    # --- 编译 ---
    if hard:
        ok, msg = static_check(src)
        if not ok:
            res.status = "compile_error" if msg.startswith("SyntaxError") else "rejected"
            res.msg = msg
            return res
        namespace = {"np": np, "math": math}
    else:
        namespace = {"np": np, "math": math, "reference": ref}
    fwd, init, err = load_candidate(src, namespace)
    if err is not None:
        res.status, res.msg = "compile_error", err
        return res
    # --- 正确性 ---
    if hard:
        param_sets = [task.make_params(rng, False), task.make_params(rng, True), task.make_params(rng, True)]
        trials = [(param_sets[i % 3], d) for i, d in enumerate(DISTS)]
    else:
        p0 = task.make_params(rng, False)
        trials = [(p0, "normal")] * N_CORRECT
    pool = CachingPool()
    prepared = {}
    for params, dist in trials:
        key = id(params)
        try:
            if key not in prepared:
                prepared[key] = init(dict(params)) if init else params
            inputs = task.make_inputs(rng, dist)
            with np.errstate(all="ignore"):
                r = np.asarray(ref(inputs, params, None), dtype=F32)
            if hard:
                expected = r
                with np.errstate(all="ignore"):
                    out = _run(fwd, inputs, prepared[key], poisoned_alloc)
            else:
                expected = r.copy()
                pool.free(r)                       # 参考的输出"还给"缓存池……
                with np.errstate(all="ignore"):
                    out = _run(fwd, inputs, prepared[key], pool.alloc)   # ……候选从同一个池里拿缓冲区
        except Exception as e:  # noqa: BLE001
            res.status, res.msg = "runtime_error", f"{type(e).__name__}: {str(e)[:80]}"
            return res
        ok, msg = check_output(out, expected, strict_type=hard)
        if not ok:
            res.status, res.msg = "wrong", f"[{dist}] {msg}"
            return res
    if not timing:
        return res
    # --- 计时 ---
    params = task.make_params(np.random.default_rng(seed + 1), False)
    p_c = init(dict(params)) if init else params
    inputs = task.make_inputs(np.random.default_rng(seed + 2))
    if not hard:
        with np.errstate(all="ignore"):
            res.t_ref = time_naive(lambda: ref(inputs, params, None), n_time)
            res.t_cand = time_naive(lambda: _run(fwd, inputs, p_c, pool.alloc), n_time)
    else:
        trng = np.random.default_rng(seed + 3)
        tr, tc, check_at = [], [], {n_time // 3, (2 * n_time) // 3}
        for i in range(n_time + 2):
            for arr in inputs:                     # 每次都重新生成输入，写进同一块内存
                trng.standard_normal(arr.shape, dtype=F32, out=arr)
            with np.errstate(all="ignore"):
                t0 = time.perf_counter()
                r = np.asarray(ref(inputs, params, None))
                t1 = time.perf_counter()
                y = _run(fwd, inputs, p_c, poisoned_alloc)
                y = np.asarray(y)                  # 强制求值（GPU 上的"同步"）放在计时区里
                t2 = time.perf_counter()
            if i >= 2:                             # 前两次当预热
                tr.append(t1 - t0)
                tc.append(t2 - t1)
            if i in check_at:                      # 核对计时时的输出
                ok, msg = check_output(y, r.astype(F32), strict_type=True)
                if not ok:
                    res.status, res.msg = "wrong", f"[计时中的输出] {msg}"
                    return res
        res.t_ref, res.t_cand = float(np.median(tr)), float(np.median(tc))
        t_sol = task.min_bytes / machine_speed()
        if res.t_cand < 0.5 * t_sol:
            res.flags.append(f"比\"光速\"还快：{res.t_cand * 1e3:.3f} ms < 一半的 {t_sol * 1e3:.3f} ms")
    res.speedup = min(SPEEDUP_CAP, res.t_ref / max(res.t_cand, 1e-12))
    if hard and res.flags:
        res.status = "rejected"
        res.msg = res.flags[0]
    return res


def task_sanity(task, seed=0, n=3):
    """
    检查任务本身是否"退化"：输出几乎不随输入变化（METR / KernelBench v0.1 筛掉了这类题）。

    【做法】用 n 组不同的随机输入（同一组参数）跑参考实现，看输出在不同输入之间的最大差别，
        和输出本身的量级比。差别 < 1e-4 × 量级 → 退化（"返回常数"就能通过，还能"无限快"）。
    【返回】(是否退化, 输出在输入之间的最大变化, 输出的最大绝对值)
    """
    rng = np.random.default_rng(seed)
    params = task.make_params(rng, False)
    outs = [np.asarray(task.ref_fn()(task.make_inputs(rng), params, None), dtype=np.float64) for _ in range(n)]
    spread = max(float(np.abs(o - outs[0]).max()) for o in outs[1:])
    scale = max(float(np.abs(o).max()) for o in outs)
    return spread < 1e-4 * max(scale, 1e-30), spread, scale


# =============================================================================
# 7. fast_p
# =============================================================================
def fast_p(results, p):
    """
    fast_p = 既正确、加速比又 > p 的任务比例（KernelBench 的定义；fast_0 就是正确率）。

    【参数】results：每个任务一个 Result（或 None 表示没有提交）
    """
    if not results:
        return 0.0
    return sum(1 for r in results if r is not None and r.correct and r.speedup > p) / len(results)


def fast_p_at_k(q, k):
    """
    重复采样 k 次、至少一次"正确且加速比 > p"的概率：1 − (1 − q)^k。

    【参数】q：每次采样成功的概率（每个任务一个，列表）；k：采样次数
    【返回】对所有任务取平均的 fast_p@k
    【边界】假设每次采样独立（真实的高温采样大致如此）；q = 0 的任务采多少次都是 0——
        KernelBench 报告 DeepSeek-V3 在 34 个卷积变体上采 100 次也没有一个正确的。
    """
    q = np.asarray(q, dtype=np.float64)
    return float(np.mean(1 - (1 - q) ** k))


def library_for(task_name):
    """某个任务在候选库里的全部候选。"""
    return [c for c in LIBRARY if c.task == task_name]


# =============================================================================
# 自检
# =============================================================================
def _selftest():
    # 参考实现能跑、形状对、有限
    rng = np.random.default_rng(0)
    for t in TASKS:
        y = np.asarray(t.ref_fn()(t.make_inputs(rng), t.make_params(rng, False), None))
        assert np.all(np.isfinite(y)), t.name
    # 退化检测：只有 softmax_after_mean 退化
    for t in TASKS:
        deg = task_sanity(t)[0]
        assert deg == (t.name == "L2_softmax_after_mean"), t.name
    # 静态检查
    assert static_check("def forward(inputs, params, alloc):\n    return reference(inputs, params, alloc)\n")[0] is False
    assert static_check("import os\ndef forward(i, p, a):\n    return 0\n")[0] is False
    assert static_check(LIBRARY[0].src)[0]
    # 缓存池确实会复用
    pool = CachingPool()
    a = pool.alloc((3,))
    a[:] = 7
    pool.free(a)
    assert pool.alloc((3,)) is a
    # 候选库：只测正确性（快）
    expect = {("L1_diag_matmul", "row_scale"): "correct", ("L1_diag_matmul", "col_scale"): "wrong",
              ("L1_diag_matmul", "syntax"): "compile_error", ("L1_softmax", "bad_axis"): "runtime_error",
              ("L2_linear_relu_scale", "fold_scale"): "correct"}
    for (tn, cn), st in expect.items():
        c = next(c for c in LIBRARY if c.task == tn and c.name == cn)
        r = evaluate(TASK[tn], c.src, "naive", timing=False, cand_name=cn)
        assert r.status == st, (tn, cn, r.status, r.msg)
    # 加固的评估器抓住"只对默认参数成立"的折叠和不稳定的 softmax
    for tn, cn in (("L2_linear_relu_scale", "fold_scale"), ("L1_softmax", "no_max")):
        c = next(c for c in LIBRARY if c.task == tn and c.name == cn)
        r = evaluate(TASK[tn], c.src, "hardened", timing=False, cand_name=cn)
        assert r.status == "wrong", (tn, cn, r.status, r.msg)
    # 一次完整的计时评估
    c = next(c for c in LIBRARY if c.task == "L1_diag_matmul" and c.name == "row_scale")
    r = evaluate(TASK["L1_diag_matmul"], c.src, "hardened", n_time=5)
    assert r.correct and r.speedup > 1.5, (r.status, r.msg, r.speedup)
    assert abs(fast_p_at_k([0.5], 2) - 0.75) < 1e-12
    print("kbench 自检通过。")


if __name__ == "__main__":
    print("== kbench 自检 ==")
    _selftest()