AI 编译器中的 SRAM 内存分配管理算法调研

文章目录

[一、AI 编译器中的 SRAM 内存分配管理算法](#一、AI 编译器中的 SRAM 内存分配管理算法)
[1. 静态图的"生存期区间(Liveness-based) 分配"](#1. 静态图的“生存期区间(Liveness-based) 分配”)
- [1.1 Linear Scan（线性扫描分配）](#1.1 Linear Scan（线性扫描分配）)
[2. 区间图着色（Interval Graph Coloring）](#2. 区间图着色（Interval Graph Coloring）)
[3. Memory Planning with Best-Fit / First-Fit（内存框选分配）](#3. Memory Planning with Best-Fit / First-Fit（内存框选分配）)
[4. 仿射模型优化（Affine Memory Planning）](#4. 仿射模型优化（Affine Memory Planning）)
[5. "Buffer Reuse"（缓冲区复用）策略](#5. “Buffer Reuse”（缓冲区复用）策略)
- - [5.1 Inplace 复用](#5.1 Inplace 复用)
  - [5.2 Alias（别名）](#5.2 Alias（别名）)
  - [5.3 Lifetime-based Reuse（生存期复用）](#5.3 Lifetime-based Reuse（生存期复用）)
  - [5.4 Pattern-based Reuse](#5.4 Pattern-based Reuse)
[6. 基于 DAG 拓扑的 Block Allocation（块级调度分配）](#6. 基于 DAG 拓扑的 Block Allocation（块级调度分配）)
[7. 动态形状（Dynamic Shape）下的分配策略](#7. 动态形状（Dynamic Shape）下的分配策略)
- [7.1 Executable Memory Planner](#7.1 Executable Memory Planner)
- [7.2 Chunk-based Allocation](#7.2 Chunk-based Allocation)
[8. Weight Streaming + Compute-Overlapping 管理](#8. Weight Streaming + Compute-Overlapping 管理)
[9. 图优化带来的间接内存减少算法](#9. 图优化带来的间接内存减少算法)
最终总结
[二、AI 编译器 SRAM 内存分配方式对比](#二、AI 编译器 SRAM 内存分配方式对比)
[1. TensorRT](#1. TensorRT)
[2. TVM](#2. TVM)
[3. IREE (Google, MLIR)](#3. IREE (Google, MLIR))
[4. TFLite](#4. TFLite)
[5. XNNPACK](#5. XNNPACK)
[6. 各类厂商专有 NPU 编译器（寒武纪 Cambricon、地平线、海思、联发科、苹果 ANE、三星 DSP 等）](#6. 各类厂商专有 NPU 编译器（寒武纪 Cambricon、地平线、海思、联发科、苹果 ANE、三星 DSP 等）)
- 通用特征
- - [6.1 双缓冲 / 多缓冲 DMA pipeline](#6.1 双缓冲 / 多缓冲 DMA pipeline)
  - [6.2 Multi-bank SRAM 冲突避免](#6.2 Multi-bank SRAM 冲突避免)
  - [6.3 Kernel 模板固定，所以内存规划固定](#6.3 Kernel 模板固定，所以内存规划固定)
  - [6.4 强制算子流水（Pipeline）](#6.4 强制算子流水（Pipeline）)
  - [6.5 全图级 Memory Planner（体现 NPU 特色）](#6.5 全图级 Memory Planner（体现 NPU 特色）)
  - [6.6 部分 NPU 有片内 L0/L1 分级缓存](#6.6 部分 NPU 有片内 L0/L1 分级缓存)
多维度对比表

一、AI 编译器中的 SRAM 内存分配管理算法

AI 编译器的 SRAM 分配目标通常包括：

最小峰值内存（peak memory）
减少带宽访问（减少 DRAM ↔ SRAM 往返）
保持计算并行度
为 kernel 调度与算子融合提供约束

在各大 NPU 编译器（TensorRT、TFLite、TVM、IREE、XNNPACK、各家 SoC NPU 编译器）中，常见的 SRAM 分配策略可分为三类：

1. 静态图的"生存期区间(Liveness-based) 分配"

这是目前 最常见、所有主流编译器都使用的基础算法。

1.1 Linear Scan（线性扫描分配）

原理：

对每个 tensor 做 liveness analysis：得到 start 与 end 点。
按时间顺序扫描分配固定区域。

特点：

O(n) 时间复杂度，编译速度快。
分配结果可预测、确定性强。
对内存块（tensor）的生命周期要求较严。

缺点：

没有全局优化，不是最优解。
在多分支 / 大模型中依然可能碎片高。

绝大多数 NPU 编译器的 baseline 都是它（包括 TVM、ONNX Runtime EP、厂商私有编译器）。

2. 区间图着色（Interval Graph Coloring）

思想来自寄存器分配，将 tensor 视为区间图上的节点，冲突区间染色融合空间。

优点：

更接近最优内存布局。
支持在极小 SRAM 中 squeeze 更大模型。

缺点：

NP-hard，需要启发式（heuristic）。
编译速度较慢，多用于高端设备（GPU/TPU 类编译器）。

TensorRT、Meta 的 Glow、部分 NPU 编译器内部采用该技术。

3. Memory Planning with Best-Fit / First-Fit（内存框选分配）

相比 linear scan 更灵活：

First Fit：找第一个足够大的空洞。
Best Fit：找最适合的空洞，减少碎片。
Worst Fit：为未来保留更多大洞（rare）。

这些算法常用于：

TFLite Micro
XNNPACK
部分嵌入式 NPU 后端（如 ARM Ethos-U）

优点：

实现简单
对碎片化有改善
很适合 运算图较小/算子固定 的 microcontroller NPU

4. 仿射模型优化（Affine Memory Planning）

IREE、MLIR、边缘 AI 编译器常用。

建立一个 tensor offset 分配问题
约束：大小、对齐、生命周期冲突
优化：最小化 peak memory

本质相当于求：

复制代码

minimize  max (offset[t] + size[t])
subject to:
    for all tensors i,j：
        if live-overlap(i,j)：
            offset[i] + size[i] <= offset[j]  or  ...

求解方式：

ILP（Integer Linear Programming）
近似 heuristics
局部优化 + 重排

适用于：

并行度高的算子融合 pipeline
编译器可使用 MLIR 约束系统

5. "Buffer Reuse"（缓冲区复用）策略

几乎所有框架都使用，但算法各家不同。

常见策略：

5.1 Inplace 复用

如 Relu、Clip、Reshape、Transpose 等不改变存储结构的算子。

5.2 Alias（别名）

如残差连接 x + f(x)，可以控制 schedule 使部分共享。

5.3 Lifetime-based Reuse（生存期复用）

结合 liveness 进行 buffer pool 分配。

5.4 Pattern-based Reuse

针对 Conv → BN → Activation 三元组做专用复用模板。

6. 基于 DAG 拓扑的 Block Allocation（块级调度分配）

Transformer 模型中常见。

原理：

以一个 block（如 attention block）为单位规划内存
每个 block 内再进一步做 local plan
block 之间只输出必要的 KV / hidden

优点：

巨大模型下编译速度线性增长
SRAM 限制强的芯片（如 1MB SRAM）仍可运行大 Transformer

属于高层策略，在 TensorRT、IREE、部分 NPU 编译器中可见。

7. 动态形状（Dynamic Shape）下的分配策略

动态 shape 下，固定 offset 失效，需要：

7.1 Executable Memory Planner

运行时根据 shape 计算 offset。

7.2 Chunk-based Allocation

SRAM 分成固定大小 chunk，以 chunk 为单位动态分配（类似 malloc）。

用于：

TensorRT dynamic shape
TVM dynamic shape backend
部分 SoC NPU RT

8. Weight Streaming + Compute-Overlapping 管理

某些 NPU SRAM 不够大（如 512 KB），常采用：

流式加载 weight（DMA）
双缓冲（ping-pong buffer）
计算与数据搬运 overlap
滑动窗口式 activation 缓冲（用于 conv）

涉及：

pipeline buffer scheduling
multi-bank SRAM 碰撞避免

这是硬件耦合度最高的部分，各家实现差异最大。

9. 图优化带来的间接内存减少算法

不直接分配内存、但影响 SRAM 占用：

算子融合（Fusion）
常量折叠（Constant Folding）
通道剪枝（Channel Pruning）
Recompute Tradeoff（用算力换内存）
Operator Tiling（切片到 SRAM 能放下）

这些通常与编译器后端的 memory planner 协同。

最终总结

算法类别	典型编译器	优点	缺点
Linear Scan	TVM/ONNX/TensorRT 基础层	快、简单	非全局最优
Graph Coloring	Glow / TensorRT	近最优	NP-hard，慢
Best/First Fit	TFLite Micro / MCU	实现简单	依赖 heuristics
ILP/Affine	MLIR/IREE/高端 NPU	全局最优	编译开销大
Buffer Reuse	所有编译器	大幅省 SRAM	依赖 schedule
Block-level Plan	Transformer 友好后端	可运行大模型	策略化较强
Dynamic Planner	TensorRT/TVM	支持动态 shape	复杂度高

二、AI 编译器 SRAM 内存分配方式对比

下面从主流六类编译器体系：TensorRT / TVM / IREE / TFLite / XNNPACK / 厂商专有 NPU 编译器

按 算法、buffer 复用、tiling、DMA pipeline、动态 shape 等维度做工程级对比。

1. TensorRT

NVIDIA 的高性能编译器，最重视 吞吐 / bandwidth / 并行调度。

核心策略

Liveness-based + 最小化峰值的全局 memory planner
内部使用 interval graph coloring（区间图着色）增强策略
对 Tensor / Workspace / Activation 进行多级分配
针对每个 engine 构建一次完整规划

强项

优秀的 buffer reuse（复用）
大量 fuse（Conv+BN+Act）、kernel auto selection
Tensor core tile 优化（影响 SRAM 分布）
支持 dynamic shape 的 Executable Memory Planner

特点

workspace（中间 scratch buffer）使用 best-fit
内存规划与 kernel schedule 紧密绑定
支持自动 recompute（算力换空间）

2. TVM

开源社区最灵活的编译器，强调算子图数学优化。

核心策略

Linear Scan 是基础（源码在 src/runtime/vm/memory_manager.cc）
结合：
- lifetime-based reuse（生存期复用）
- region-based allocation（区域分配）
Relay 层 IR 做 global memory planning

强项

算子级 tiling 高度可控（对 SRAM 布局影响巨大）
可以通过 schedule 或 pass 改写内存需求
可以自定义 memory planner（用户可 hack）

特点

无强制 kernel 融合，用户 schedule 决定 SRAM 行为
动态 shape 依赖 runtime allocator（chunk-based）

3. IREE (Google, MLIR)

目前结构最现代的编译器，其内存规划主要基于 MLIR affine 分析。

核心策略

Affine Memory Planning（线性规划近似）
通过 MLIR 的 flow.dispatch 和 HAL 层做两级规划
高质量的 buffer aliasing

强项

最完整的 operand & result aliasing 系统
→ 显著减少 SRAM 占用
MLIR 提供准确的 liveness 与 shape 推理
全球最先进的 dynamic buffer suballocation 模型之一

特点

以 dispatch region 为单元规划
利用 tensor, linalg, flow, vm 多层 IR 协作

4. TFLite

面向移动端（手机），对 SRAM 比 PC 显著敏感。

核心策略

Simple Liveness + Best Fit 为基本算法（见 lite/core）
offline memory planner（静态图）生成 offset
operator pattern 级别的 inplace 支持

强项

针对常见算子（Conv, Depthwise, FC）有定制复用模板
频繁使用 inplace（如 ReLU inplace + reshape 共享 buffer）
对 activation buffer 有特殊优化（减少峰值）

特点

对大型模型不如 TensorRT 灵活
为移动端芯片（DSP/NPU）预留 delegate 空间

5. XNNPACK

Google Android NN 的 CPU kernel 库（高性能 ARM SIMD），其内存方式非常像 MCU。

核心策略

静态 liveness + Best-Fit / First-Fit
非常依赖 operator pattern

强项

针对 common operator 提供数十种复用简化（特别是 ReLU/concat/pool）
计算流图结构固定、可预测
分配器数据结构简单（为速度优化）

特点

不追求全局峰值最小化，只追求快
适合高并发实时场景（例如相机 pipeline）

6. 各类厂商专有 NPU 编译器（寒武纪 Cambricon、地平线、海思、联发科、苹果 ANE、三星 DSP 等）

这些才是 SRAM 分配最激进的体系，一般不对外公开。共同特点是对硬件结构高度耦合：

通用特征

6.1 双缓冲 / 多缓冲 DMA pipeline

weight streaming（权重流式加载）
activation streaming
compute overlap memory staging

6.2 Multi-bank SRAM 冲突避免

算法包括：

bank-aware tiling
conflict-free layout mapping
stride normalization

6.3 Kernel 模板固定，所以内存规划固定

例如：

Conv tile shape 固定
GEMM tile shape 固定
相关 SRAM occupancy 可以提前写死

6.4 强制算子流水（Pipeline）

例如：

复制代码

DMA load → MAC → Store → 下一 tile

因此 SRAM 分配还包括：

ping-pong buffer（A/B）
sliding window buffer（用于 conv）
partial sum accumulation buffer

6.5 全图级 Memory Planner（体现 NPU 特色）

aggressive buffer reuse（非常激进）
activation recompute
tensor slicing 为 SRAM 限制适配

6.6 部分 NPU 有片内 L0/L1 分级缓存

分配器需要为每级 cache 决定：

tile decision
spatial/temporal reuse
scratchpad vs streaming

多维度对比表

编译器	分配核心算法	复用能力	是否做 ILP/Affine	Tiling 影响	动态 shape 支持
TensorRT	Liveness + Graph Coloring	⭐⭐⭐⭐	中	强依赖 Tensor Core tile	强
TVM	Linear Scan + Region	⭐⭐⭐	可选	强（用户 schedule）	中
IREE	Affine Planner（MLIR）	⭐⭐⭐⭐⭐	强	中等	非常强
TFLite	Liveness + Best-Fit	⭐⭐	无	弱	弱
XNNPACK	Liveness + First-Fit	⭐⭐	无	弱	无
NPU 编译器	强硬件耦合：pipeline + tile SRAM	⭐⭐⭐⭐⭐	少	极强	中等