内存优化

本文引用的文件 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/pjrt/compiled_memory_stats.cc - xla/backends/gpu/runtime/host_memory_pool.cc - xla/hlo/analysis/hlo_alias_analysis.cc - xla/hlo/analysis/alias_info.cc - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc - xla/service/heap_simulator/heap_simulator.h

目录

1. 简介
2. 项目结构
3. 核心组件
4. 架构总览
5. 详细组件分析
6. 依赖关系分析
7. 性能考量
8. 故障排查指南
9. 结论
10. 附录

简介

本文件系统性阐述XLA在HLO（High-Level Optimizer）层面的内存优化机制与实现，重点覆盖以下方面： - 内存空间传播算法：确保形状布局中的内存空间一致性，避免跨内存域的数据不一致。 - 缓冲区分配策略：通过点集分析与别名分析，识别可复用与可别名的缓冲区，降低峰值内存占用。 - 复制消除技术：基于别名信息与输入输出别名配置，消除不必要的拷贝与临时缓冲区。 - 生命周期调度：以贪心列表调度为核心，结合点集分析与堆模拟器估算峰值内存，指导指令执行顺序。 - 内存布局与池化：CPU侧连续段内存管理器与GPU主机内存池，提升内存分配效率与对齐。

本文件同时提供HLO优化前后的示意流程与图示，帮助读者理解内存优化如何减少设备内存占用并提升访问效率，以及对后续代码生成的影响。

项目结构

围绕内存优化的关键代码主要分布在如下模块： - HLO变换层：内存空间传播、内存放置注解转换、内存调度与别名优化。 - 分析层：别名分析、点集分析、别名信息维护。 - 后端运行时：CPU连续段内存管理器、GPU主机内存池。 - PJRT：编译后内存统计，区分设备/主机、参数/输出/临时与别名缓冲区。

graph TB
subgraph "HLO变换层"
MSP["内存空间传播<br/>memory_space_propagation.cc"]
MPA["内存放置注解转换<br/>convert_memory_placement_to_internal_annotations.cc"]
HMS["内存调度器<br/>hlo_memory_scheduler.cc"]
OIOA["输入输出缓冲区别名优化<br/>optimize_input_output_buffer_alias.cc"]
end
subgraph "分析层"
HAA["Hlo别名分析<br/>hlo_alias_analysis.cc"]
AI["别名信息<br/>alias_info.cc"]
end
subgraph "后端运行时"
CSM["CPU连续段内存管理器<br/>contiguous_section_memory_manager.cc"]
HMP["GPU主机内存池<br/>host_memory_pool.cc"]
end
subgraph "PJRT"
CMS["编译内存统计<br/>compiled_memory_stats.cc"]
end
MSP --> HAA
HMS --> HAA
HMS --> AI
MPA --> HAA
OIOA --> HAA
CSM --> |"内存对齐/分配"| CPU["CPU执行路径"]
HMP --> |"主机内存复用"| GPU["GPU执行路径"]
CMS --> |"统计设备/主机/别名"| 统计["内存报告"]

图表来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc - xla/hlo/analysis/hlo_alias_analysis.cc - xla/hlo/analysis/alias_info.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc - xla/pjrt/compiled_memory_stats.cc

章节来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc - xla/pjrt/compiled_memory_stats.cc

核心组件

