第 2 章　JAX 代码性能优化

import os, re, sys, time, gzip, glob, json, logging, tempfile, textwrap, subprocess, functools
from collections import Counter

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
import jax.extend as je

import jaxlib
print("JAX", jax.__version__, "| jaxlib", jaxlib.__version__, "| backend:", jax.default_backend(), "|", jax.devices())

def bench(f, *args, n=20):
    """先预热（含编译），再测 n 次的平均耗时（秒）；每次都 block_until_ready。"""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        jax.block_until_ready(f(*args))
    return (time.perf_counter() - t0) / n

_TABLE_HDR = re.compile(r'^(FileNames|FunctionNames|FileLocations|StackFrames)$')
def clean_hlo(text):
    """去掉 HLO 里的源码位置表、metadata 和空行，只留下计算本身。"""
    out = []
    for line in text.splitlines():
        if _TABLE_HDR.match(line) or re.match(r'^\d+ ', line) or not line.strip():
            continue
        out.append(re.sub(r', metadata=\{[^}]*\}', '', line))
    return "\n".join(out)

def count_fusions(hlo_text):
    """统计优化后 HLO 中 fusion 指令（= 融合后的“内核”）的条数。"""
    return len(re.findall(r' fusion\(', hlo_text))

def compiled_hlo(f, *args, **jit_kw):
    """返回 jit(f) 针对 args 编译后的 HLO 文本。"""
    return jax.jit(f, **jit_kw).lower(*args).compile().as_text()

class CompileCounter(logging.Handler):
    """with CompileCounter("f") as cc: ...  之后 cc.count 就是 jit(f) 真正编译的次数。"""
    def __init__(self, name=None):
        super().__init__()
        self.name, self.records = name, []
    def emit(self, record):
        msg = record.getMessage()
        if msg.startswith("Compiling jit("):
            fn = msg[len("Compiling jit("):].split(")")[0]
            if self.name is None or fn == self.name:
                self.records.append(msg)
    def __enter__(self):
        self._logger = logging.getLogger("jax")
        self._logger.addHandler(self)
        # 暂时调高 jax 自带 StreamHandler 的级别，避免 log_compiles 刷屏
        self._saved = [(h, h.level) for h in self._logger.handlers if h is not self]
        for h, _ in self._saved:
            h.setLevel(logging.ERROR)
        self._ctx = jax.log_compiles(True)
        self._ctx.__enter__()
        return self
    def __exit__(self, *exc):
        self._ctx.__exit__(*exc)
        for h, lvl in self._saved:
            h.setLevel(lvl)
        self._logger.removeHandler(self)
    @property
    def count(self):
        return len(self.records)

SCRATCH = tempfile.mkdtemp(prefix="adv02_")
def run_py(src, env_extra=None, args=()):
    """把 src 写成脚本，在全新的 Python 子进程里运行，返回 stdout+stderr 文本。"""
    path = os.path.join(SCRATCH, f"sub_{abs(hash(src)) % 10**8}.py")
    with open(path, "w") as fh:
        fh.write(textwrap.dedent(src))
    env = dict(os.environ)
    env.update(env_extra or {})
    p = subprocess.run([sys.executable, path, *args], capture_output=True, text=True, env=env, timeout=300)
    if p.returncode != 0:
        raise RuntimeError(p.stdout + p.stderr)
    return p.stdout + p.stderr

print("临时目录:", SCRATCH)

# ❌ 旧写法（过时）：jax.random.PRNGKey 返回“原始 uint32 数组”形式的 key，0.12 推荐类型化 key
old_key = jax.random.PRNGKey(0)
new_key = jax.random.key(0)
print("PRNGKey(0):", old_key, "| dtype =", old_key.dtype, "| shape =", old_key.shape)
print("key(0)    :", new_key.dtype, "| shape =", new_key.shape)
try:
    # 类型化 key 不能被当成普通整数数组做算术，这能防止把 key 误当数据用
    _ = new_key + 1
    print("类型化 key 居然可以做加法？")
except Exception as e:
    print("类型化 key 做算术会报错：", type(e).__name__)
assert old_key.dtype == jnp.uint32 and old_key.shape == (2,)
assert jnp.issubdtype(new_key.dtype, jax.dtypes.prng_key) and new_key.shape == ()

# ✅ 推荐写法：jax.random.key + split，x、y 用不同的子 key
@jax.jit
def complex_computation(x, y):
    z = jnp.dot(x, y)
    z = jnp.sin(z)
    return jnp.mean(z * x + y)

kx, ky = jax.random.split(jax.random.key(0))
x = jax.device_put(jax.random.normal(kx, (2000, 2000)))
y = jax.device_put(jax.random.normal(ky, (2000, 2000)))

# 预热：触发编译，并等它真正完成
complex_computation(x, y).block_until_ready()

