在现代AI系统中，尤其是自动驾驶和具身智能场景下，输入数据的维度经常是动态变化的。车辆检测到的目标数量、语音指令的长度、点云数据的密度都是运行时才能确定的。本章将深入探讨AI编译器如何优雅地处理动态shape，以及如何通过JIT（Just-In-Time）编译技术在运行时性能和灵活性之间取得平衡。
在传统的深度学习框架中，张量的形状（shape）通常在图构建时就已确定。然而，现实世界的AI应用经常面临形状不确定的情况：
静态Shape示例：
输入: [batch_size=32, height=224, width=224, channels=3]
卷积输出: [32, 112, 112, 64] # 所有维度编译时已知
动态Shape示例：
输入: [batch_size=?, seq_len=?, embedding_dim=768]
注意力输出: [?, ?, 768] # batch_size和seq_len运行时确定
在自动驾驶场景中，动态shape无处不在：
动态shape给编译器优化带来了根本性挑战，这些挑战贯穿整个编译流程：
1. 内存分配不确定性
静态shape允许编译器预先计算所有中间张量的大小，实现精确的内存池管理。而动态shape下：
在自动驾驶场景中，这种不确定性尤为突出。考虑一个多目标跟踪系统：
帧1：检测到5辆车 → 需要 5×feature_dim 内存
帧2：检测到20辆车 → 需要 20×feature_dim 内存
帧3：高速路段，50辆车 → 需要 50×feature_dim 内存
内存管理策略：
- 悲观策略：按最大可能分配 → 浪费严重
- 乐观策略：按平均值分配 → 可能OOM
- 自适应策略：动态扩容 → 碎片化和延迟
编译器必须在这些策略间权衡，通常采用分级内存池：
2. 优化决策困难
许多编译器优化依赖于具体的tensor维度：
静态循环（N=4）：
for i in range(4):
process(data[i])
可展开为：
process(data[0])
process(data[1])
process(data[2])
process(data[3])
动态循环（N=?）：
for i in range(N):
process(data[i])
无法完全展开，只能部分展开：
for i in range(0, N, 4):
if i+3 < N:
process(data[i])
process(data[i+1])
process(data[i+2])
process(data[i+3])
else:
# 处理剩余元素
SIMD宽度选择问题：
AVX2: 256位，可处理8个float32
AVX512: 512位，可处理16个float32
动态shape下的策略：
if N % 16 == 0:
use_avx512() # 完美对齐
elif N % 8 == 0:
use_avx2() # 次优对齐
else:
use_scalar_with_remainder() # 标量处理尾部
Conv + BatchNorm + ReLU融合：
静态shape：可以精确计算中间buffer大小
动态shape：需要最坏情况预留或动态分配
3. 调度复杂性
考虑矩阵乘法的并行化：
C[M, N] = A[M, K] @ B[K, N]
静态shape（M=1024, K=512, N=2048）：
- 块大小：64×64（经过调优的最优值）
- 线程分配：16×32线程网格
- 负载均衡：每个线程处理相同工作量
动态shape：
- M, K, N运行时才知道
- 需要动态选择分块策略
- 可能出现负载不均衡
动态调度算法：
def dynamic_gemm_schedule(M, K, N):
# 根据问题规模选择策略
if M * N < 1000: # 小矩阵
return single_thread_gemm
elif M * N < 100000: # 中等矩阵
block_size = min(32, M, N)
return blocked_gemm(block_size)
else: # 大矩阵
# 自适应选择块大小
block_m = find_best_factor(M, target=64)
block_n = find_best_factor(N, target=64)
block_k = find_best_factor(K, target=32)
return tiled_gemm(block_m, block_n, block_k)
4. 寄存器分配挑战
动态shape影响寄存器的使用模式：
静态shape优化：
- 编译时确定每个变量的生命周期
- 精确的寄存器分配
- 最小化spill/reload
动态shape问题：
- 循环迭代次数未知，生命周期不确定
- 可能需要保守的寄存器分配
- 更多的寄存器溢出
5. 指令调度困难
现代处理器的指令级并行依赖于静态分析：
静态shape：
load A[0] | load B[0] | mul | add | store C[0]
load A[1] | load B[1] | mul | add | store C[1]
load A[2] | load B[2] | mul | add | store C[2]
可以流水线化，隐藏延迟
动态shape：
需要插入边界检查，破坏流水线：
if (i < N) load A[i]
if (i < N) load B[i]
if (i < N) mul
...
条件分支影响指令预取和流水线效率
尽管充满挑战，动态shape也带来了新的优化机遇：
1. 自适应优化
编译器可以根据实际运行时shape选择最优策略：
2. 特化与版本化
通过收集运行时信息，编译器可以：
3. 延迟优化决策
某些优化可以推迟到运行时：
即使存在动态维度，编译器仍可进行部分shape推断，这是优化动态shape程序的关键技术：
1. 符号化Shape表示
编译器使用符号变量表示未知维度，构建符号表达式：
符号化Shape分析：
输入: [N, 768] # N是符号变量
Linear层: W[768, 512]
输出: [N, 512] # 可以推断出输出shape为[N, 512]
复杂案例 - 多头注意力：
输入: [B, L, D] # B=batch, L=seq_len, D=hidden_dim
拆分: [B, L, H, D/H] # H=num_heads
转置: [B, H, L, D/H]
注意力: [B, H, L, L] @ [B, H, L, D/H] = [B, H, L, D/H]
合并: [B, L, D]
符号约束系统：
- D % H == 0 # hidden_dim必须被num_heads整除
- L <= MAX_SEQ_LEN # 序列长度上界
- B >= 1 # 批大小至少为1
2. Shape约束传播
编译器通过数据流分析传播shape约束：
约束传播示例：
如果 A[M, K] @ B[K, N] = C[M, N]
且已知 M <= 1024（批大小上限）
则可以传播约束：
- C的第一维 <= 1024
- 内存需求 <= 1024 * N * sizeof(dtype)
传播规则：
1. 相等约束：Reshape操作要求元素总数相等
[B, L, D] → [B*L, D] 要求输入输出元素数相同
2. 广播约束：element-wise操作的广播规则
[B, 1, D] + [1, L, D] → [B, L, D]
3. 缩减约束：reduction操作消除特定维度
sum([B, L, D], axis=1) → [B, D]
4. 切片约束：索引操作的范围限制
A[0:K, :] 要求 K <= A.shape[0]
3. 区间分析与范围推断
对动态维度进行区间分析，缩小可能值范围：
区间算术：
N ∈ [1, 100] # 批大小范围
L ∈ [10, 512] # 序列长度范围
MatMul([N, 256], [256, 512]):
输出shape: [N, 512]
内存范围: [512, 51200] * sizeof(float32)
条件分支的区间细化：
if N < 32:
# 此分支内 N ∈ [1, 31]
small_batch_kernel()
else:
# 此分支内 N ∈ [32, 100]
large_batch_kernel()
4. Shape特化点识别
通过值谱分析(value profiling)识别常见shape：
运行时统计：
shape_histogram = {
(32, 128, 768): 45%, # 最常见
(16, 256, 768): 30%, # 次常见
(1, 512, 768): 15%, # 推理模式
其他: 10%
}
特化策略：
- 为前3个高频shape生成特化代码
- 其他shape使用通用版本
- 定期更新统计，调整特化策略
5. 依赖shape的优化决策
基于shape信息做出智能优化选择：
def shape_dependent_optimization(shape_info):
if shape_info.is_static():
# 完全静态，使用激进优化
return full_optimization_pipeline()
elif shape_info.has_bounds():
# 有界动态，基于上界优化
if shape_info.max_size() < SMALL_THRESHOLD:
return small_tensor_optimization()
else:
return bounded_optimization()
else:
# 完全动态，保守优化
return conservative_optimization()
6. Shape不变量检测
识别程序中的shape不变量，用于优化：
循环不变量：
for i in range(num_iterations):
x = process(x) # x的shape在循环中不变
# 编译器可以提升shape相关计算到循环外
相对不变量：
if x.shape[0] == y.shape[0]:
# 在此作用域内，两者第一维相等
# 可以共享内存分配，优化计算
结构不变量：
class ResidualBlock:
def forward(self, x):
# 输入输出shape相同（残差连接的特性）
return x + self.transform(x)
这种符号化分析使得编译器能够：
JIT（Just-In-Time）编译是一种延迟编译策略，在程序运行时根据实际执行情况进行编译优化。与AOT（Ahead-Of-Time）编译相比，JIT能够利用运行时信息做出更好的优化决策。
编译策略对比：
AOT编译流程：
源代码 → 编译器 → 优化后的机器码 → 部署 → 执行
↑
编译时优化（保守）
JIT编译流程：
源代码 → 字节码/IR → 运行时 → 监控执行 → 热点检测 → JIT编译 → 优化代码
↑ ↓
收集profile信息
1. 解释器/基线编译器
JIT系统通常从解释执行或简单编译开始：
2. Profile收集器
运行时信息收集是JIT优化的基础：
3. 优化编译器
基于profile信息进行激进优化：
4. 代码缓存
管理编译后的代码：
JIT编译的触发时机对性能至关重要，需要在编译开销和执行效率间找到平衡点：
1. 基于计数的触发
最直观的触发策略是统计执行次数：
函数执行计数器：
def forward(x, shape):
counter[shape] += 1
if counter[shape] > THRESHOLD:
compiled_fn = jit_compile(forward, shape)
cache[shape] = compiled_fn
return execute(x, shape)
自适应阈值：
class AdaptiveThreshold:
def __init__(self):
self.base_threshold = 50
self.compile_history = []
def should_compile(self, exec_count, compile_cost_estimate):
# 根据历史编译效果调整阈值
if self.compile_history:
avg_speedup = mean([h.speedup for h in self.compile_history])
if avg_speedup > 2.0:
threshold = self.base_threshold * 0.8 # 更激进
elif avg_speedup < 1.2:
threshold = self.base_threshold * 1.5 # 更保守
else:
threshold = self.base_threshold
else:
threshold = self.base_threshold
return exec_count > threshold
2. 基于采样的触发
定期采样比每次计数更高效：
采样策略：
class SamplingProfiler:
def __init__(self, sample_rate=0.01):
self.sample_rate = sample_rate
self.samples = defaultdict(int)
def record(self, function, shape):
if random.random() < self.sample_rate:
self.samples[(function, shape)] += 1
def get_hot_functions(self, threshold=10):
# 根据采样频率推算实际执行次数
hot = []
for (func, shape), count in self.samples.items():
estimated_count = count / self.sample_rate
if estimated_count > threshold:
hot.append((func, shape, estimated_count))
return hot
优点：
3. 分层编译策略
渐进式优化，避免过早的重编译开销：
执行层级：
Level 0: 解释执行
Level 1: 快速编译（无优化）
Level 2: 基础优化编译
Level 3: 完全优化编译
升级策略：
class TieredCompilation:
def __init__(self):
self.levels = {
0: (10, self.interpret),
1: (100, self.quick_compile),
2: (1000, self.optimized_compile),
3: (10000, self.fully_optimized_compile)
}
self.current_level = {}
self.exec_count = {}
def execute(self, func, args):
key = (func, self.get_shape_signature(args))
self.exec_count[key] += 1
current = self.current_level.get(key, 0)
# 检查是否需要升级
threshold, compiler = self.levels[current]
if current < 3 and self.exec_count[key] > threshold:
# 异步编译下一级
self.async_compile(func, args, current + 1)
# 执行当前版本
return self.get_executor(key, current)(args)
实际案例 - 自动驾驶感知模块：
- Level 0：首次执行，解释模式
- Level 1：检测到目标的常规处理（10次后）
- Level 2：高频场景如车道线检测（100次后）
- Level 3：核心循环如NMS算法（1000次后）
4. 基于延迟的触发
考虑执行时间而非仅次数：
延迟敏感触发：
class LatencyBasedTrigger:
def __init__(self, latency_budget_ms=100):
self.budget = latency_budget_ms
self.exec_times = defaultdict(list)
def should_compile(self, func, shape, last_exec_time_ms):
key = (func, shape)
self.exec_times[key].append(last_exec_time_ms)
if len(self.exec_times[key]) < 5:
return False # 需要更多样本
# 如果平均执行时间超过预算的10%，考虑编译
avg_time = mean(self.exec_times[key][-10:])
if avg_time > self.budget * 0.1:
# 估算编译后的时间
estimated_speedup = 2.0 # 保守估计2倍加速
if avg_time / estimated_speedup < self.budget * 0.05:
return True
return False
5. 基于内存压力的触发
在内存受限环境下的特殊考虑：
内存感知触发：
class MemoryAwareTrigger:
def __init__(self, memory_limit_mb=1024):
self.limit = memory_limit_mb
self.compiled_size = {}
def should_compile(self, func, shape):
# 估算编译后代码大小
estimated_size = self.estimate_code_size(func, shape)
current_usage = self.get_current_memory_usage()
if current_usage + estimated_size > self.limit:
# 内存不足，可能需要驱逐其他编译代码
if self.can_evict(estimated_size):
self.evict_lru(estimated_size)
return True
return False # 无法腾出足够空间
return True # 有足够内存
6. 混合触发策略
实际系统通常结合多种策略：
综合决策系统：
class HybridTrigger:
def __init__(self):
self.count_trigger = CountBasedTrigger()
self.sampling_trigger = SamplingTrigger()
self.latency_trigger = LatencyBasedTrigger()
self.memory_trigger = MemoryAwareTrigger()
def should_compile(self, context):
# 权重投票
votes = [
(0.3, self.count_trigger.vote(context)),
(0.2, self.sampling_trigger.vote(context)),
(0.3, self.latency_trigger.vote(context)),
(0.2, self.memory_trigger.vote(context))
]
score = sum(weight * vote for weight, vote in votes)
# 根据系统状态调整阈值
threshold = 0.5
if context.is_training:
threshold = 0.7 # 训练时更保守
elif context.is_inference:
threshold = 0.3 # 推理时更激进
return score > threshold
JIT系统必须处理优化假设失效的情况：
1. 守护（Guards）机制
编译时假设：
- Shape = [32, 128, 768]
- Dtype = float32
运行时守护：
if shape != [32, 128, 768] or dtype != float32:
fallback_to_interpreter() # 去优化
else:
execute_optimized_code()
2. 去优化点（Deoptimization Points）
在优化代码中插入检查点：
3. On-Stack Replacement (OSR)
允许在循环执行过程中切换到优化版本：
Shape专门化是指为特定的张量形状生成优化的代码版本。这种技术特别适合处理虽然动态但倾向于若干固定模式的场景。
通用版本 vs 专门化版本：
通用矩阵乘法：
for i in range(M):
for j in range(N):
for k in range(K):
C[i,j] += A[i,k] * B[k,j]
专门化版本（M=32, K=64, N=128）：
// 可以完全展开内层循环
// 优化的内存访问模式
// SIMD向量化
在自动驾驶场景中，常见的shape模式：
1. 离散模式
批处理大小：1, 4, 8, 16, 32
图像分辨率：[640,480], [1280,720], [1920,1080]
2. 范围模式
序列长度：10-100（语音指令）
目标数量：0-50（检测框）
点云数量：1000-100000
3. 关联模式
如果 batch_size = 32，则 hidden_size = 256
如果 batch_size = 1，则 hidden_size = 512
1. 版本选择树
Root
|
batch_size?
/ \
<16 >=16
/ \
seq_len? seq_len?
/ \ / \
<50 >=50 <50 >=50
V1 V2 V3 V4
2. 版本淘汰策略
3. 渐进式特化
阶段1：收集所有shape
阶段2：聚类分析，识别主要模式
阶段3：为高频模式生成特化版本
阶段4：持续监控，动态调整
1. 重编译触发条件
2. 编译成本模型
收益评估：
Benefit = (预期执行次数 × 单次性能提升) - 编译开销
决策：
if Benefit > THRESHOLD:
trigger_recompilation()
else:
use_generic_version()
3. 异步编译
避免阻塞主执行流：
在自动驾驶系统中，变长序列处理无处不在：
1. 语音交互系统
驾驶员指令：
"导航到最近的充电站" (8个token)
"帮我找一个有特斯拉超充、营业到晚上10点、评分4星以上的充电站" (25个token)
系统需要处理：
- 不同长度的语音输入
- 实时语音识别的流式输入
- 多轮对话的上下文拼接
2. 轨迹预测
历史轨迹序列：
车辆A: [(x1,y1,t1), (x2,y2,t2), ..., (xn,yn,tn)] # n=5-50不等
行人B: [(x1,y1,t1), (x2,y2,t2), ..., (xm,ym,tm)] # m=3-20不等
预测挑战：
- 不同对象的历史长度不同
- 遮挡导致的序列中断
- 多模态预测的分支数量变化
3. 地图匹配与路径规划
道路网络表示：
路径1: [节点1, 节点2, ..., 节点k] # k取决于距离
路径2: [节点1, 节点3, ..., 节点j] # j可能不等于k
需要处理：
- 不同长度的候选路径
- 动态的道路网络图
- 实时的路况更新序列
处理变长序列的两种主要策略：
1. Padding策略
原始序列：
seq1: [A, B, C] 长度=3
seq2: [D, E, F, G, H] 长度=5
seq3: [I, J] 长度=2
Padding后：
seq1: [A, B, C, PAD, PAD]
seq2: [D, E, F, G, H]
seq3: [I, J, PAD, PAD, PAD]
批处理张量: [3, 5, hidden_dim]
优点：
缺点：
2. Packing策略
序列打包：
packed_data: [A, B, C, D, E, F, G, H, I, J]
batch_indices: [0, 0, 0, 1, 1, 1, 1, 1, 2, 2]
positions: [0, 1, 2, 0, 1, 2, 3, 4, 0, 1]
紧凑存储: [10, hidden_dim]
优点：
缺点：
1. 桶化策略（Bucketing）
长度桶定义：
Bucket1: 1-10 tokens
Bucket2: 11-20 tokens
Bucket3: 21-50 tokens
Bucket4: 51-100 tokens
动态分配：
incoming_sequences → 分类到桶 → 桶内批处理 → 并行执行
2. 自适应批大小
延迟约束下的批处理：
while True:
batch = []
deadline = current_time + MAX_LATENCY
while current_time < deadline and len(batch) < MAX_BATCH:
if new_request_available():
batch.append(get_request())
wait_or_timeout()
if batch:
process_batch(batch)
3. 序列长度预测
基于历史模式预测序列长度：
利用预测优化调度和内存分配。
Transformer中的注意力计算对序列长度敏感：
1. 注意力复杂度分析
标准注意力: O(n²·d)
其中 n=序列长度, d=隐藏维度
当n变化时：
n=10: 100·d 的计算量
n=100: 10000·d 的计算量
n=1000: 1000000·d 的计算量
2. 动态选择注意力算法
if seq_len < 64:
use_standard_attention() # 简单高效
elif seq_len < 512:
use_flash_attention() # 内存优化
else:
use_sparse_attention() # 稀疏模式
3. KV Cache管理
在自回归生成中：
缓存策略：
- 固定大小缓存：预分配最大长度
- 动态增长缓存：按需扩展
- 滑动窗口缓存：只保留最近k个token
内存布局优化：
连续存储 vs 分块存储
行优先 vs 列优先
自动驾驶的实时性要求给序列处理带来额外挑战：
1. 流式处理
语音识别流水线：
音频输入 → 特征提取 → 编码器 → 解码器 → 文本输出
↓ ↓ ↓ ↓
10ms 20ms 30ms 40ms
必须在下一帧到达前完成处理
2. 增量计算
轨迹更新：
时刻t: process([p1, p2, p3])
时刻t+1: process([p1, p2, p3, p4])
优化：重用t时刻的中间结果
避免重复计算[p1, p2, p3]部分
3. 优先级调度
紧急度分级：
优先级1: 碰撞检测相关序列（<10ms）
优先级2: 路径规划序列（<100ms）
优先级3: 舒适性优化序列（<1s）
优先级4: 信息娱乐系统（best effort）
投机执行允许系统在完整信息可用前开始计算，特别适合自动驾驶的预测场景：
1. 多假设预测
行人运动预测：
假设1: 继续直行 → 预计算路径1
假设2: 转向过马路 → 预计算路径2
假设3: 停止等待 → 预计算路径3
当获得更多观测后，选择最可能的分支
2. 投机解码
自回归生成的投机：
主模型: 生成token的同时
小模型: 快速预测后续k个token
验证: 主模型验证预测，接受或拒绝
1. 分支预测提示
if likely(is_highway_scenario):
# 编译器优化这个分支
process_highway_logic()
else:
process_urban_logic()
2. 投机内存分配 预分配可能需要的缓冲区，减少运行时开销。
3. 计算图的投机展开 提前展开可能的计算路径，支持快速切换。
本章深入探讨了AI编译器处理动态shape和JIT编译的核心技术：
动态Shape的本质：运行时才能确定的张量维度，在自动驾驶场景中普遍存在（目标数量、序列长度、点云密度等）
JIT编译原理：通过延迟编译和运行时信息收集，实现更精准的优化决策
Shape专门化：为常见shape模式生成优化版本，平衡通用性和性能
变长序列处理：Padding vs Packing的权衡，动态批处理策略，注意力机制的自适应优化
投机执行：在不确定场景下的预计算策略，特别适合预测任务
练习7.1：Shape推断 给定以下计算图，其中N和M是动态维度，推断每个操作后的shape：
输入: X[N, 768]
操作1: Linear(768, 512)
操作2: Reshape(-1, 64, 8)
操作3: Transpose(1, 2)
操作4: MatMul with Y[N, 8, 32]
💡 提示：考虑维度的约束关系和广播规则
练习7.2：JIT触发策略 设计一个JIT编译触发策略，考虑以下因素：
💡 提示：考虑收益break-even点
练习7.3：Padding效率分析 批处理中有100个序列，长度分布为：
比较Padding和Packing策略的内存使用和计算效率。
💡 提示：计算有效计算比例和内存占用
练习7.4：动态批处理优化 自动驾驶系统需要处理三类请求：
设计一个动态批处理策略。
💡 提示：考虑优先级队列和延迟约束
练习7.5：Shape特化的成本模型 设计一个成本模型来决定是否为特定shape生成专门化代码。考虑：
💡 提示：考虑LRU缓存和编译成本预测
练习7.6：JIT编译的正确性保证 在JIT编译系统中，如何保证动态优化不会破坏程序语义？设计一个验证框架。
💡 提示：考虑守护条件、状态一致性和回滚机制
练习7.7：增量JIT编译 设计一个增量JIT编译系统，支持在已编译代码基础上增加新的优化，而不需要完全重新编译。
💡 提示：考虑优化的依赖关系和代码patch机制
问题：为每个见过的shape都生成特化版本 后果：编译开销爆炸，缓存污染 解决：聚类相似shape，设置特化阈值
问题：频繁触发重编译 后果：编译时间超过执行节省 解决：建立成本模型，延迟编译决策
问题：JIT缓存无限增长 后果：内存耗尽，系统崩溃 解决：实现缓存淘汰策略，监控内存使用
问题：优化基于错误的shape不变量 后果：计算错误或崩溃 解决：完善的守护机制，保守的假设
问题：多线程同时触发编译 后果：重复编译，竞态条件 解决：编译锁，共享缓存，原子操作