第 7 章　用 Pallas 写自定义内核

# 公共准备：导入、版本检查、对拍与计时工具
import sys, time, re, logging, functools, inspect, contextlib
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| 后端", jax.default_backend(), jax.devices())
assert jax.__version__.startswith("0.12")
assert jax.default_backend() == "cpu", "本章在 CPU 上用 interpret 模式运行（在 Colab 上请使用 CPU 运行时）"

def f64(a):
    """转成 NumPy float64，用来做高精度参考实现"""
    return np.asarray(a, dtype=np.float64)

def max_err(a, ref):
    """最大绝对误差（都转成 float64 后再比较）"""
    return float(np.max(np.abs(f64(a) - f64(ref))))

def bench(f, *args, n=5):
    """平均耗时（毫秒）：先预热一次，再循环 n 次。计时结论只打印，不做 assert"""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n * 1e3

@contextlib.contextmanager
def quiet_callback_logs():
    """TPU 模拟器内部用 io_callback 报错时，会先往日志里打一整段 traceback；演示 ❌ 写法时把它静音"""
    lg = logging.getLogger("jax._src.callback")
    old = lg.disabled
    lg.disabled = True
    try:
        yield
    finally:
        lg.disabled = old

rng = np.random.default_rng(0)   # 只用于少量 NumPy 侧数据；JAX 随机数一律用 jax.random.key

# 探针 1：导入路径与关键名字
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu            # TPU（Mosaic）专用：VMEM、CompilerParams、InterpretParams……
from jax.experimental.pallas import triton as pltriton      # GPU · Triton 后端
from jax.experimental.pallas import mosaic_gpu as plgpu     # GPU · Mosaic GPU 后端

must_have = ["pallas_call", "BlockSpec", "GridSpec", "program_id", "num_programs", "when",
             "cdiv", "ds", "CompilerParams", "CostEstimate", "debug_print", "Squeezed"]
for name in must_have:
    assert hasattr(pl, name), name
assert hasattr(pltpu, "InterpretParams") and hasattr(pltpu, "CompilerParams") and hasattr(pltpu, "VMEM")
assert hasattr(plgpu, "InterpretGPUParams") and hasattr(pltriton, "CompilerParams")

# 旧教程里的名字：0.12 中已不存在
gone = {"pl.load": hasattr(pl, "load"), "pl.store": hasattr(pl, "store"),
        "pl.InterpretParams": hasattr(pl, "InterpretParams"),
        "pltpu.TPUCompilerParams": hasattr(pltpu, "TPUCompilerParams")}
print("已不存在的名字：", gone)
assert not any(gone.values())

sig = inspect.signature(pl.pallas_call)
print("pallas_call 参数：", list(sig.parameters))
print("interpret 默认值：", sig.parameters["interpret"].default)
assert sig.parameters["interpret"].default is False
assert {"out_shape", "grid", "in_specs", "out_specs", "scratch_shapes", "compiler_params",
        "input_output_aliases", "cost_estimate", "interpret", "name"} <= set(sig.parameters)

# 探针 2：interpret 参数到底接受什么？三种写法都在 CPU 上实测
def add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]

x = jnp.arange(8, dtype=jnp.float32)
y = jnp.full((8,), 10.0, jnp.float32)
sds = jax.ShapeDtypeStruct(x.shape, x.dtype)

results = {}
for label, interp in [("interpret=True", True),
                      ("interpret=pltpu.InterpretParams()", pltpu.InterpretParams()),
                      ("interpret=plgpu.InterpretGPUParams()", plgpu.InterpretGPUParams())]:
    out = pl.pallas_call(add_kernel, out_shape=sds, interpret=interp)(x, y)
    results[label] = np.asarray(out)
    print(f"{label:40s} -> {results[label]}")
    np.testing.assert_array_equal(results[label], f64(x) + f64(y))

# ❌ 错误写法：在 CPU 后端上不开 interpret —— CPU 没有 Pallas 的编译后端
try:
    pl.pallas_call(add_kernel, out_shape=sds)(x, y)
    err = None
except Exception as e:
    err = e
    print(type(e).__name__, ":", e)
assert isinstance(err, ValueError) and "Only interpret mode is supported on CPU" in str(err)

# ✅ 正确写法：本机统一用一个全局开关；换到 GPU/TPU 上时只需改成 False
INTERPRET = True
out = pl.pallas_call(add_kernel, out_shape=sds, interpret=INTERPRET)(x, y)
np.testing.assert_array_equal(out, f64(x) + f64(y))
print(out)

