第4章 JIT编译深度优化

第4章：JIT编译深度优化：静态参数与动态参数的处理、三值规则、编译缓存与手动控制编译流程

JIT编译，说白了就是让JAX在背后偷偷帮你把Python代码变成高效的机器码。但很多同学用起来总觉得「玄学」——有时候快得飞起，有时候又莫名其妙重新编译。今天我就把这层窗户纸捅破，聊聊JIT编译的那些坑和优化技巧。

4.1 静态参数与动态参数：编译器的「已知」与「未知」

先问个问题：为什么同样的函数，换了个参数值就要重新编译？

答案很简单——JAX需要知道哪些参数是「静态的」，哪些是「动态的」。静态参数的值在编译时就固定下来，编译器可以针对这个值做极致优化。动态参数则保留为运行时变量。

举个例子：

import jax
import jax.numpy as jnp

@jax.jit
def f(x, n):
    return x + jnp.arange(n)

# 第一次调用，n=5
f(jnp.array([1.0]), 5)  # 编译一次

# 第二次调用，n=10
f(jnp.array([1.0]), 10) # 又编译一次！

为什么会这样？因为n默认是动态参数，但jnp.arange(n)的输出形状依赖于n的值。JAX没法在编译时确定输出形状，只能每次重新编译。

核心原则：影响数组形状的参数，必须标记为静态。

修复方法很简单：

@jax.jit(static_argnums=(1,))  # 标记第1个参数为静态
def f(x, n):
    return x + jnp.arange(n)

# 现在n=5和n=10都会触发重新编译
# 但n=5时，编译器知道arange(5)的形状是(5,)

我个人习惯把static_argnums和static_argnames配合使用。参数多了用名字标记更清晰：

@jax.jit(static_argnames=['n'])
def f(x, n):
    return x + jnp.arange(n)

小技巧：静态参数不要滥用。每个静态参数都会产生一份独立的编译缓存。如果你有10个不同的n值，就会编译10次。我见过有人把batch_size也设为静态，结果训练时每个batch size都重新编译一次，反而更慢。

4.2 三值规则（Three-value rule）：JAX的「编译三态」

嗯，这里要讲一个比较抽象但很重要的概念——三值规则。JAX在JIT编译时，每个参数其实有三种状态：

状态	含义	示例
静态已知值	编译时值完全确定	`static_argnums`标记的参数
动态未知值	运行时才确定，但形状已知	普通数组参数
动态未知形状	运行时形状都可能变化	需要`jnp.where`等动态操作

你想想看，编译器最怕什么？最怕「未知的形状」。因为所有内存分配、循环展开、向量化优化都依赖于形状信息。

我在项目中遇到过这样一个坑：

@jax.jit
def process_batch(batch):
    # batch中可能包含不同长度的序列
    lengths = jnp.sum(batch != 0, axis=1)
    max_len = jnp.max(lengths)
    return batch[:, :max_len]  # 动态切片！

这段代码每次调用都会重新编译，因为max_len的值在变化，导致切片后的形状不确定。解决方案是用jnp.pad统一形状，或者用lax.while_loop处理动态逻辑。

避坑指南：我曾经花了一整天调试一个模型，发现每次推理都慢得离谱。最后发现是jnp.nonzero返回了动态形状的数组，导致JIT无法缓存编译结果。记住：任何产生动态形状的操作，都会破坏JIT的优化效果。

4.3 编译缓存机制：JAX的「记忆」能力

JAX的编译缓存，说白了就是一个哈希表。每次调用JIT函数时，JAX会计算输入参数的「签名」（包括形状、数据类型、静态参数的值），然后查缓存。命中就直接执行编译后的代码，没命中就重新编译。

缓存的关键要素：

参数形状：(batch, 784)和(batch, 1024)是不同的缓存条目
数据类型：float32和float16是不同的
静态参数值：每个不同的静态值对应一个缓存
设备类型：CPU和GPU的编译结果不同

如何查看缓存状态？

import jax

@jax.jit
def f(x):
    return x * 2

# 查看缓存信息
print(f.lower(jnp.array([1.0])).compiler_ir())
# 或者用jax.jit的cache_info()
# 注意：jax没有直接暴露cache_info，但可以通过traceback判断

实际项目中，我建议用jax.jit的cache参数控制缓存行为：

from functools import partial

@partial(jax.jit, cache=True)  # 显式开启缓存
def f(x):
    return x * 2

经验之谈：如果你的模型有多个不同的输入形状（比如推理服务中batch size变化），可以考虑用jax.jit的static_argnums把batch size设为静态。这样每个batch size编译一次，后续直接命中缓存。但要注意控制静态参数的数量，避免缓存爆炸。

4.4 使用lower与compile手动控制编译流程

JAX的JIT编译其实分两步：lower和compile。默认情况下，@jax.jit把这两步合在一起了。但如果你想精细控制，可以手动拆开。

为什么要手动控制？

调试时想看看中间IR（中间表示）长什么样
想复用编译结果，避免重复编译
想修改编译选项（比如优化级别）

来看个完整例子：

import jax
import jax.numpy as jnp

def f(x, y):
    return jnp.dot(x, y) + 1.0

# 第一步：lower，生成HLO IR
x = jnp.ones((4, 4))
y = jnp.ones((4, 4))

lowered = jax.jit(f).lower(x, y)
print(lowered.compiler_ir())  # 打印HLO IR

# 第二步：compile，生成可执行代码
compiled = lowered.compile()

# 第三步：执行
result = compiled(x, y)

手动控制的好处是什么？

# 可以复用编译结果
compiled_func = jax.jit(f).lower(x, y).compile()

# 多次调用，只编译一次
for _ in range(100):
    result = compiled_func(x, y)  # 直接执行，不重新编译

我在项目中遇到过这样一个场景：模型推理服务需要处理不同batch size的请求。如果用@jax.jit，每次新batch size都要重新编译。我的做法是：

# 预编译几个常用batch size
batch_sizes = [1, 4, 8, 16, 32]
compiled_funcs = {}

for bs in batch_sizes:
    x_sample = jnp.ones((bs, 784))
    compiled_funcs[bs] = jax.jit(model).lower(x_sample).compile()

# 推理时直接选择对应的编译结果
def predict(batch):
    bs = batch.shape[0]
    # 找到最接近的预编译batch size
    target_bs = min(batch_sizes, key=lambda x: abs(x - bs))
    # 如果batch size不匹配，需要padding或截断
    if bs < target_bs:
        batch = jnp.pad(batch, ((0, target_bs - bs), (0, 0)))
    elif bs > target_bs:
        batch = batch[:target_bs]
    return compiled_funcs[target_bs](batch)

注意：手动编译时，lower和compile之间可以插入一些优化步骤。比如用jax.jit的donate_argnums参数标记哪些参数可以被覆盖，减少内存分配。但这是进阶用法，新手建议先用默认的@jax.jit。

4.5 知识体系总览

下面这张图总结了本章的核心逻辑：

这张图把四个知识点串起来了。从左到右，其实是一个从「基础控制」到「高级优化」的递进关系。静态参数是基础，三值规则帮你理解编译器的「思维」，缓存机制是性能保障，手动控制则是终极武器。

最后说一句：JIT编译优化没有银弹。每个模型、每个场景都需要你亲手去调、去试。但理解了这些底层机制，你至少知道问题出在哪，而不是对着「重新编译」的日志干瞪眼。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录