log_dir = tempfile.mkdtemp(prefix="jax_profiling_")   # 原示例用 /tmp/jax_profiling
jax.profiler.start_trace(log_dir)
for _ in range(5):
    result = complex_computation(x, y)
    result.block_until_ready()   # 重要：阻塞，确保计算在停止分析前完成
jax.profiler.stop_trace()
print("trace 已保存到:", log_dir)

files = sorted(p for p in glob.glob(log_dir + "/**", recursive=True) if os.path.isfile(p))
for p in files:
    print(f"  {os.path.relpath(p, log_dir):70s} {os.path.getsize(p):>8d} 字节")
assert any(p.endswith(".xplane.pb") for p in files)
assert any(p.endswith(".trace.json.gz") for p in files)

@jax.jit
def annotated_step(x, y):
    with jax.named_scope("my_matmul"):        # 作用于 jaxpr/HLO 中的操作名
        z = x @ y
    with jax.named_scope("my_epilogue"):
        out = jnp.mean(jnp.sin(z))
    return out

a = jax.random.normal(kx, (512, 512)); b = jax.random.normal(ky, (512, 512))
annotated_step(a, b).block_until_ready()       # 预热

log_dir2 = tempfile.mkdtemp(prefix="jax_trace_")
with jax.profiler.trace(log_dir2, create_perfetto_trace=True):
    for i in range(3):
        with jax.profiler.TraceAnnotation(f"step_{i}"):   # 作用于宿主端时间线
            annotated_step(a, b).block_until_ready()

files2 = sorted(os.path.relpath(p, log_dir2) for p in glob.glob(log_dir2 + "/**", recursive=True) if os.path.isfile(p))
print("\n".join(files2))
assert any(f.endswith("perfetto_trace.json.gz") for f in files2)

# trace.json.gz 是标准 Chrome trace 格式：{"traceEvents": [...]}
tj = glob.glob(log_dir2 + "/**/*.trace.json.gz", recursive=True)[0]
events = json.load(gzip.open(tj))["traceEvents"]
names = Counter(e["name"] for e in events if e.get("ph") == "X")
print("\n事件数最多的前 12 种：")
for k, v in names.most_common(12):
    print(f"  {v:4d}  {k}")
assert all(names[f"step_{i}"] == 1 for i in range(3))   # TraceAnnotation 写进了 trace
assert names["PjitFunction(annotated_step)"] >= 3       # 每次调用 jit 函数的宿主端事件

# xplane.pb 也可以直接用 jax.profiler.ProfileData 解析（不需要 TensorBoard）
xp = glob.glob(log_dir2 + "/**/*.xplane.pb", recursive=True)[0]
pdata = jax.profiler.ProfileData.from_file(xp)
for plane in pdata.planes:
    lines = list(plane.lines)
    print(f"plane {plane.name!r}: {len(lines)} 条 line")
    for ln in lines:
        if "XLAExecute" in ln.name or ln.name == "":
            evs = [e.name for e in ln.events]
            kern = sorted({n for n in evs if "fusion" in n or n.startswith("wrapped")})
            print(f"   line {ln.name!r:28s} 事件 {len(evs):3d} 个；其中内核名: {kern}")

# named_scope 不改变计算，但会出现在 HLO 的 op_name 元数据里（调试/对照 trace 时很有用）
lowered = annotated_step.lower(a, b)
hlo_ops = re.findall(r'op_name="([^"]*my_[^"]*)"', lowered.compile().as_text())
print("\n带 named_scope 前缀的 HLO op_name：", sorted(set(hlo_ops)))
assert any("my_matmul" in s for s in hlo_ops) and any("my_epilogue" in s for s in hlo_ops)
assert "my_matmul" not in lowered.as_text()                 # 默认 StableHLO 文本不带位置信息
assert "my_matmul" in lowered.as_text(debug_info=True)      # debug_info=True 才会显示

mem_path = os.path.join(SCRATCH, "memory.prof")
blob_before = jax.profiler.device_memory_profile()          # 返回 bytes，不落盘
big = jnp.ones((4000, 4000)); big.block_until_ready()          # 新增一个约 64 MB 的存活缓冲区
jax.profiler.save_device_memory_profile(mem_path)             # 落盘版本
blob_after = open(mem_path, "rb").read()
print("CPU 上可用；文件大小:", os.path.getsize(mem_path), "字节；gzip 魔数:", blob_after[:2])
print("解压后 protobuf 大小：之前", len(gzip.decompress(blob_before)), "→ 之后", len(gzip.decompress(blob_after)))
assert blob_after[:2] == b"\x1f\x8b"     # gzip 格式的 pprof
del big
print("\n查看方法：  pprof -http=:8080", mem_path, "   （或 go tool pprof ...）")

def example_function(x, y):
    a = jnp.sin(x)
    b = jnp.cos(y)
    return a + b * 2.0

