Mini-Infer (23): 内存优化的黑魔法 — 静态内存规划与贪心着色

Mini-Infer (23): 内存优化的黑魔法 — 静态内存规划

1. 核心架构：内存规划的四个步骤

我们的内存规划器 (MemoryPlanner) 并不是在运行时动态申请释放内存，而是在推理之前（编译期） 对图进行静态分析，计算出一份最优的内存分配方案。

这个过程分为四个严谨的步骤：

生命周期分析 (Liveness Analysis)：确定每个 Tensor 何时生、何时死。
构建冲突图 (Interference Graph)：确定哪些 Tensor 绝对不能共用内存。
贪心着色 (Greedy Coloring)：分配内存池 ID，让不冲突的 Tensor 复用同一个 ID。
生成方案 (Plan Generation)：输出最终的内存池配置。

2. 第一步：生命周期分析 (`LivenessAnalyzer`)

要复用内存，首先要知道 Tensor 的有效期。

出生 (Birth): 生产该 Tensor 的算子开始执行的时间点。
死亡 (Death): 最后一个消费该 Tensor 的算子执行结束的时间点。

我们定义了 TensorLifetime 结构体来存储这些信息：

struct TensorLifetime {
    std::string name;
    size_t size_bytes;   // 字节大小
    int birth_time;      // 拓扑序中的出生时间
    int death_time;      // 拓扑序中的死亡时间
    bool is_persistent;  // 是否为持久化Tensor（如权重、图输入输出），这类不能复用
    int pool_id;         // 最终分配的内存池ID
};

实现逻辑 (analyze 函数):

拓扑排序: 我们首先对 Graph 进行拓扑排序，得到一个线性的节点执行序列。节点的下标（0, 1, 2…）就是逻辑上的“时间”。
确定生产者/消费者: 遍历图，找出每个 Tensor 是由谁产生的，又被谁消费了。
计算时间区间:
- birth_time = 生产者节点的下标。
- death_time = 所有消费者节点下标的最大值（即最后一次使用的时间）。
持久化标记: 图的输入 (graph->is_input)、输出 (graph->is_output) 以及权重 (node->inputs().empty()) 被标记为 is_persistent。这些 Tensor 需要一直存在，不能参与内存复用。

3. 第二步：构建冲突图 (`InterferenceGraph`)

知道了生命周期后，我们需要判断哪些 Tensor 冲突 (Interference)。冲突意味着：两个 Tensor 的生命周期有重叠，它们必须同时存在于内存中，因此不能复用同一块内存。

冲突判断逻辑 (lifetimes_overlap): 如果 Tensor A 和 Tensor B 满足以下条件，则它们冲突：

A 在 B 死亡之前出生且 A 在 B 出生之后死亡。
或者简单来说：两个时间区间 [start, end] 有交集。

bool MemoryPlanner::lifetimes_overlap(const TensorLifetime& a, const TensorLifetime& b) const {
    // 经典区间重叠判断
    return !(a.death_time < b.birth_time || b.death_time < a.birth_time);
}

我们构建了一个 InterferenceGraph 类，它本质上是一个无向图。图中的节点是 Tensor，边表示两个 Tensor 冲突。

4. 第三步：贪心着色算法 (`greedy_coloring`)

现在问题转化为了经典的 图着色问题 (Graph Coloring)：

节点: Tensor
颜色: 内存池 (Memory Pool)
规则: 相邻节点（冲突的 Tensor）不能涂相同的颜色。
目标: 使用最少的颜色（内存池），且尽量减少总内存大小。

由于图着色是 NP-Hard 问题，我们使用 贪心算法 来获取近似最优解。

算法流程：

排序策略: 我们将所有 Tensor 按 大小降序 (Size Descending) 排列。
- 原因: 这是一个经典的启发式策略（Best Fit）。先安排大块内存，小块内存更容易塞进大块内存留下的空隙（时间间隙）中。如果先安排小的，可能会导致大块内存无法复用，从而增加碎片。
分配循环: 遍历排序后的 Tensor：
- 检查当前已创建的所有内存池 Pool 0, Pool 1...。
- 对于每个池，检查：“这个池子里目前分配的所有 Tensor，有没有和当前 Tensor 冲突的？”
- 利用 InterferenceGraph::has_edge 进行快速查找。
- 如果找到一个不冲突的池，就复用它，并更新该池的大小（max(old_size, current_tensor_size)）。
- 如果所有池都冲突，则新建一个内存池。

// 核心代码片段
MemoryPlan MemoryPlanner::greedy_coloring(...) {
    // 1. 排序
    std::sort(sorted_lifetimes.begin(), sorted_lifetimes.end(),
              [](const TensorLifetime* a, const TensorLifetime* b) {
                  return a->size_bytes > b->size_bytes; 
              });

    // 2. 贪心分配
    for (auto* tensor : sorted_lifetimes) {
        int pool_id = find_available_pool(*tensor, graph, plan);
        if (pool_id == -1) {
            // 新建池
            // ...
        } else {
            // 复用池
            plan.tensor_to_pool[tensor->name] = pool_id;
            // 池的大小可能会因为塞入了一个更大的 Tensor 而膨胀
            plan.pools[pool_id].size_bytes = std::max(plan.pools[pool_id].size_bytes, aligned_size);
        }
    }
}

5. 内存对齐 (`Alignment`)

在分配内存大小时，我们必须考虑对齐。CPU 上的 AVX/SSE 指令通常要求内存地址是 16 或 32 字节对齐的；在 GPU 上通常要求 256 字节对齐。

Mini-Infer 默认使用 256 字节对齐，这既满足了 CPU SIMD 的需求，也为未来支持 CUDA 预留了空间。

1
2
3

size_t MemoryPlanner::align_size(size_t size) const {
    return ((size + alignment_ - 1) / alignment_) * alignment_;
}

6. 效果与总结

当我们在推理流程中启用 MemoryPlanner 后，系统会输出详细的规划报告：

╔════════════════════════════════════════════════════════════════════╗
║              Static Memory Planning Result                         ║
╚════════════════════════════════════════════════════════════════════╝

Memory Pools: 3
----------------------------------------
Pool 0: 1024.00 KB
  Tensors (5):
    - conv1_output
    - pool1_output
    - fc1_output
    ...

Pool 1: 512.00 KB
  Tensors (2):
    - conv2_output
    ...

----------------------------------------
Original Memory:  5650.00 KB
Optimized Memory: 1536.00 KB
Memory Saving:    72.81%