# ❌ 错误写法：照旧教程用 pl.load / pl.store 读写 Ref —— 0.12 中这两个函数已从 pl 命名空间移除
def add_kernel_old(x_ref, y_ref, o_ref):
    pl.store(o_ref, ..., pl.load(x_ref, ...) + pl.load(y_ref, ...))
try:
    pl.pallas_call(add_kernel_old, out_shape=sds, interpret=INTERPRET)(x, y)
    err = None
except AttributeError as e:
    err = e
    print(type(e).__name__, ":", e)
assert isinstance(err, AttributeError)

# ✅ 正确写法：直接对 Ref 做 NumPy 风格的索引读写；要取一段动态切片就用 pl.ds(start, size)
def add_kernel_halves(x_ref, y_ref, o_ref):
    o_ref[pl.ds(0, 4)] = x_ref[pl.ds(0, 4)] + y_ref[pl.ds(0, 4)]   # 前半段
    o_ref[4:] = x_ref[4:] + y_ref[4:]                              # 后半段：静态切片也可以
out = pl.pallas_call(add_kernel_halves, out_shape=sds, interpret=INTERPRET)(x, y)
np.testing.assert_array_equal(out, f64(x) + f64(y))
print(out)

# 1) jaxpr 里的 pallas_call 原语
add_call = pl.pallas_call(add_kernel, out_shape=sds, interpret=INTERPRET)
cj = jax.make_jaxpr(add_call)(x, y)
prims = [str(e.primitive) for e in cj.jaxpr.eqns]
print("外层 jaxpr 的原语：", prims)
inner = cj.jaxpr.eqns[0].params["jaxpr"]          # pallas_call 里嵌着 kernel 自己的 jaxpr
print("kernel 的 jaxpr：\n", inner)
assert prims == ["pallas_call"]
assert all("Ref" in str(v.aval) for v in inner.invars)   # kernel 的参数全是 Ref

# 2) interpret 模式：编译后的 HLO 里只有普通操作和 while 循环，没有 custom-call
hlo_interp = jax.jit(add_call).lower(x, y).compile().as_text()
print("interpret 模式 HLO 中 custom-call 的数量：", hlo_interp.count("custom-call"))
assert "custom-call" not in hlo_interp

# 3) 交叉 lowering 到 TPU（不执行！只看生成的 StableHLO）
add_tpu = pl.pallas_call(add_kernel, out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32))  # 不开 interpret
tpu_text = jax.jit(add_tpu).trace(jax.ShapeDtypeStruct((8, 128), jnp.float32),
                                  jax.ShapeDtypeStruct((8, 128), jnp.float32)).lower(lowering_platforms=("tpu",)).as_text()
line = [l.strip() for l in tpu_text.splitlines() if "tpu_custom_call" in l][0]
print("TPU lowering 结果：", line[:110], "...")
assert "tpu_custom_call" in tpu_text

# 多输出：out_shape 是一个元组，kernel 就多接一个输出 Ref
def sum_diff_kernel(x_ref, y_ref, s_ref, d_ref):
    x, y = x_ref[...], y_ref[...]          # 先读出来，避免重复读
    s_ref[...] = x + y
    d_ref[...] = x - y

kx, ky = jax.random.split(jax.random.key(0))
a = jax.random.normal(kx, (4, 6))
b = jax.random.normal(ky, (4, 6))
s, d = pl.pallas_call(sum_diff_kernel,
                      out_shape=(jax.ShapeDtypeStruct(a.shape, a.dtype),
                                 jax.ShapeDtypeStruct(a.shape, a.dtype)),
                      interpret=INTERPRET)(a, b)
np.testing.assert_allclose(f64(s), f64(a) + f64(b), rtol=1e-6)
np.testing.assert_allclose(f64(d), f64(a) - f64(b), rtol=1e-6)
print("多输出 kernel 正确：", s.shape, d.shape)

# ❌ 错误写法：kernel 返回结果（像普通函数那样），而不是写进输出 Ref
def double_kernel_bad(x_ref, o_ref):
    return x_ref[...] * 2
try:
    pl.pallas_call(double_kernel_bad, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a)
    err = None
except ValueError as e:
    err = e
    print(type(e).__name__, ": … kernel 函数", re.search(r"should return None.*", str(e)).group(0))
assert err is not None and "should return None" in str(err)

# ✅ 正确写法：写进 o_ref，函数本身不返回任何东西
def double_kernel(x_ref, o_ref):
    o_ref[...] = x_ref[...] * 2
out = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a)
np.testing.assert_allclose(f64(out), 2 * f64(a), rtol=1e-7)
print("✅", out.shape)

# ❌ 错误写法：kernel 通过闭包捕获一个 JAX 数组（比如权重）
W_bias = jnp.arange(6, dtype=jnp.float32)
def bias_kernel_bad(x_ref, o_ref):
    o_ref[...] = x_ref[...] + W_bias          # W_bias 是外部的数组常量
try:
    pl.pallas_call(bias_kernel_bad, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a)
    err = None
except ValueError as e:
    err = e
    print(type(e).__name__, ": … kernel 函数", re.search(r"captures constants.*", str(e)).group(0))
assert err is not None and "captures constants" in str(err)

# ✅ 正确写法：数组一律作为输入传进来（Python 标量可以闭包捕获，它会被当成编译期常数）
scale = 0.5                                   # Python float：允许
def bias_kernel(x_ref, w_ref, o_ref):
    o_ref[...] = x_ref[...] * scale + w_ref[...]
out = pl.pallas_call(bias_kernel, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a, W_bias)
np.testing.assert_allclose(f64(out), f64(a) * 0.5 + f64(W_bias), rtol=1e-6)
print("✅ 权重作为参数传入，结果正确")

def add2d(x, y, bm=64, bn=128):
    """二维逐元素加法：每个 program 处理一个 (bm, bn) 块"""
    M, N = x.shape
    spec = pl.BlockSpec((bm, bn), lambda i, j: (i, j))     # 第 (i, j) 个 program 取第 (i, j) 个块
    return pl.pallas_call(add_kernel,                        # kernel 与 §0 相同：它只关心“一个块”
                          out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
                          grid=(M // bm, N // bn),
                          in_specs=[spec, spec], out_specs=spec,
                          interpret=INTERPRET)(x, y)

kx, ky = jax.random.split(jax.random.key(1))
X = jax.random.normal(kx, (256, 512))
Y = jax.random.normal(ky, (256, 512))
Z = add2d(X, Y)
np.testing.assert_allclose(f64(Z), f64(X) + f64(Y), rtol=1e-6)
print("grid=(4, 4) 的分块加法正确，输出形状", Z.shape)

# 让每个 program 把自己的 (program_id(0), program_id(1)) 写进它负责的块，直观看到 grid → 块的映射
def pid_kernel(o_ref):
    i, j = pl.program_id(0), pl.program_id(1)
    o_ref[...] = jnp.full(o_ref.shape, 10 * i + j, o_ref.dtype)   # 十位是 i，个位是 j

pid_map = pl.pallas_call(pid_kernel, out_shape=jax.ShapeDtypeStruct((4, 6), jnp.int32),
                         grid=(2, 3), out_specs=pl.BlockSpec((2, 2), lambda i, j: (i, j)),
                         interpret=INTERPRET)()
print(pid_map)
expected = np.kron(np.array([[0, 1, 2], [10, 11, 12]]), np.ones((2, 2), int))
np.testing.assert_array_equal(pid_map, expected)

# ❌ 错误写法：index_map 返回“元素偏移”而不是“块下标”—— 不报错，静默算错
xs = jnp.arange(64, dtype=jnp.float32)
bad = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(xs.shape, xs.dtype), grid=(4,),
                     in_specs=[pl.BlockSpec((16,), lambda i: (i * 16,))],    # 错：i*16 是元素偏移
                     out_specs=pl.BlockSpec((16,), lambda i: (i,)),
                     interpret=INTERPRET)(xs)
wrong = ~np.isclose(f64(bad), 2 * f64(xs))
print("错误元素个数：", int(wrong.sum()), "/ 64；第 16~19 个输出：", np.asarray(bad[16:20]))
assert wrong.sum() > 0      # 越界的块下标在 interpret 模式下被夹到最后一块，结果悄悄出错

# ✅ 正确写法：index_map 返回块下标，偏移由 Pallas 乘上 block_shape
good = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(xs.shape, xs.dtype), grid=(4,),
                      in_specs=[pl.BlockSpec((16,), lambda i: (i,))],
                      out_specs=pl.BlockSpec((16,), lambda i: (i,)),
                      interpret=INTERPRET)(xs)
np.testing.assert_array_equal(good, 2 * f64(xs))
print("✅", np.asarray(good[16:20]))