kx_, ky_ = jax.random.split(jax.random.key(0))       # 用推荐的 jax.random.key（旧写法是 PRNGKey）
x_example = jax.random.normal(kx_, (10,))
y_example = jax.random.normal(ky_, (10,))

closed = jax.make_jaxpr(example_function)(x_example, y_example)
print(closed.jaxpr)

j = closed.jaxpr
print(type(closed).__name__, "→ .jaxpr 是", type(j).__name__)
print("invars   :", [(str(v), v.aval.str_short()) for v in j.invars])
print("constvars:", j.constvars, "| consts:", closed.consts)
print("outvars  :", [(str(v), v.aval.str_short()) for v in j.outvars])
for eqn in j.eqns:
    print(f"  eqn: {eqn.primitive.name:4s} in={[str(v) for v in eqn.invars]} out={[str(v) for v in eqn.outvars]} params={eqn.params}")

# 0.12：ClosedJaxpr 已并入 Jaxpr，.jaxpr 是返回自身的兼容属性
assert type(closed) is je.core.ClosedJaxpr is type(j) and closed.jaxpr is closed
assert [e.primitive.name for e in j.eqns] == ["sin", "cos", "mul", "add"]
assert len(j.invars) == 2 and len(j.outvars) == 1 and j.constvars == []
# 字面量 2.0 不是变量，而是 Literal
assert type(j.eqns[2].invars[1]).__name__ == "Literal"

W_small = jnp.arange(9.0).reshape(3, 3)     # 小矩阵，只为演示 jaxpr 结构
def uses_closure(v):
    return jnp.dot(W_small, v) + jnp.pi

cj = jax.make_jaxpr(uses_closure)(jnp.ones(3))
print(cj)
print("constvars:", [v.aval.str_short() for v in cj.jaxpr.constvars], "| consts 形状:", [c.shape for c in cj.consts])
assert len(cj.jaxpr.constvars) == 1 and cj.consts[0].shape == (3, 3)
# Python 标量 jnp.pi 则直接成为字面量
assert type(cj.jaxpr.eqns[1].invars[1]).__name__ == "Literal"

def count_primitives(jaxpr, counter=None, depth=0, show=False):
    """递归统计 jaxpr（含所有子 jaxpr）中的原语出现次数。"""
    counter = Counter() if counter is None else counter
    for eqn in jaxpr.eqns:
        counter[eqn.primitive.name] += 1
        subs = list(je.core.jaxprs_in_params(eqn.params))
        if show and subs:
            print("  " * depth + f"└ {eqn.primitive.name}（含 {len(subs)} 个子 jaxpr）")
        for sub in subs:
            count_primitives(sub, counter, depth + 1, show)
    return counter

@jax.jit
def inner(x):
    return jax.nn.relu(x) * 2.0

def outer(x):
    return inner(x).sum() + jnp.logaddexp(x, 1.0).sum()

co = jax.make_jaxpr(outer)(jnp.ones(3))
print("顶层原语:", [e.primitive.name for e in co.jaxpr.eqns])
print("嵌套结构：")
cnt = count_primitives(co.jaxpr, show=True)
print("递归统计:", dict(cnt))

top = co.jaxpr.eqns[0]
assert top.primitive.name == "jit" and top.params["name"] == "inner"
# 0.12 中嵌套 jit 的原语名是 "jit"（旧版本叫 "pjit"）
assert "pjit" not in cnt
# relu 本身是 custom_jvp 函数，其内部又包了一层 jit；logaddexp 也是 custom_jvp
assert cnt["custom_jvp_call"] == 2 and cnt["jit"] == 2 and cnt["max"] >= 1

# ❌ 常见误解：“依赖值的 Python 控制流会让 jaxpr 改变，从而触发重新编译”
@jax.jit
def value_branch(x, flag):
    if flag > 0.5:          # 对 tracer 求 bool
        return jnp.sin(x)
    return jnp.cos(x)

try:
    value_branch(jnp.ones(3), 0.7)
    caught = None
except Exception as e:
    caught = type(e).__name__
    print("实测：", caught, "|", str(e).splitlines()[0])
assert caught == "TracerBoolConversionError"

# ✅ 推荐写法：确实要让“值”决定图结构时，标为静态参数；此时不同的值 → 不同 jaxpr → 各自编译一次
def value_branch_static(x, flag):
    return jnp.sin(x) if flag > 0.5 else jnp.cos(x)

jp_hi = jax.make_jaxpr(value_branch_static, static_argnums=1)(jnp.ones(3), 0.7)
jp_lo = jax.make_jaxpr(value_branch_static, static_argnums=1)(jnp.ones(3), 0.3)
print(jp_hi, jp_lo, sep="\n")
assert [e.primitive.name for e in jp_hi.jaxpr.eqns] == ["sin"]
assert [e.primitive.name for e in jp_lo.jaxpr.eqns] == ["cos"]

