AI编译器内存规划器（Memory Planner）：中间张量复用必须证明生命周期不重叠
在AI编译器的优化流程中，内存规划器（Memory Planner）虽常被忽视，却扮演着关键角色。它负责为每个中间张量（Tensor）分配内存位置，直接决定模型能否在资源受限设备上顺利部署。本文将从生命周期分析、分配策略、动态Shape处理等多方面，深入探讨如何正确高效地实现中间张量复用，从而避免因内存规划不当导致的部署失败或性能瓶颈。
一、中间张量：隐藏的内存占用大户
在AI编译器优化中，算子融合、常量折叠、图裁剪等优化技术备受关注，但内存规划器同样至关重要。推理过程中，中间张量的生命周期短暂，若每个张量都独立占用缓冲区，内存峰值将急剧上升。特别是在边缘设备或多并发推理场景下，中间张量的复用直接决定了模型的可部署性。
然而，内存复用不能仅凭直觉判断。每一次复用都必须严格证明生命周期不重叠。
二、生命周期分析是基础
flowchart TD
A[计算图] --> B[拓扑排序]
B --> C[Tensor 首次使用]
C --> D[Tensor 最后使用]
D --> E[内存区间分配]每个张量都有其诞生和消亡时刻：即由哪个算子生成，以及被哪个算子最后一次使用。只有生命周期不重叠的张量，才能共享同一块内存区域。这是内存复用的核心前提。
tensor_lifetime:
tensor_17:
birth_op: conv_3
last_use_op: relu_3
size_bytes: 262144生命周期信息最好能够导出（dump）以供分析。当调试内存峰值时，工程师需要清楚哪些张量同时处于活跃状态。这有助于定位内存瓶颈，并为后续优化提供依据。
三、分配策略必须考虑对齐约束
fn align_up(v: usize, align: usize) -> usize {
(v + align - 1) & !(align - 1)
}NPU、SIMD和DMA等硬件通常要求内存对齐。内存规划器如果仅按大小进行区间染色，可能会生成硬件无法接受的地址。对齐方式、步长、内存bank、读写权限等都可能成为约束条件。
此外，还需要区分常量权重和临时激活张量。权重通常为只读且生命周期贯穿整个模型，而临时激活张量才是需要复用的主要对象。将两者混入同一个内存池，会导致规划逻辑复杂且难以审查。
对齐约束在不同硬件上差异显著。例如，NVIDIA GPU的kernel通常要求float4访问需16字节对齐，而在Tensor Core路径下甚至要求128字节对齐。某些NPU的片上SRAM按bank分组，同一bank内的连续地址不能同时访问，内存规划器需对地址进行哈希或填充以避免bank冲突。在Rust实现的编译器中，对齐规划可采用“先排序再贪心”策略：按大小降序排列待分配的张量，从偏移0开始，每次通过维护的空闲区间列表寻找最小可用偏移量，分配后更新区间，确保大张量优先占据对齐友好的位置。对于encoder-decoder等多子图模型，还需考虑子图切换时的生命周期断层：encoder的输出作为decoder的输入需要保留，但encoder内部的中间张量在decoder阶段完全可以被覆盖。这意味着内存规划器需要跨子图进行全局生命周期分析，而不是逐子图独立规划。
四、动态Shape使内存规划更复杂
在静态Shape情况下，张量大小可在编译期确定；而在动态Shape情况下，输入长度、batch大小或图分支的变化会导致内存需求动态变化。此时，内存规划器需要保守估计，或在运行时根据Shape重新规划。
memory_planner_policy:
static_shape: compile_time_plan
dynamic_shape: runtime_plan_with_cap
dump_peak_memory: true动态规划必须设定上限。否则，一次超长输入可能导致运行时临时申请大量内存，破坏服务稳定性。推理服务通常应限制Shape范围，然后在此范围内进行内存规划。
内存复用还需配合算子融合。融合后，某些中间张量不再物化，其生命周期会发生变化；而融合失败时，内存峰值可能再次升高。各优化Pass之间应重新计算生命周期，切勿复用已过期的分析结果。
五、验证与报告：保障内存规划正确性
最后，内存规划器的正确性需借助工具验证。可以在调试模式下向共享内存区填充特定模式，检查张量是否被提前覆盖。内存规划出错时，输出可能仅出现轻微漂移，很难从日志中直接发现。
此外，还需输出峰值解释报告。仅给出一个“峰值512MB”的数字远远不够，工程师需要了解峰值发生在哪个算子区间、哪些张量同时存活、哪些缓冲区因对齐而浪费了空间。在优化内存时，解释报告远比肉眼查看计算图有效。
memory_plan_report:
show_peak_interval: true
show_live_tensors_at_peak: true
show_alignment_waste: true
export_dot_graph: true如果计划将内存规划器用于多个硬件后端，则应将后端约束作为输入参数，而非在编译器中硬编码。CPU、GPU、NPU的对齐要求和可访问内存域各不相同，同一计算图可能需要不同的内存规划方案。
常见问题解答（FAQ）
- 问：如何快速定位导致内存峰值的具体张量？
答：启用memory_plan_report中的show_peak_interval和show_live_tensors_at_peak配置，即可输出峰值区间的详细信息，结合导出的生命周期文件，精准定位问题。 - 问：在动态Shape场景下，如果输入超出预期范围应如何处理？
答：在推理服务层设置Shape上限，并在内存规划器中配置最大容量。一旦超出，应拒绝请求或触发降级策略，避免内存爆炸。 - 问：不同硬件后端的对齐要求差异较大，如何统一处理？
答：建议将后端约束抽象为配置输入，而非在编译器代码中硬编码。内存规划器根据后端配置动态调整对齐策略，并可通过导出报告验证对齐是否生效。
AI编译器内存规划器需要基于张量生命周期、对齐约束、动态Shape以及优化Pass顺序来进行内存复用。中间张量复用具有重要价值，但每一次共享缓冲区都必须有生命周期证明。在节省内存的同时，必须确保正确性。