def matmul_rowcol(a, b, bm=64, bn=128):
    """不沿 K 维切块的 matmul：A 取 (bm, K) 行块，B 取 (K, bn) 列块"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref):
        o_ref[...] = jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
                          grid=(M // bm, N // bn),
                          in_specs=[pl.BlockSpec((bm, K), lambda i, j: (i, 0)),   # 第 i 个行块，K 维只有一块
                                    pl.BlockSpec((K, bn), lambda i, j: (0, j))],  # 第 j 个列块
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j: (i, j)),
                          interpret=INTERPRET)(a, b)

ka, kb = jax.random.split(jax.random.key(2))
A = jax.random.normal(ka, (256, 384))
B = jax.random.normal(kb, (384, 512))
C_rc = matmul_rowcol(A, B)
ref_AB = f64(A) @ f64(B)                       # NumPy float64 参考
print("最大绝对误差：", max_err(C_rc, ref_AB))
np.testing.assert_allclose(f64(C_rc), ref_AB, rtol=1e-4, atol=1e-4)

visit_order = []
def recorder(token, grid_point, core_id):
    visit_order.append(tuple(int(c) for c in grid_point))
    return token                               # 必须把 token 原样返回，否则报 "Mismatched number of outputs"

def copy_kernel(x_ref, o_ref):
    o_ref[...] = x_ref[...]

spec = pl.BlockSpec((8, 128), lambda i, j: (i, j))
xg = jax.random.normal(jax.random.key(3), (16, 384))
yg = pl.pallas_call(copy_kernel, out_shape=jax.ShapeDtypeStruct(xg.shape, xg.dtype), grid=(2, 3),
                    in_specs=[spec], out_specs=spec,
                    interpret=pltpu.InterpretParams(grid_point_recorder=recorder))(xg)
jax.block_until_ready(yg)
print("访问顺序：", visit_order)
np.testing.assert_array_equal(yg, xg)
assert visit_order == [(i, j) for i in range(2) for j in range(3)]     # 字典序

x100 = jnp.arange(1, 101, dtype=jnp.float32)          # 长度 100，块大小 32 → 4 个块，最后一块只有 4 个有效元素
peek = pl.pallas_call(copy_kernel, out_shape=jax.ShapeDtypeStruct((4, 32), jnp.float32), grid=(pl.cdiv(100, 32),),
                      in_specs=[pl.BlockSpec((32,), lambda i: (i,))],
                      out_specs=pl.BlockSpec((None, 32), lambda i: (i, 0)),   # None：这一维被挤掉（squeeze），kernel 看到的是 (32,)
                      interpret=INTERPRET)(x100)
tail = np.asarray(peek[3])
print("第 4 个块读到的内容：", tail[:8], "...")
assert np.array_equal(tail[:4], [97, 98, 99, 100]) and np.isnan(tail[4:]).all()

# 但逐元素 kernel 不受影响：越界部分写出时被丢弃
dbl = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(x100.shape, x100.dtype), grid=(pl.cdiv(100, 32),),
                     in_specs=[pl.BlockSpec((32,), lambda i: (i,))], out_specs=pl.BlockSpec((32,), lambda i: (i,)),
                     interpret=INTERPRET)(x100)
np.testing.assert_array_equal(dbl, 2 * f64(x100))
print("逐元素 kernel 在非整除长度上仍然正确")

kx, ky = jax.random.split(jax.random.key(4))
Xb = jax.random.normal(kx, (1024, 1024))
Yb = jax.random.normal(ky, (1024, 1024))
add_jnp = jax.jit(lambda x, y: x + y)
add_pl_1 = jax.jit(functools.partial(add2d, bm=1024, bn=1024))   # 1 个块
add_pl_64 = jax.jit(functools.partial(add2d, bm=128, bn=128))    # 64 个块
np.testing.assert_allclose(f64(add_pl_64(Xb, Yb)), f64(Xb) + f64(Yb), rtol=1e-6)
print(f"jnp + XLA            : {bench(add_jnp, Xb, Yb):8.3f} ms")
print(f"Pallas interpret 1 块 : {bench(add_pl_1, Xb, Yb):8.3f} ms")
print(f"Pallas interpret 64 块: {bench(add_pl_64, Xb, Yb):8.3f} ms   ← 块越多，解释循环越长（只打印，不作结论）")

def matmul_acc(a, b, bm=64, bk=128, bn=128):
    """沿 K 维累加的分块 matmul，输出 float32，直接用 o_ref 当累加器"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref):
        @pl.when(pl.program_id(2) == 0)            # 每个输出块的第一个 K 步：清零
        def _init():
            o_ref[...] = jnp.zeros_like(o_ref)
        o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
                          grid=(M // bm, N // bn, K // bk),       # 归约维 k 放最后
                          in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
                          interpret=INTERPRET)(a, b)

