AI编译内存规划：中间张量复用需证明生命周期

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顺序来进行内存复用。中间张量复用具有重要价值，但每一次共享缓冲区都必须有生命周期证明。在节省内存的同时，必须确保正确性。