jit_static = jax.jit(value_branch_static, static_argnums=1)
with CompileCounter("value_branch_static") as cc:
    for flag in [0.7, 0.7, 0.3, 0.3]:
        jit_static(jnp.ones(3), flag)
print("4 次调用实际编译次数:", cc.count)
assert cc.count == 2

def conditional_function(use_sin, x):
    pred = use_sin > 0.5          # 标量布尔
    def true_fun(v):
        return jnp.sin(v)
    def false_fun(v):
        return jnp.cos(v)
    return lax.cond(pred, true_fun, false_fun, x)

xs3 = jnp.arange(3.0)
jc_07 = jax.make_jaxpr(conditional_function)(0.7, xs3)
jc_03 = jax.make_jaxpr(conditional_function)(0.3, xs3)
print(jc_07)
assert str(jc_07) == str(jc_03)            # 结构与谓词的值无关

cond_eqn = [e for e in jc_07.jaxpr.eqns if e.primitive.name == "cond"][0]
branch_prims = [[e.primitive.name for e in br.jaxpr.eqns] for br in cond_eqn.params["branches"]]
print("branches 顺序:", branch_prims)
prev = jc_07.jaxpr.eqns[1]
print("谓词先被转换:", prev.primitive.name, prev.params["new_dtype"])
assert branch_prims == [["cos"], ["sin"]]   # 第 0 个分支是 false_fun
assert prev.primitive.name == "convert_element_type" and prev.params["new_dtype"] == jnp.int32

def f_grad_demo(x):
    return jnp.sum(jnp.sin(x) ** 2)

gj = jax.make_jaxpr(jax.grad(f_grad_demo))(jnp.ones(4))
print(gj)
prims = count_primitives(gj.jaxpr)
assert prims["sin"] == 1 and prims["cos"] == 1        # 反向用到 d sin = cos

# jaxpr_as_fun：把 ClosedJaxpr 当普通函数调用，结果应与原函数一致
fn_from_jaxpr = je.core.jaxpr_as_fun(gj)
np.testing.assert_allclose(fn_from_jaxpr(jnp.arange(4.0))[0], jax.grad(f_grad_demo)(jnp.arange(4.0)), rtol=1e-6)
print("jaxpr_as_fun 结果与 jax.grad 一致")

def simple_computation(x, a, b):
    y = jnp.sin(x)
    z = a * y
    w = z + b
    return w

xs = jnp.ones((1000, 1000))
lowered = jax.jit(simple_computation).lower(xs, 2.0, 1.0)
stablehlo_txt = lowered.as_text()
print("==== StableHLO（lower 之后、XLA 优化之前）====")
print(stablehlo_txt)

compiled = lowered.compile()
opt_hlo = compiled.as_text()
print("==== 优化后的 HLO（compile 之后）====")
print(clean_hlo(opt_hlo))

assert "stablehlo.sine" in stablehlo_txt and "HloModule" not in stablehlo_txt
assert opt_hlo.startswith("HloModule")
# 三个逐元素操作（sin / mul / add）被融合成一条 fusion 指令
assert count_fusions(opt_hlo) == 1
print("优化后 fusion 条数:", count_fusions(opt_hlo))

ca = compiled.cost_analysis()
ma = compiled.memory_analysis()
print("cost_analysis 类型:", type(ca).__name__)
for k in ["flops", "transcendentals", "bytes accessed"]:
    print(f"  {k:16s} = {ca[k]:,.0f}")
print("memory_analysis:")
for k in ["argument_size_in_bytes", "output_size_in_bytes", "temp_size_in_bytes", "alias_size_in_bytes"]:
    print(f"  {k:24s} = {getattr(ma, k):,}")

n_el = 1000 * 1000
assert ca["flops"] == 2 * n_el            # 一次乘法 + 一次加法
assert ca["transcendentals"] == n_el      # 一次 sin
# 融合后只读 x、写 w（外加两个标量）：两遍 4MB
assert ca["bytes accessed"] == 2 * 4 * n_el + 8
assert ma.output_size_in_bytes == 4 * n_el and ma.temp_size_in_bytes == 0

# 常量折叠：jnp.sin(2.0)、arange(4).sum() 只依赖常量 → 优化后 HLO 里已经算好
fold = lambda v: v * jnp.sin(2.0) + jnp.arange(4.0).sum()
v4 = jnp.ones(4)
low_txt = jax.jit(fold).lower(v4).as_text()
opt_txt = compiled_hlo(fold, v4)
print("StableHLO 里有 sine / reduce:", "stablehlo.sine" in low_txt, "stablehlo.reduce" in low_txt)
print("优化后 HLO 里有 sine / reduce:", " sine(" in opt_txt, " reduce(" in opt_txt)
assert "stablehlo.sine" in low_txt and " sine(" not in opt_txt and " reduce(" not in opt_txt