C_acc = matmul_acc(A, B)                    # A: (256, 384), B: (384, 512) → grid (4, 4, 3)
print("最大绝对误差（对 float64）：", max_err(C_acc, ref_AB))
np.testing.assert_allclose(f64(C_acc), ref_AB, rtol=1e-4, atol=1e-4)

# ❌ 错误写法：用 Python 的 if 判断 program_id —— 它是追踪期的 tracer，不能转成 bool
def mm_kernel_pyif(a_ref, b_ref, o_ref):
    if pl.program_id(2) == 0:
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
try:
    pl.pallas_call(mm_kernel_pyif, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
                   in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)),
                             pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
                   out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=INTERPRET)(A, B)
    err = None
except jax.errors.TracerBoolConversionError as e:
    err = e
    print(type(e).__name__, ":", str(e).splitlines()[0])
assert isinstance(err, jax.errors.TracerBoolConversionError)

# ✅ 正确写法：用 @pl.when(条件) 装饰一个无参函数（matmul_acc 里就是这么写的）
def mm_kernel_when(a_ref, b_ref, o_ref):
    @pl.when(pl.program_id(2) == 0)
    def _init():
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
C_when = pl.pallas_call(mm_kernel_when, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
                        in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)),
                                  pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
                        out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=INTERPRET)(A, B)
np.testing.assert_allclose(f64(C_when), ref_AB, rtol=1e-4, atol=1e-4)
print("✅ pl.when 版本正确，最大误差", max_err(C_when, ref_AB))

# ❌ 错误写法：忘记在 k == 0 时清零，直接 +=
def mm_kernel_noinit(a_ref, b_ref, o_ref):
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
mm_noinit = lambda interp: pl.pallas_call(
    mm_kernel_noinit, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
    in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)), pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
    out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=interp)(A, B)
try:
    r_nan = mm_noinit(INTERPRET)                                          # 未初始化内存 = NaN
    r_zero = mm_noinit(pltpu.InterpretParams(uninitialized_memory="zero")) # 即使初始内存恰好是 0……
    print("interpret=True：NaN 比例 =", float(np.isnan(f64(r_nan)).mean()))
    print("TPU 模拟器（未初始化内存填 0）：最大误差 =", max_err(r_zero, ref_AB))
except Exception as e:
    print(type(e).__name__, e)
    raise
assert np.isnan(f64(r_nan)).all()
assert max_err(r_zero, ref_AB) > 1.0     # 输出缓冲区在块之间被复用，上一个块的结果被累加了进来

# ✅ 正确写法：k == 0 时显式清零（matmul_acc）。在两种解释器、两种未初始化填充值下都正确
for label, interp in [("interpret=True", INTERPRET), ("InterpretParams()", pltpu.InterpretParams()),
                      ('InterpretParams(uninitialized_memory="zero")', pltpu.InterpretParams(uninitialized_memory="zero"))]:
    C_ok = pl.pallas_call(mm_kernel_when, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
                          in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=interp)(A, B)
    np.testing.assert_allclose(f64(C_ok), ref_AB, rtol=1e-4, atol=1e-4)
    print(f"✅ {label:46s} 最大误差 {max_err(C_ok, ref_AB):.2e}")

