当前位置：首页 > news >正文

生存周期的锚定者：深入解析 MLIR 中的内存分配与释放（alloc/dealloc）

news 2026/7/14 0:17:34

生存周期的锚定者：深入解析 MLIR 中的内存分配与释放（alloc/dealloc）

在完成了图层面的算子融合（DRR）、函数的无壁化解构（Inliner）以及仿射多面体循环重构（Affine Lowering）之后，大模型计算图里流动的不再是虚无缥缈的数学张量，而是具备清晰物理边界的内存块（MemRefType）。

然而，要在异构算力芯片（如 GPU、ASIC、自研 NPU）上真正跑起来，编译器必须在图的中后期解决两个极其冷酷且敏感的底层问题：这块数据在显存的什么地方诞生？又在什么时候被彻底销毁？

在 MLIR 的原生体系中，memref.alloc（内存分配）与memref.dealloc（内存释放）扮演着控制物理存储生命周期的绝对主角。在大模型高并发、显存极度紧张的Serving（Serving）场景下，这两个算子的精细化调度直接决定了系统是否会发生 OOM（显存溢出）。

本文将带你深度剖析 MLIR 中内存开辟与释放的文本形态、C++ 生成范式，以及幕后极为硬核的动态显存优化 Pass。

1. 文本 IR 视角：`alloc`与`dealloc`的标准面貌

在 MLIR 的memref方言中，分配与释放成对出现，遵循严格的静态单赋值（SSA）约束。它们在.mlir文件里的标准长相如下：

① 静态显存开辟（适合固定维度的权重或隐藏层）

// 开辟一块存储 128x4096 个 bf16 元素的连续物理显存，并将其绑定在高速静态缓存 (@SRAM) 上 %alloc_static = memref.alloc() : memref<128x4096xbf16, space: @SRAM> // ... 对该内存块 %alloc_static 进行矩阵乘法读写计算 ... // 显式销毁该内存句柄，释放物理空间 memref.dealloc %alloc_static : memref<128x4096xbf16, space: @SRAM>

② 动态显存开辟（适合大模型变长 Token 推理）

当维度中含有未知量?时，memref.alloc必须显式地将前端在运行时计算出来的动态轴长度作为操作数（Operands）喂进来：

// %seq_len 是运行时动态拿到的 Token 长度 %alloc_dynamic = memref.alloc(%seq_len) : memref<?x4096xf32> // ... 伴随计算 ... memref.dealloc %alloc_dynamic : memref<?x4096xf32>

2. C++ 源码视角：在 Lowering 阶段精准发射分配指令

在编写将高层张量方言降级到物理内存层的 Pass（即Bufferization Pass）时，工程师需要使用OpBuilder在 C++ 侧亲手织就这一对生命周期算子。

下面展示了如何捕获一个高层的数学计算算子，并为其在前后两端分别编织alloc和dealloc的经典 C++ 范式：

#include"mlir/Dialect/MemRef/IR/MemRef.h"usingnamespacemlir;// 假设我们正在将一个自研的高层算子降级，并为其规划物理内存LogicalResultlowerToPhysicalBuffer(Operation*op,PatternRewriter&rewriter){Location loc=op->getLoc();// 1. 确定输出张量的精确几何外形 (假设推导出的结果是静态的 1x4096)SmallVector<int64_t,2>shape({1,4096});automemrefType=MemRefType::get(shape,rewriter.getF32Type());// 2. 在计算发生前，抢先发射 memref::AllocOp 算子，在显存池里“占座”autoallocOp=rewriter.create<memref::AllocOp>(loc,memrefType);Value bufferResult=allocOp.getResult();// 拿到了真真切切的内存引用句柄// 3. 呼叫下游算子（例如发射底层的计算 Kernel），将计算结果灌入 bufferResult 中rewriter.create<my_npu::ComputeKernelOp>(loc,...,bufferResult);// 4. 将焊枪口（Insertion Point）暂时移到当前 Block 的最末尾（如 return 算子之前）// 确保在该计算块生命周期结束的最后一刻，精准发射销毁指令，防止显存泄漏OpBuilder::InsertionGuardguard(rewriter);rewriter.setInsertionPoint(op->getBlock()->getTerminator());rewriter.create<memref::DeallocOp>(loc,bufferResult);returnsuccess();}

3. 大模型时代的硬核对抗：Buffer 反置化与自动分配优化

在大模型百亿参数的高频吞吐下，如果编译器傻傻地每遇到一个算子就无脑发射一对alloc/dealloc，那底层硬件将会面临两场灾难：

频繁调用cudaMalloc/cudaFree导致严重的运行时耗时（Runtime Overhead）。
内存碎片化导致显存利用率暴跌。

为了解决这一痛点，MLIR 在alloc/dealloc层级沉淀了两个常驻的工业级硬核优化 Pass：

① 自动生命周期提升（Buffer Deallocation Pass）

有时候，前端转译出来的alloc它的配套dealloc位置放得并不合理，或者根本漏掉了（导致显存泄漏）。
MLIR 官方提供了一个极其强悍的通用数据流分析通道——BufferDeallocationPass。它通过在控制流图（CFG）上执行深度的不透明逃逸分析（Escape Analysis），顺藤摸瓜地追踪每一个%alloc产生的值穿透了哪些分支、走过了哪些循环。随后，它会在该变量生命周期彻底死掉（Liveness Dead）的最优绝对临界点，自动帮你插入memref.dealloc。这就让编译器具备了类似 Rust 语言的自动生命周期管理能力（RAII），既安全又极致紧凑。

② 显存原地重用（In-place & Reuse Optimization）

在大模型（如 Transformer 架构）的层级推进中，中间流转的临时张量（激活值 Activation）非常大，但很多算子的生存周期是高度互斥的（即算子 B 启动的时候，算子 A 已经计算完毕死掉了）。

[ 原始混乱状态 ] [ 经内存重分配 Pass 优化后 ] %A = memref.alloc() : memref<4096xf32> %A = memref.alloc() : memref<4096xf32> // ... %A 的计算与终结 ... // ... %A 的计算与终结 ... memref.dealloc %A %B = memref.alloc() : memref<4096xf32> // 完美复用 A 的物理空间，省去了 // ... %B 的计算 ... // 一次分配开销与 4096 字节的物理容积！ memref.dealloc %B // 直接对 %A 的句柄进行指针复用与写覆盖 memref.dealloc %A

在alloc/dealloc层级，MLIR 的分配优化器（Memory Allocation Optimizer）会绘制一张全局的内存区间重叠图（Interference Graph）。如果发现%A和%B容积相同且生命周期完全没有交集，编译器会在 Lowering 阶段把%B的开辟算子抹除，直接让下游算子复用%A的物理指针地址，实现零开销的原地写覆盖（In-place Override）。