• 内存空间传播（Memory Space Propagation）
• 在模块内构建数据流分析，遍历参数、根节点、融合输入/输出等叶形状，将源形状的内存空间与拆分配置传播到目标形状布局，保证一致性。
• 支持动态更新切片、融合嵌套场景下的跨层级传播。
• 内存调度器（Memory Scheduler）
• 基于点集分析与别名分析，采用贪心列表调度优先释放大缓冲区并定义小输出的指令序列，最小化峰值内存。
• 提供DFS/BFS/后序等多种备选调度策略，并通过堆模拟器评估不同方案的最小内存需求。
• 输入输出缓冲区别名优化（Optimize Input/Output Buffer Alias）
• 利用输入输出别名配置，将参数直接作为输出返回或复用中间缓冲区，消除复制。
• 内存放置注解转换（Convert Memory Placement to Internal Annotations）
• 将前端自定义调用上的外部内存放置注解转换为内部注解，驱动后续传输或放置逻辑。
• 运行时内存管理
• CPU连续段内存管理器：按代码/只读/读写段连续分配，支持最终保护与缓存失效。
• GPU主机内存池：固定大小环形空闲链表，复用主机侧临时缓冲区，降低频繁分配开销。
• 编译后内存统计（Compiled Memory Stats）
• 从分配集合中重新计算设备/主机、参数/输出/临时与别名的内存统计，区分可别名区域。

章节来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc - xla/pjrt/compiled_memory_stats.cc

架构总览

下图展示了从HLO模块到执行路径的内存优化关键通路：内存空间传播与放置注解转换确保布局一致性；内存调度器与别名分析共同决定指令顺序与缓冲区复用；运行时内存管理器负责对齐与分配；编译后统计汇总设备/主机内存使用情况。

sequenceDiagram
participant Mod as "HloModule"
participant MSP as "内存空间传播"
participant MPA as "内存放置注解转换"
participant HMS as "内存调度器"
participant HAA as "Hlo别名分析"
participant AI as "别名信息"
participant HS as "堆模拟器"
participant RT as "运行时内存管理器"
Mod->>MSP : 遍历参数/根/融合输入/输出叶形状
MSP-->>Mod : 设置形状布局内存空间与拆分配置
Mod->>MPA : 转换外部内存放置注解为内部注解
Mod->>HMS : 构建点集分析与别名分析
HMS->>HAA : 使用别名信息
HMS->>AI : 使用别名信息
HMS->>HS : 评估不同调度的最小内存
HS-->>HMS : 返回最优方案
HMS-->>Mod : 写入HloSchedule
Mod->>RT : 按需分配/保护内存段

图表来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/service/heap_simulator/heap_simulator.h

详细组件分析

组件A：内存空间传播（Memory Space Propagation）

• 目标：在HLO形状布局中传播内存空间与拆分配置，确保参数、根节点、融合输入/输出之间的布局一致。
• 关键流程
• 构建数据流分析，遍历每个计算的参数与根形状的叶形状索引。
• 对融合计算，遍历融合参数与融合表达式根的叶形状。
• 递归传播：遇到动态更新切片、融合根/参数、嵌套融合时，向上/向下传播至调用者/被调用者。
• 更新目标形状布局的内存空间与拆分配置，若发生变化则标记修改。
• 复杂度与优化
• 时间复杂度近似与值集与使用边数量线性相关；通过访问集合避免重复传播。
• 与布局工具协作，保持拆分配置一致性，减少后续布局错误。

flowchart TD
Start(["开始"]) --> BuildDF["构建数据流分析"]
BuildDF --> ForEachParam["遍历参数叶形状"]
BuildDF --> ForEachRoot["遍历根叶形状"]
ForEachParam --> Propagate["传播内存空间与拆分配置"]
ForEachRoot --> Propagate
Propagate --> DynUpdate{"是否动态更新切片?"}
DynUpdate --> |是| PropagateChild["传播到操作数0"]
DynUpdate --> |否| FusionCheck{"是否融合节点?"}
FusionCheck --> |是| PropagateFuse["传播到融合根/参数"]
FusionCheck --> |否| NextUse["处理下一个使用"]
PropagateChild --> NextUse
PropagateFuse --> NextUse
NextUse --> End(["结束"])

图表来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/memory_space_propagation.cc

章节来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/memory_space_propagation.cc

组件B：内存调度器（Memory Scheduler）