def matmul_bf16(a, b, bm=64, bk=128, bn=128):
    """bf16 输入/输出，float32 scratch 累加"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref, acc_ref):         # 最后一个参数是 scratch
        k = pl.program_id(2)
        @pl.when(k == 0)
        def _init():
            acc_ref[...] = jnp.zeros_like(acc_ref)
        acc_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
        @pl.when(k == pl.num_programs(2) - 1)        # 最后一个 K 步：写出
        def _store():
            o_ref[...] = acc_ref[...].astype(o_ref.dtype)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), a.dtype),
                          grid=(M // bm, N // bn, K // bk),
                          in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
                          scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],
                          interpret=INTERPRET)(a, b)

A16, B16 = A.astype(jnp.bfloat16), B.astype(jnp.bfloat16)
C16 = matmul_bf16(A16, B16)
ref16 = f64(A16) @ f64(B16)                # 参考：同样的 bf16 输入，float64 计算
rel = np.max(np.abs(f64(C16) - ref16) / (np.abs(ref16) + 1.0))
print("输出 dtype：", C16.dtype, "| 最大相对误差：", rel, "（bf16 的舍入单位约为 2^-8 ≈ 3.9e-3）")
assert C16.dtype == jnp.bfloat16
assert rel < 2 ** -7

def compute_dtype(dt):
    """块内计算精度：bf16/f16 提升到 float32；float32、float64 保持不变（直接写 astype(float32) 会把 float64 降精度）"""
    return jnp.promote_types(dt, jnp.float32)

def softmax_kernel(x_ref, o_ref):
    x = x_ref[...].astype(compute_dtype(x_ref.dtype))
    m = jnp.max(x, axis=-1, keepdims=True)          # 数值稳定：先减去行最大值
    e = jnp.exp(x - m)
    o_ref[...] = (e / jnp.sum(e, axis=-1, keepdims=True)).astype(o_ref.dtype)

def row_call(kernel, x, *extra, br=8):
    """按行分块调用：每个块是 br 行 × 整行；extra 是按列广播的参数（形状 (1, C)，所有块共用）"""
    R, Cc = x.shape
    row_spec = pl.BlockSpec((br, Cc), lambda i: (i, 0))
    vec_spec = pl.BlockSpec((1, Cc), lambda i: (0, 0))      # 每个 program 都取同一块（第 0 块）
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
                          grid=(pl.cdiv(R, br),),
                          in_specs=[row_spec] + [vec_spec] * len(extra), out_specs=row_spec,
                          interpret=INTERPRET)(x, *extra)

def softmax_ref64(x):
    x = f64(x); e = np.exp(x - x.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)

xs_sm = jax.random.normal(jax.random.key(5), (50, 256)) * 3       # 50 行：不能被 8 整除
sm = row_call(softmax_kernel, xs_sm)
print("softmax 最大误差：", max_err(sm, softmax_ref64(xs_sm)))
np.testing.assert_allclose(f64(sm), softmax_ref64(xs_sm), rtol=1e-5, atol=1e-7)
np.testing.assert_allclose(f64(sm).sum(-1), 1.0, rtol=1e-5)

# ❌ 错误写法：softmax 不减最大值 —— 输入较大时 exp 上溢为 inf，inf/inf = NaN
def softmax_kernel_naive(x_ref, o_ref):
    e = jnp.exp(x_ref[...])
    o_ref[...] = e / jnp.sum(e, axis=-1, keepdims=True)
big = xs_sm + 100.0                                     # softmax 对平移不变，正确答案不变
try:
    sm_bad = row_call(softmax_kernel_naive, big)
    print("NaN 比例：", float(np.isnan(f64(sm_bad)).mean()))
except Exception as e:
    print(type(e).__name__, e); raise
assert np.isnan(f64(sm_bad)).any()

# ✅ 正确写法：减去行最大值后，平移 100 的结果与原结果一致
sm_big = row_call(softmax_kernel, big)
np.testing.assert_allclose(f64(sm_big), softmax_ref64(xs_sm), rtol=1e-5, atol=1e-7)
print("✅ 平移 100 后结果不变，最大误差", max_err(sm_big, softmax_ref64(xs_sm)))

EPS = 1e-5
def layernorm_kernel(x_ref, g_ref, b_ref, o_ref):
    x = x_ref[...].astype(compute_dtype(x_ref.dtype))
    mu = jnp.mean(x, axis=-1, keepdims=True)
    xc = x - mu
    var = jnp.mean(xc * xc, axis=-1, keepdims=True)        # 两遍法：先中心化，再求平方的均值
    o_ref[...] = (xc * lax.rsqrt(var + EPS) * g_ref[...] + b_ref[...]).astype(o_ref.dtype)

def layernorm_ref64(x, g, b):
    x = f64(x); mu = x.mean(-1, keepdims=True); var = ((x - mu) ** 2).mean(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + EPS) * f64(g) + f64(b)

kx, kg, kb = jax.random.split(jax.random.key(6), 3)
xln = jax.random.normal(kx, (50, 256))
gamma = 1.0 + 0.1 * jax.random.normal(kg, (1, 256))        # 权重作为参数传入
beta = 0.1 * jax.random.normal(kb, (1, 256))
ln = row_call(layernorm_kernel, xln, gamma, beta)
print("LayerNorm 最大误差：", max_err(ln, layernorm_ref64(xln, gamma, beta)))
np.testing.assert_allclose(f64(ln), layernorm_ref64(xln, gamma, beta), rtol=1e-4, atol=1e-5)

# ❌ 错误写法：方差用 E[x²] − E[x]²。数据带有很大的偏移（这里是 1e4）时，float32 发生灾难性抵消
def layernorm_kernel_onepass(x_ref, g_ref, b_ref, o_ref):
    x = x_ref[...]
    mu = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.mean(x * x, axis=-1, keepdims=True) - mu * mu          # 1e8 量级的两个数相减
    o_ref[...] = (x - mu) * lax.rsqrt(var + EPS) * g_ref[...] + b_ref[...]
x_off = xln + 1e4
ref_off = layernorm_ref64(f64(xln) + 1e4, gamma, beta)          # float64 参考
try:
    ln_bad = row_call(layernorm_kernel_onepass, x_off, gamma, beta)
    err_bad = np.nanmax(np.abs(f64(ln_bad) - ref_off)) if not np.isnan(f64(ln_bad)).all() else np.inf
    print("一遍法最大误差：", err_bad, "| NaN 个数：", int(np.isnan(f64(ln_bad)).sum()))
except Exception as e:
    print(type(e).__name__, e); raise
assert err_bad > 0.1 or np.isnan(f64(ln_bad)).any()

# ✅ 正确写法：两遍法（先中心化）
ln_good = row_call(layernorm_kernel, x_off, gamma, beta)
err_good = max_err(ln_good, ref_off)
print("两遍法最大误差：", err_good)
assert err_good < 5e-3     # 剩余误差来自 1e4 + x 本身在 float32 里的表示误差（约 1e-3 的 ulp）

def matmul_masked(a, b, bm=64, bk=64, bn=64, mask_a=True, mask_b=True):
    """任意形状的分块 matmul：grid 用 cdiv 向上取整，K 方向用掩码"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref):
        k = pl.program_id(2)
        @pl.when(k == 0)
        def _init():
            o_ref[...] = jnp.zeros_like(o_ref)
        a_blk, b_blk = a_ref[...], b_ref[...]
        if mask_a:   # 这里的 if 判断的是 Python 布尔值（构建期常量），不是 tracer，可以用
            ka = k * bk + lax.broadcasted_iota(jnp.int32, a_blk.shape, 1)   # A 块的列 = 全局 K 下标
            a_blk = jnp.where(ka < K, a_blk, 0.0)
        if mask_b:
            kb_ = k * bk + lax.broadcasted_iota(jnp.int32, b_blk.shape, 0)  # B 块的行 = 全局 K 下标
            b_blk = jnp.where(kb_ < K, b_blk, 0.0)
        o_ref[...] += jnp.dot(a_blk, b_blk, preferred_element_type=jnp.float32)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
                          grid=(pl.cdiv(M, bm), pl.cdiv(N, bn), pl.cdiv(K, bk)),
                          in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
                          interpret=INTERPRET)(a, b)