# 代数简化 ①：x * 1.0 + 0.0 → 直接复制 x
simp = clean_hlo(compiled_hlo(lambda v: v * 1.0 + 0.0, xs))
print("\nx*1.0+0.0 优化后:\n", simp.split("ENTRY")[-1])
assert " multiply(" not in simp and " add(" not in simp

# ❌ 常见误解：以为 XLA 会把 (x + y) - x 化简为 y
alg = clean_hlo(compiled_hlo(lambda p, q: (p + q) - p, xs, xs))
print("\n".join(l for l in alg.splitlines() if " add(" in l or " subtract(" in l))
try:
    assert " subtract(" not in alg, "HLO 里仍然有 subtract"
    print("被化简了")
except AssertionError as e:
    print("实测：没有化简 ——", e)

# ✅ 正确理解：浮点下 (x+y)-x ≠ y（会丢失精度），XLA 默认不做这种不保值的改写
p_, q_ = jnp.float32(1e8), jnp.float32(1.0)
exact_y = q_
via_xla = jax.jit(lambda p, q: (p + q) - p)(p_, q_)
print("(1e8 + 1) - 1e8 =", via_xla, "（数学上应为 1）")
assert " subtract(" in alg and float(via_xla) != float(exact_y)

dump_dir = os.path.join(SCRATCH, "xla_dump")
out = run_py("""
    import jax, jax.numpy as jnp
    def simple_computation(x, y):
        return jnp.exp(jnp.log(x) + y)
    x = jnp.ones((1000, 1000))
    print(float(jax.jit(simple_computation)(x, x).sum()))
""", env_extra={"XLA_FLAGS": f"--xla_dump_to={dump_dir} --xla_dump_hlo_as_text"})
dumped = sorted(os.listdir(dump_dir))
mine = [f for f in dumped if "jit_simple_computation" in f]
print(f"共转储 {len(dumped)} 个文件，其中 jit_simple_computation 相关：")
print("\n".join("  " + f for f in mine))

before = open(os.path.join(dump_dir, [f for f in mine if f.endswith("before_optimizations.txt")][0])).read()
after = open(os.path.join(dump_dir, [f for f in mine if f.endswith("cpu_after_optimizations.txt")][0])).read()
print("\n---- before_optimizations（ENTRY 部分）----\n", clean_hlo(before).split("ENTRY")[-1])
print("---- after_optimizations（ENTRY 部分）----\n", clean_hlo(after).split("ENTRY")[-1])
assert count_fusions(before) == 0 and count_fusions(after) == 1

rng = np.random.default_rng(0)
n = 128
W_np = rng.standard_normal((n, n)).astype(np.float32)
X_np = rng.standard_normal((n, n)).astype(np.float32)
ref = float(np.sum(W_np.astype(np.float64) @ X_np.astype(np.float64)))   # float64 参考值

W_const = jnp.asarray(W_np)
closure_fn = jax.jit(lambda x: jnp.sum(W_const @ x))      # W 作为闭包常量
param_fn   = jax.jit(lambda x, w: jnp.sum(w @ x))         # W 作为参数传入

r_closure = float(closure_fn(jnp.asarray(X_np)))
r_param   = float(param_fn(jnp.asarray(X_np), W_const))
print(f"float64 参考 = {ref:.4f}")
print(f"传参版本     = {r_param:.4f}   误差 {abs(r_param - ref):.2e}")
print(f"闭包版本     = {r_closure:.4f}   误差 {abs(r_closure - ref):.2e}")
print("闭包版本是否触发了 bug:", abs(r_closure - ref) > 1.0)
assert abs(r_param - ref) < 1e-2     # 传参版本永远可信

h_closure = clean_hlo(closure_fn.lower(jnp.asarray(X_np)).compile().as_text())
h_param   = clean_hlo(param_fn.lower(jnp.asarray(X_np), W_const).compile().as_text())
def fused_body(h):
    body = h.split("%fused_computation")[1]
    return body[:body.index("\n}")]
print("==== 闭包版本的 fused_computation ====", fused_body(h_closure), sep="")
print("==== 传参版本的 fused_computation ====", fused_body(h_param), sep="")
print("ENTRY 指令:\n ", [l.strip() for l in h_closure.splitlines() if "ROOT %ynn" in l][0][:120])

assert '"kind":"__ynn_fusion"' in h_closure and '"kind":"__ynn_fusion"' in h_param
# 关键差异：闭包版本把 128×128 的常量直接嵌进了 YNN fusion 内部，并作为 dot 的左操作数
assert re.search(r"constant\.\d+ = f32\[128,128\]\{1,0\} constant\(\{\.\.\.\}\)", fused_body(h_closure))
assert "constant({...})" not in fused_body(h_param)