• 目标：通过贪心列表调度最小化峰值内存，结合点集分析与堆模拟器评估不同序列。
• 关键流程
• 构建点集分析映射每个指令使用的逻辑缓冲区集合。
• 计算每个逻辑缓冲区尚未调度的使用次数，统计“调度后可释放字节数”与“用户数量”作为优先级。
• 处理特殊指令（如入/出站）以改善通信与I/O时间窗口。
• 通过堆模拟器评估模块级最小内存需求，选择最优调度方案。
• 数据结构与复杂度
• 使用多映射按优先级组织就绪队列，O(logN)插入/删除。
• 通过一次拓扑扫描与增量更新，整体接近线性复杂度。

flowchart TD
S(["开始"]) --> PTS["构建点集分析"]
PTS --> Init["初始化缓冲区未调度使用计数"]
Init --> ReadyQ["构造就绪队列(按优先级)"]
ReadyQ --> Loop{"就绪队列非空?"}
Loop --> |是| PopBest["弹出最高优先级指令"]
PopBest --> Schedule["加入执行序列"]
Schedule --> DecUse["对使用的缓冲区计数减一"]
DecUse --> UpdateReady["更新相关指令入度/就绪状态"]
UpdateReady --> Adjust["必要时调整就绪队列优先级"]
Adjust --> Loop
Loop --> |否| HeapSim["堆模拟器评估最小内存"]
HeapSim --> Choose["选择最优调度方案"]
Choose --> E(["结束"])

图表来源 - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc

章节来源 - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc

组件C：输入输出缓冲区别名优化（Optimize Input/Output Buffer Alias）

• 目标：通过输入输出别名配置，使某些运算直接复用参数或中间缓冲区，避免复制。
• 实现要点
• 读取模块级输入输出别名配置，识别可别名的参数与输出。
• 在满足数据依赖与生命周期约束的前提下，替换使用关系，减少临时缓冲区与拷贝。
• 影响
• 显著降低峰值内存与带宽占用，尤其在大模型与长序列任务中收益明显。

章节来源 - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc

组件D：内存放置注解转换（Convert Memory Placement to Internal Annotations）

• 目标：将前端自定义调用上的外部内存放置注解转换为内部注解，便于后续传输或放置逻辑处理。
• 关键流程
• 遍历非融合计算中的指令，识别具有外部内存放置注解的自定义调用。
• 根据注解类型映射到对应的内部自定义调用目标，替换使用并移除原指令。
• 效果
• 统一注解语义，简化后续传输与放置决策。

章节来源 - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc

组件E：运行时内存管理

• CPU连续段内存管理器
• 按页对齐预留代码/只读/读写段空间，统一分配后分别设置保护属性，最后进行指令缓存失效。
• 适合JIT代码生成与运行时内存布局。
• GPU主机内存池
• 固定大小的主机内存块，内部维护元素指针队列，支持句柄式借用与自动归还，降低频繁分配成本。

classDiagram
class ContiguousSectionMemoryManager {
+reserveAllocationSpace(...)
+allocateCodeSection(...)
+allocateDataSection(...)
+finalizeMemory(...)
}
class HostMemoryPool {
+Create(executor, type)
+Acquire()
+Release(ptr)
}
class Handle {
+Handle(Handle&&)
+operator=(Handle&&)
+~Handle()
}
HostMemoryPool --> Handle : "返回句柄"

图表来源 - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc

章节来源 - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc

组件F：编译后内存统计（Compiled Memory Stats）

• 目标：从分配集合中重新计算设备/主机、参数/输出/临时与别名的内存统计，区分可别名区域。
• 关键逻辑
• 遍历所有分配，校验默认着色行为（颜色=内存空间），累加各类尺寸。
• 区分参数（含别名）、输出、预分配临时缓冲区的设备/主机用量。
• 应用
• 为性能分析与资源规划提供准确的内存视图，辅助定位内存热点。

章节来源 - xla/pjrt/compiled_memory_stats.cc

依赖关系分析