ka, kb = jax.random.split(jax.random.key(7))
Ao = jax.random.normal(ka, (100, 200))        # M=100, K=200, N=130：都不能被 64 整除
Bo = jax.random.normal(kb, (200, 130))
ref_o = f64(Ao) @ f64(Bo)

# ❌ 错误写法：K 方向不掩码，或者只掩码 A —— NaN 污染所有输出
for label, ma, mb in [("不掩码", False, False), ("只掩码 A", True, False)]:
    try:
        r = matmul_masked(Ao, Bo, mask_a=ma, mask_b=mb)
        print(f"{label:8s}：NaN 比例 = {float(np.isnan(f64(r)).mean()):.2f}")
    except Exception as e:
        print(type(e).__name__, e); raise
    assert np.isnan(f64(r)).all()

# ✅ 正确写法：A 的列和 B 的行都用 jnp.where 掩码
C_o = matmul_masked(Ao, Bo)
print("形状", C_o.shape, "| 最大误差：", max_err(C_o, ref_o))
np.testing.assert_allclose(f64(C_o), ref_o, rtol=1e-4, atol=1e-4)

def rowsum(x, br=32, bc=128, reduce_last=True, interpret=INTERPRET):
    """按行求和：x (R, C) -> (R,)；reduce_last=False 用来演示错误的 grid 顺序"""
    R, Cc = x.shape
    red_axis = 1 if reduce_last else 0
    def kernel(x_ref, o_ref):
        j = pl.program_id(red_axis)
        @pl.when(j == 0)
        def _init():
            o_ref[...] = jnp.zeros_like(o_ref)
        col = j * bc + lax.broadcasted_iota(jnp.int32, x_ref.shape, 1)
        o_ref[...] += jnp.sum(jnp.where(col < Cc, x_ref[...], 0.0), axis=1)   # 掩码后再求和
    if reduce_last:
        grid, xmap, omap = (pl.cdiv(R, br), pl.cdiv(Cc, bc)), (lambda i, j: (i, j)), (lambda i, j: (i,))
    else:   # 错误顺序：归约维放在最前面
        grid, xmap, omap = (pl.cdiv(Cc, bc), pl.cdiv(R, br)), (lambda j, i: (i, j)), (lambda j, i: (i,))
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((R,), jnp.float32), grid=grid,
                          in_specs=[pl.BlockSpec((br, bc), xmap)], out_specs=pl.BlockSpec((br,), omap),
                          compiler_params=pltpu.CompilerParams(dimension_semantics=("parallel", "arbitrary")) if reduce_last else None,
                          interpret=interpret)(x)