probe_src = """
    import sys, numpy as np, jax, jax.numpy as jnp
    rng = np.random.default_rng(0); n = 128
    W = rng.standard_normal((n, n)).astype(np.float32); X = rng.standard_normal((n, n)).astype(np.float32)
    Wc = jnp.asarray(W)
    f = jax.jit(lambda x: jnp.sum(Wc @ x))
    hlo = f.lower(jnp.asarray(X)).compile().as_text()
    print(float(f(jnp.asarray(X))), float(np.sum(W.astype(np.float64) @ X.astype(np.float64))), "__ynn_fusion" in hlo)
"""
res_default = run_py(probe_src).split()
res_noynn = run_py(probe_src, env_extra={"XLA_FLAGS": "--xla_cpu_experimental_ynn_fusion_type="}).split()
print("默认设置     : 结果 %s  参考 %s  使用 YNN=%s" % tuple(res_default[-3:]))
print("关闭 YNN 融合: 结果 %s  参考 %s  使用 YNN=%s" % tuple(res_noynn[-3:]))
val, refv, ynn = float(res_noynn[-3]), float(res_noynn[-2]), res_noynn[-1]
assert ynn == "False" and abs(val - refv) < 1e-2      # 关闭 YNN 融合后闭包版本也正确

A = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
A_F = np.asfortranarray(A)
print("行主序线性内存:", A.ravel(order="K"), " strides =", A.strides)
print("列主序线性内存:", A_F.ravel(order="K"), " strides =", A_F.strides)
assert A.ravel(order="K").tolist() == [1, 2, 3, 4, 5, 6]
assert A_F.ravel(order="K").tolist() == [1, 4, 2, 5, 3, 6]
# 偏移公式（以元素为单位）
m_, n_ = A.shape
for i in range(m_):
    for j_ in range(n_):
        assert A.strides[0] // 4 * i + A.strides[1] // 4 * j_ == i * n_ + j_
        assert A_F.strides[0] // 4 * i + A_F.strides[1] // 4 * j_ == i + j_ * m_

M2 = jax.random.normal(jax.random.key(1), (2000, 2000))
print("jax.Array.format:", M2.format.layout)
h_id = compiled_hlo(lambda t: t * 2.0, M2)
print("HLO 入口布局:", re.search(r"entry_computation_layout=\{([^}]*\}[^}]*)\}", h_id).group(1))
assert M2.format.layout.major_to_minor == (0, 1)      # 行主序：第 0 维最“慢”
assert "f32[2000,2000]{1,0}" in h_id                   # HLO 的 minor_to_major 写法

B2 = jax.random.normal(jax.random.key(2), (2000, 2000))
cases = {
    "① 只做转置 a.T":          (lambda a: a.T,            (M2,)),
    "② 转置后接逐元素 a.T*2+1": (lambda a: a.T * 2 + 1,    (M2,)),
    "③ 转置后接矩阵乘 a.T @ b": (lambda a, b: a.T @ b,     (M2, B2)),
    "④ 两次转置 a.T.T":        (lambda a: a.T.T,          (M2,)),
}
hlos = {}
for name, (fn, args) in cases.items():
    jp = jax.make_jaxpr(fn)(*args)
    h = clean_hlo(compiled_hlo(fn, *args))
    hlos[name] = h
    entry = [l.strip()[:110] for l in h.split("ENTRY")[-1].splitlines() if "ROOT" in l][0]
    print(f"{name}\n   jaxpr 中 transpose 次数: {count_primitives(jp.jaxpr)['transpose']}\n   优化后 ROOT: {entry}")

assert " transpose(" in hlos["① 只做转置 a.T"]                       # 真正搬数据（在 fusion 里）
assert count_fusions(hlos["② 转置后接逐元素 a.T*2+1"]) == 1           # 转置与逐元素融合成一个内核
assert " transpose(" not in hlos["③ 转置后接矩阵乘 a.T @ b"]          # 转置被吸收进 dot 的收缩维度
assert "lhs_contracting_dims={0}" in hlos["③ 转置后接矩阵乘 a.T @ b"]
assert " transpose(" not in hlos["④ 两次转置 a.T.T"]                  # 两次转置互相抵消

sum_rows = jax.jit(lambda a: a.sum(axis=1))   # 沿最后一维：连续访问
sum_cols = jax.jit(lambda a: a.sum(axis=0))   # 沿第一维：跨步访问（XLA 可能会重排循环）
copy_c   = jax.jit(lambda a: a * 1.0 + 0.0)   # 纯复制（被化简为 copy）
copy_t   = jax.jit(lambda a: a.T)             # 物化转置

t_rows, t_cols = bench(sum_rows, M2), bench(sum_cols, M2)
t_copy, t_tr = bench(copy_c, M2), bench(copy_t, M2)
print(f"XLA  sum(axis=1) {t_rows*1e3:7.3f} ms | sum(axis=0) {t_cols*1e3:7.3f} ms")
print(f"XLA  复制        {t_copy*1e3:7.3f} ms | 物化转置    {t_tr*1e3:7.3f} ms")

