跳到主要内容
极客日志极客日志
首页博客AI提示词GitHub精选代理工具
|注册
博客列表

目录

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

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

调研了 AI 编译器中 SRAM 内存分配的核心算法，包括线性扫描、区间图着色、最佳/首次适配、仿射模型优化及缓冲区复用等策略。对比分析了 TensorRT、TVM、IREE、TFLite 及厂商 NPU 编译器的实现差异，涵盖静态图与动态形状下的规划方法，以及权重流式加载与计算重叠管理技术。总结指出不同场景下需权衡编译速度、峰值内存与硬件耦合度。

ServerBase发布于 2026/3/27更新于 2026/4/166.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（内存框选分配）

相比 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 ScanTVM/ONNX/TensorRT 基础层快、简单非全局最优
Graph ColoringGlow / TensorRT近最优NP-hard，慢
Best/First FitTFLite Micro / MCU实现简单依赖 heuristics
ILP/AffineMLIR/IREE/高端 NPU全局最优编译开销大
Buffer Reuse所有编译器大幅省 SRAM依赖 schedule
Block-level PlanTransformer 友好后端可运行大模型策略化较强
Dynamic PlannerTensorRT/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/AffineTiling 影响动态 shape 支持
TensorRTLiveness + Graph Coloring⭐⭐⭐⭐中强依赖 Tensor Core tile强
TVMLinear Scan + Region⭐⭐⭐可选强（用户 schedule）中
IREEAffine Planner（MLIR）⭐⭐⭐⭐⭐强中等非常强
TFLiteLiveness + Best-Fit⭐⭐无弱弱
XNNPACKLiveness + First-Fit⭐⭐无弱无
NPU 编译器强硬件耦合：pipeline + tile SRAM⭐⭐⭐⭐⭐少极强中等
极客日志微信公众号二维码

微信扫一扫，关注极客日志

微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog

更多推荐文章

查看全部
  • Python 爬虫实战：爬取飞猪旅行酒店套餐信息
  • 机器人坐标转换中的 RT 矩阵详解
  • 朴素贝叶斯算法详解
  • AI 核心机制深度解析：Skills 如何驱动生产力转型
  • 码云（Gitee）代码推送全流程与实操指南
  • 腾讯 AI 双工具对比：QClaw 与 WorkBuddy 解析
  • Deep Agents：用于复杂任务自动化的 AI 代理框架
  • 地瓜机器人 RDK 系列选型指南：X3 vs X5 vs S100 vs S100P
  • 【选型】地瓜机器人RDK系列选型指南：X3 vs X5 vs S100 vs S100P（含资源对比图）
  • 通义万相 2.1 多模态 AI 生成技术与应用场景分析
  • C++ 最大堆模拟：最后一块石头的重量
  • 通义万相 2.1 多模态 AI 生成技术解析与应用前景
  • llama-cpp-python 完整安装与配置指南
  • VRChat 跨语言交流工具 VRCT 使用指南
  • Java equals() 与 hashCode() 契约详解及容器类（List/Set/Map）核心整理
  • Spring Boot集成WebSocket，实现后台向前端推送信息
  • 多模态赋能情绪理解：Qwen3-VL+LLaMA-Factory 的人脸情绪识别实战
  • webdav-server 轻量级 WebDAV 服务器部署与配置指南
  • WPF + .NET6 WebAPI + SqlSugar 权限管理系统架构与实现
  • Happy Coder：Claude Code 移动端与 Web 客户端

相关免费在线工具

  • 加密/解密文本

    使用加密算法（如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