xr = jax.random.normal(jax.random.key(8), (96, 300))     # 300 列不能被 128 整除
rs = rowsum(xr)
print("行求和最大误差：", max_err(rs, f64(xr).sum(1)))
np.testing.assert_allclose(f64(rs), f64(xr).sum(1), rtol=1e-5, atol=1e-4)

# ❌ 错误写法：把归约维放在 grid 的第一维 —— 同一个输出块被“不连续”地访问
r_true = rowsum(xr, reduce_last=False, interpret=True)     # 通用解释器：碰巧算对，掩盖了问题
print("interpret=True 下误差：", max_err(r_true, f64(xr).sum(1)), "← 碰巧正确")
assert max_err(r_true, f64(xr).sum(1)) < 1e-3
try:
    with quiet_callback_logs():
        jax.block_until_ready(rowsum(xr, reduce_last=False, interpret=pltpu.InterpretParams()))
    err = None
except Exception as e:
    err = e
    msg = [l for l in str(e).splitlines() if "Revisited" in l][0]
    print("TPU 模拟器：", type(e).__name__, "→", msg.strip())
finally:
    pltpu.reset_tpu_interpret_mode_state()                  # 模拟器抛异常后必须重置
assert err is not None and "Revisited block" in str(err)

# ✅ 正确写法：归约维放最后，TPU 模拟器也能通过
rs_tpu = rowsum(xr, reduce_last=True, interpret=pltpu.InterpretParams())
np.testing.assert_allclose(f64(rs_tpu), f64(xr).sum(1), rtol=1e-5, atol=1e-4)
print("✅ TPU 模拟器下误差：", max_err(rs_tpu, f64(xr).sum(1)))

def logsumexp_rows(x, br=16, bc=128):
    """按行在线 logsumexp：x (R, C) -> (R,)，C 方向分块、跨块在线更新"""
    R, Cc = x.shape
    def kernel(x_ref, o_ref, m_ref, l_ref):
        j = pl.program_id(1)
        @pl.when(j == 0)
        def _init():
            m_ref[...] = jnp.full_like(m_ref, -jnp.inf)
            l_ref[...] = jnp.zeros_like(l_ref)
        col = j * bc + lax.broadcasted_iota(jnp.int32, x_ref.shape, 1)
        xb = jnp.where(col < Cc, x_ref[...], -jnp.inf)         # 求最大值的单位元是 -inf
        m_old = m_ref[...]
        m_new = jnp.maximum(m_old, jnp.max(xb, axis=1))
        l_ref[...] = l_ref[...] * jnp.exp(m_old - m_new) + jnp.sum(jnp.exp(xb - m_new[:, None]), axis=1)
        m_ref[...] = m_new
        @pl.when(j == pl.num_programs(1) - 1)
        def _store():
            o_ref[...] = m_ref[...] + jnp.log(l_ref[...])
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((R,), jnp.float32),
                          grid=(pl.cdiv(R, br), pl.cdiv(Cc, bc)),
                          in_specs=[pl.BlockSpec((br, bc), lambda i, j: (i, j))],
                          out_specs=pl.BlockSpec((br,), lambda i, j: (i,)),
                          scratch_shapes=[pltpu.VMEM((br,), jnp.float32), pltpu.VMEM((br,), jnp.float32)],
                          interpret=INTERPRET)(x)

from scipy.special import logsumexp as sp_logsumexp
xl = jax.random.normal(jax.random.key(9), (40, 1000)) * 4 + 50.0   # 行很长、数值偏大；1000 列不能被 128 整除
lse = logsumexp_rows(xl)
ref_lse = sp_logsumexp(f64(xl), axis=1)
print("在线 logsumexp 最大误差：", max_err(lse, ref_lse))
np.testing.assert_allclose(f64(lse), ref_lse, rtol=1e-5)
np.testing.assert_allclose(f64(lse), f64(jax.nn.logsumexp(xl, axis=1)), rtol=1e-5)