Mn = np.asarray(M2)
t0 = time.perf_counter(); [np.ascontiguousarray(Mn) .copy() for _ in range(10)]; t_np_c = (time.perf_counter()-t0)/10
t0 = time.perf_counter(); [np.ascontiguousarray(Mn.T) for _ in range(10)];          t_np_t = (time.perf_counter()-t0)/10
print(f"NumPy 连续复制    {t_np_c*1e3:7.3f} ms | 转置后变连续 {t_np_t*1e3:7.3f} ms")
print("（XLA 会为 reduce 选择更好的遍历顺序，所以两种 sum 的差距往往比朴素 C 循环小）")

big_c = np.random.default_rng(0).random((3000, 3000), dtype=np.float32)
big_f = np.asfortranarray(big_c)
t_put_c = bench(lambda a: jax.device_put(a), big_c, n=10)
t_put_f = bench(lambda a: jax.device_put(a), big_f, n=10)
print(f"device_put 行主序 {t_put_c*1e3:8.3f} ms | 列主序 {t_put_f*1e3:8.3f} ms  （约 {t_put_f/t_put_c:.0f}×）")
arr_from_f = jnp.asarray(big_f)
assert arr_from_f.format.layout.major_to_minor == (0, 1)     # 进入 JAX 后统一是行主序
assert np.array_equal(np.asarray(arr_from_f), big_c)          # 数值不变

# ❌ 旧写法（代码错误）：被追踪参数上的 Python if 无法追踪，第一次调用就报错，根本谈不上“第三次调用重新编译”
@jax.jit
def example_function(x, size_param):
    # Python 控制流依赖 size_param 的值
    if size_param > 5:
        return jnp.sum(x * 2.0)
    else:
        return jnp.sum(x + 1.0)

data = jax.random.normal(jax.random.key(0), (10,))
try:
    example_function(data, 3).block_until_ready()
    err_a = None
except Exception as e:
    err_a = type(e).__name__
    print("实测：", err_a)
assert err_a == "TracerBoolConversionError"

# ✅ 推荐写法（课程代码块 B）：标记为静态；用 CompileCounter 代替“看计时猜编译”
@jax.jit(static_argnums=(1,))   # 将 size_param（索引 1）标记为静态（0.12 支持这种装饰器工厂写法）
def example_function_static(x, size_param):
    print(f"Compiling for size_param = {size_param}")   # 只在追踪时执行
    if size_param > 5:
        return jnp.sum(x * 2.0)
    else:
        return jnp.sum(x + 1.0)

timings = []
with CompileCounter("example_function_static") as cc:
    for sp in [3, 3, 7, 7]:
        t0 = time.perf_counter()
        example_function_static(data, sp).block_until_ready()
        timings.append(time.perf_counter() - t0)
for sp, t in zip([3, 3, 7, 7], timings):
    print(f"size_param={sp}: {t*1e3:8.3f} ms")
print("实际编译次数:", cc.count)
assert cc.count == 2          # 3 和 7 各编译一次

@jax.jit
def scale(v):
    return v * 2

with CompileCounter("scale") as cc:
    scale(jnp.ones(3, jnp.float32))
    scale(np.ones(3, np.float64))          # x32 模式下被规范化为 float32
    n_after_f64 = cc.count
    scale(1)                               # int32, weak_type=True
    scale(1.0)                             # float32, weak_type=True
    scale(jnp.float32(1.0))                # float32, weak_type=False
for r in cc.records:
    print(" ", r.split(" with global shapes and types ")[1].split(". Argument")[0])
print("float32 数组 + float64 NumPy 数组 → 编译次数:", n_after_f64, "；总编译次数:", cc.count)
assert n_after_f64 == 1 and cc.count == 4
assert not jax.config.jax_enable_x64

@jax.jit
def example_function_lax(x, size_param_val):
    # size_param_val 是 0 维数组（或可被追踪的 Python 标量）
    pred = size_param_val > 5
    def true_fun(operand):
        return jnp.sum(operand * 2.0)
    def false_fun(operand):
        return jnp.sum(operand + 1.0)
    return lax.cond(pred, true_fun, false_fun, x)

size_val_3, size_val_7 = jnp.array(3), jnp.array(7)
with CompileCounter("example_function_lax") as cc:
    r3 = example_function_lax(data, size_val_3).block_until_ready()
    r7 = example_function_lax(data, size_val_7).block_until_ready()
print("编译次数:", cc.count, "| r3 =", r3, "| r7 =", r7)
assert cc.count == 1
np.testing.assert_allclose(r3, np.sum(np.asarray(data) + 1.0), rtol=1e-5)
np.testing.assert_allclose(r7, np.sum(np.asarray(data) * 2.0), rtol=1e-5)

def base(v, a):
    return v * a

v3 = jnp.ones(3)
with CompileCounter("<lambda>") as c_lambda:
    for _ in range(3):
        jax.jit(lambda v: v * 2)(v3)                    # 每次都是新的 lambda 对象