• 内存空间传播依赖数据流分析与布局工具，确保形状布局一致性。
• 内存调度器依赖点集分析与别名分析，结合堆模拟器评估不同调度方案。
• 输入输出缓冲区别名优化依赖模块级别名配置与别名信息。
• 运行时内存管理器与编译后统计分别服务于执行期分配与编译期统计。

graph LR
MSP["内存空间传播"] --> HAA["Hlo别名分析"]
MSP --> LUT["布局工具"]
HMS["内存调度器"] --> HAA
HMS --> AI["别名信息"]
HMS --> HS["堆模拟器"]
OIOA["输入输出别名优化"] --> HAA
MPA["内存放置注解转换"] --> HAA
CSM["CPU连续段内存管理器"] --> RT["执行路径"]
HMP["GPU主机内存池"] --> RT
CMS["编译内存统计"] --> RT

图表来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc - xla/pjrt/compiled_memory_stats.cc

章节来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc - xla/hlo/transforms/convert_memory_placement_to_internal_annotations.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc - xla/pjrt/compiled_memory_stats.cc

性能考量

• 内存空间传播
• 通过布局一致性避免跨内存域访问带来的额外拷贝与同步。
• 在融合嵌套场景中正确传播，减少布局不一致导致的回退与重排。
• 内存调度
• 贪心列表调度优先释放大缓冲区，显著降低峰值内存；对入/出站指令的特殊处理改善I/O时间窗。
• 多策略对比（DFS/后序/列表）与堆模拟器评估，选择更优序列。
• 别名与复制消除
• 输入输出别名与参数别名可直接复用缓冲区，减少临时分配与带宽压力。
• 运行时内存管理
• CPU连续段内存管理器按段分配与保护，减少碎片与保护开销。
• GPU主机内存池通过固定块与句柄借用，降低频繁分配与锁竞争。

[本节为通用性能讨论，无需列出具体文件来源]

故障排查指南

• 内存空间传播后布局不一致
• 检查融合嵌套层级的传播是否完整，确认动态更新切片与融合根/参数的传播路径。
• 参考：内存空间传播实现与递归传播逻辑。
• 峰值内存过高
• 检查是否存在大量标量指令未聚合成寄存器使用；确认列表调度是否启用且策略选择合理。
• 参考：内存调度器的优先级与堆模拟器评估。
• 参数/输出未别名
• 检查输入输出别名配置是否正确，确认别名信息与别名分析结果。
• 参考：输入输出缓冲区别名优化与别名分析。
• 运行时分配失败
• CPU：检查连续段预留空间与对齐要求；确认最终保护与缓存失效步骤。
• GPU：检查主机内存池容量与并发借用上限。
• 参考：CPU连续段内存管理器与GPU主机内存池。

章节来源 - xla/hlo/transforms/memory_space_propagation.cc - xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc - xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias.cc - xla/backends/cpu/codegen/contiguous_section_memory_manager.cc - xla/backends/gpu/runtime/host_memory_pool.cc

结论

XLA在HLO层面通过内存空间传播、内存调度、别名分析与输入输出缓冲区别名优化，形成一套完整的内存优化体系。配合运行时内存管理器与编译后统计，能够在不牺牲正确性的前提下显著降低设备内存占用、减少带宽压力并提升执行效率。对于后续代码生成，这些优化直接影响缓冲区分配、布局与指令顺序，从而提升整体吞吐与稳定性。

[本节为总结性内容，无需列出具体文件来源]

附录

• HLO示例：内存优化前后对比（概念性说明）
• 优化前：存在多个临时缓冲区与复制，峰值内存高，访问模式分散。
• 优化后：通过内存空间传播统一布局，通过内存调度最小化峰值，通过别名消除复制，访问模式更局部化。
• 关键算法实现路径
• 内存空间传播：见“内存空间传播”章节。
• 内存调度：见“内存调度器”章节。
• 别名分析与信息：见“分析层”文件。
• 运行时内存管理：见“运行时内存管理”章节。
• 编译后统计：见“编译后内存统计”章节。

[本节为概念性说明，无需列出具体文件来源]