静态内存规划与分配

执行机器学习模型，尤其是大型模型或部署在内存受限设备上的模型，需要细致的内存资源管理。尽管通用编程中，使用 malloc/free 或类似机制进行动态内存分配很常见，但ML计算图通常的静态结构为更积极的编译时优化：静态内存规划，提供了可能性。

与运行时分配（其根据内存请求的出现做出反应）不同，静态内存规划会提前分析整个计算图，以确定每个张量（中间结果）的精确生命周期。基于此分析，编译器可以设计一种内存分配策略，尽可能地重用内存缓冲区，显著降低执行期间所需的峰值内存占用。本节详细阐述了ML编译器中静态内存规划的原理和技术。

理解张量生命周期

静态内存规划的基础是理解每个张量的内存缓冲区何时实际需要。当生成张量的操作完成时，张量缓冲区变为活跃；而在最后一个消费张量的操作完成读取后，它变为非活跃。这两个时间点之间的间隔是张量的活跃范围。

考虑一个简单的序列：

1. T1 = Conv(Input)
2. T2 = ReLU(T1)
3. T3 = MaxPool(T1)
4. T4 = Add(T2, T3)

这里，T1由Conv生成，并由ReLU和MaxPool消费。其活跃范围在Conv完成后开始，并且仅在ReLU和MaxPool都完成后结束。T2从ReLU后直到Add完成后活跃。T3从MaxPool后直到Add完成后活跃。如果我们为Input、T1、T2、T3和T4分配独立的、唯一的缓冲区，所需的总内存是它们大小的总和。然而，静态分析可以显示重用的机会。

活跃度分析与干扰

编译器执行活跃度分析以正式确定这些活跃范围。此分析通常以逆向执行顺序遍历计算图，传播活跃度信息。结果识别出在执行计划的任何时间点，哪些张量必须驻留在内存中。

如果两个张量的活跃范围重叠，则它们干扰。如果张量A在张量B需要计算和存储时仍然活跃，那么A和B会相互干扰，不能共享同一个内存缓冲区。这种干扰关系可以通过干扰图来表示：

• 节点： 表示张量（或它们所需的内存缓冲区）。
• 边： 如果对应的张量相互干扰（它们的活跃范围重叠），则连接两个节点。

一个干扰图示例。两个张量（例如T1和T2）之间的边表示它们的活跃范围重叠，阻止它们共享同一内存缓冲区。具体的重叠取决于执行计划。

通过共享进行内存分配

静态内存规划的目标是将张量分配到物理内存缓冲区，使得任意两个相互干扰的张量不会被分配到同一个缓冲区，同时最小化所有已分配缓冲区的总大小。这个问题类似于干扰图上的图着色：

• 为每个节点（张量）分配一个“颜色”（代表一个独立的内存缓冲区）。
• 约束：通过边连接的节点（相互干扰的张量）必须获得不同的颜色。
• 目标：最小化与所用颜色相关的“成本”。

尽管标准图着色旨在最小化颜色数量，但内存分配旨在最小化峰值内存使用。如果所有张量大小相同，最小化缓冲区数量就能达到此目标。然而，张量的大小差异很大。因此，算法必须最小化所有执行点上并发分配的最大内存。这使得问题比标准图着色更复杂，尽管干扰图仍然是一个核心思想。

由于找到绝对最优的内存分配是NP难问题，编译器采用启发式算法。常见策略包括：

1. 贪心分配（例如，首次适应）：
   • 按特定顺序处理张量（例如，按活跃范围的开始时间、结束时间或大小排序）。
   • 对于每个张量，遍历已分配的缓冲区。将张量分配给第一个足够大且当前可用（即在此张量活跃范围内不包含任何干扰张量）的缓冲区。
   • 如果未找到合适的缓冲区，则分配一个新缓冲区。
   • 存在“最佳适应”（选择最小的足够大的缓冲区）或“最差适应”等变体，以计算复杂度换取潜在的碎片减少。
2. 基于偏移的分配：
   • 不使用独立的逻辑缓冲区，而是维护一个（或少量）大型内存池。
   • 根据估计的峰值使用量计算所需内存池的大小。
   • 为每个张量在内存池中分配一个特定的字节偏移量。
   • 主要挑战仍然是为相互干扰的张量找到不重叠的偏移量分配，同时最小化总内存池大小（所有张量的max(offset + size)）。这通常使用类似于贪心方法但操作于偏移量的启发式算法来解决。

计算峰值内存和偏移量

考虑张量T1 (100KB)、T2 (50KB)、T3 (80KB) 和 T4 (100KB)，其活跃范围如下（简化时间单位）：

• T1: [0, 10)
• T2: [2, 12)
• T3: [3, 8)
• T4: [10, 15)

干扰：

• T1与T2和T3相互干扰。
• T2与T1和T3相互干扰。
• T3与T1和T2相互干扰。
• T4不与T1、T2或T3相互干扰。

简单分配需要 100 + 50 + 80 + 100 = 330KB。

使用静态规划（例如，基于偏移）：

• 时间 [0, 2)：活跃：{T1}。内存：100KB。将T1分配在偏移量0。内存池大小：100KB。
• 时间 [2, 3)：活跃：{T1, T2}。内存：100+50=150KB。T1在偏移量0。将T2分配在偏移量100。内存池大小：150KB。
• 时间 [3, 8)：活跃：{T1, T2, T3}。内存：100+50+80=230KB。T1在0，T2在100。将T3分配在偏移量150。内存池大小：230KB（峰值）。
• 时间 [8, 10)：活跃：{T1, T2}。内存：100+50=150KB。T3的空间 [150, 230) 现在空闲。
• 时间 [10, 12)：活跃：{T2, T4}。内存：50+100=150KB。T1的空间 [0, 100) 空闲。T4需要100KB。可以完美地复用T1的空间（偏移量0）。内存池大小保持230KB。
• 时间 [12, 15)：活跃：{T4}。内存：100KB。T2的空间 [100, 150) 空闲。T4在偏移量0。内存池大小保持230KB。

优化后的峰值内存为230KB，相对于简单分配方式有显著降低。

示例场景中，简单策略（分配唯一缓冲区）与内存重用静态规划之间的峰值内存使用比较。

实际考量

• 对齐： 硬件通常要求内存访问对齐到特定边界（例如，32、64或128字节）。内存规划器必须确保为每个张量分配的偏移量符合其所需的对齐要求，这可能会留下小的未使用空间（内部碎片）。
• 与融合的相互作用： 算子融合减少了中间张量的数量，简化了干扰图，并且通常在明确规划之前就显著降低了峰值内存需求。内存规划通常在融合之后运行。
• 控制流： 带有条件分支或循环的图使静态活跃度分析复杂化，因为确切的执行路径在编译时可能无法得知。通常使用保守分析（假设如果张量在任何可能路径上活跃，则其活跃），这可能会高估内存需求。更先进的技术可能涉及为特定的常见路径规划或使用运行时信息。
• 重新计算与存储： 在内存极度受限的情况下，重新计算张量可能比将其保持在内存中更优。这代表了计算与内存之间的一种权衡，通常由单独的优化阶段处理。

静态内存规划是ML编译器中图级别的重要优化。通过分析张量生命周期，并通过基于活跃度分析和干扰图的技术精心组织缓冲区重用，编译器可以大幅减少ML模型的峰值内存消耗。这使得在内存资源有限的硬件上部署更大、更复杂的模型成为可能，有时还能通过增加重用缓冲区留在更快缓存层级的可能性来提高性能。