with CompileCounter("base") as c_partial:
    for _ in range(3):
        jax.jit(functools.partial(base, a=2.0))(v3)     # 每次都是新的 partial 对象
with CompileCounter("base") as c_rejit:
    for _ in range(3):
        jax.jit(base)(v3, 2.0)                          # 新的 jit 包装，但被包装的函数是同一个
p_once = jax.jit(functools.partial(base, a=2.0))        # partial 只创建一次
with CompileCounter("base") as c_once:
    for _ in range(3):
        p_once(v3)
print(f"循环内新建 lambda: {c_lambda.count} 次 | 循环内新建 partial: {c_partial.count} 次 | "
      f"循环内 jax.jit(base): {c_rejit.count} 次 | partial 只建一次: {c_once.count} 次")
assert (c_lambda.count, c_partial.count, c_rejit.count, c_once.count) == (3, 3, 1, 1)

@jax.jit
def explain_me(v, w):
    return v @ w

class _Grab(logging.Handler):
    def __init__(self):
        super().__init__(); self.msgs = []
    def emit(self, record):
        self.msgs.append(record.getMessage())

grab = _Grab(); jl = logging.getLogger("jax"); jl.addHandler(grab)
saved = [(h, h.level) for h in jl.handlers if h is not grab]
for h, _ in saved: h.setLevel(logging.ERROR)       # 不刷屏，只收集
jax.config.update("jax_explain_cache_misses", True)
try:
    explain_me(jnp.ones((4, 4)), jnp.ones((4, 4)))
    explain_me(jnp.ones((4, 4)), jnp.ones((4, 8)))          # 只改了 w 的形状
finally:
    jax.config.update("jax_explain_cache_misses", False)
    for h, lvl in saved: h.setLevel(lvl)
    jl.removeHandler(grab)
miss = [m for m in grab.msgs if "CACHE MISS" in m and "explain_me" in m]
print(miss[-1])
assert "at w, now f32[4,8] and before f32[4,4]" in miss[-1]

cache_dir = os.path.join(SCRATCH, "jax_cache")
cache_src = """
    import sys, time, logging
    import jax, jax.numpy as jnp
    jax.config.update("jax_compilation_cache_dir", sys.argv[1])
    jax.config.update("jax_persistent_cache_min_compile_time_secs", 0.0)   # 演示用：小函数也写盘
    jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
    logging.basicConfig(level=logging.WARNING, format="%(message)s")
    logging.getLogger("jax._src.compiler").setLevel(logging.DEBUG)       # 打印 HIT / MISS

    def model(x):
        for _ in range(20):
            x = jnp.tanh(x @ x) + 0.1
        return x

    t0 = time.perf_counter()
    exe = jax.jit(model).lower(jnp.ones((64, 64))).compile()
    print("COMPILE_SECONDS", time.perf_counter() - t0)
"""
out1 = run_py(cache_src, args=(cache_dir,))
out2 = run_py(cache_src, args=(cache_dir,))
pick = lambda s: [l for l in s.splitlines() if "jit_model" in l or "COMPILE_SECONDS" in l]
print("---- 进程 1 ----", *pick(out1), sep="\n")
print("---- 进程 2 ----", *pick(out2), sep="\n")
print("\n缓存目录内容:", sorted(f[:40] + "…" for f in os.listdir(cache_dir) if f.endswith("-cache")))

assert "PERSISTENT COMPILATION CACHE MISS for 'jit_model'" in out1
assert "Persistent compilation cache hit for 'jit_model'" in out2
assert any(f.startswith("jit_model-") for f in os.listdir(cache_dir))

# ❌ 旧写法（代码错误）：x、y 用同一个 PRNGKey 生成 → 两个数组完全相同
key_old = jax.random.PRNGKey(0)
x_old = jax.random.uniform(key_old, (1000, 1000))
y_old = jax.random.uniform(key_old, (1000, 1000))
print("x 与 y 完全相同:", bool(jnp.array_equal(x_old, y_old)))
assert bool(jnp.array_equal(x_old, y_old))

# ✅ 推荐写法：jax.random.key + split
import jax
import jax.numpy as jnp

def simple_computation(x, y):
    a = jnp.log(x)
    b = a + y
    c = jnp.exp(b)
    return c

# JIT 编译函数
compiled_computation = jax.jit(simple_computation)

# 示例数据
kx, ky = jax.random.split(jax.random.key(0))
x = jax.random.uniform(kx, (1000, 1000))
y = jax.random.uniform(ky, (1000, 1000))
assert not bool(jnp.array_equal(x, y))

# 执行
result = compiled_computation(x, y).block_until_ready()
np.testing.assert_allclose(result, np.exp(np.log(np.asarray(x, np.float64)) + np.asarray(y, np.float64)), rtol=1e-5)
print("结果正确，形状", result.shape)