跳到主要内容AI 编译器 SRAM 内存分配管理算法调研 | 极客日志C++AI算法
AI 编译器 SRAM 内存分配管理算法调研
调研了 AI 编译器中 SRAM 内存分配的核心算法，包括线性扫描、区间图着色、最佳/首次适配、仿射模型优化及缓冲区复用等策略。对比分析了 TensorRT、TVM、IREE、TFLite 及厂商 NPU 编译器的实现差异，涵盖静态图与动态形状下的规划方法，以及权重流式加载与计算重叠管理技术。总结指出不同场景下需权衡编译速度、峰值内存与硬件耦合度。
ServerBase6.1K 浏览 一、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（内存框选分配）
- First Fit：找第一个足够大的空洞。
- Best Fit：找最适合的空洞，减少碎片。
- Worst Fit：为未来保留更多大洞（rare）。
- TFLite Micro
- XNNPACK
- 部分嵌入式 NPU 后端（如 ARM Ethos-U）
- 实现简单
- 对碎片化有改善
- 很适合 运算图较小/算子固定 的 microcontroller NPU
4. 仿射模型优化（Affine Memory Planning）
- 建立一个 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（块级调度分配）
- 以一个 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
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. 图优化带来的间接内存减少算法
- 算子融合（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
- 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 分级缓存
- 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 | ⭐⭐⭐⭐⭐ | 少 | 极强 | 中等 |
微信扫一扫，关注极客日志
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
相关免费在线工具
- 加密/解密文本
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
- RSA密钥对生成器
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
- Mermaid 预览与可视化编辑
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
- Base64 字符串编码/解码
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
- Base64 文件转换器
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online
- Markdown转HTML
将 Markdown（GFM）转为 HTML 片段，浏览器内 marked 解析；与 HTML转Markdown 互为补充。 在线工具，Markdown转HTML在线工具，online
