从抽象张量到物理内存的惊险惊险跃迁深入浅出 MLIR 中的 MemRef Lowering在完成了高级算子的代数化简DRR、内联边界擦除Inliner以及多维循环的仿射重构Affine Lowering之后编译管线里原本悬浮在空中的“张量数学形态”已经具备了清晰的循环轨迹。然而无论是多维的有阶张量类型RankedTensorType还是看似有了空间轮廓的仿射算子在真正的计算机体系结构里都是不存在的。底层的物理硬件无论是 CPU、GPU 还是自研 NPU 芯片的片上 SRAM只认两件事一个具体的内存起始地址指针指针和一串连续的、由硬件寻址计算出来的字节偏移量Offset。在 MLIR 的多级方言谱系中将纯数学、无副作用的张量语义彻底打碎并解构为绑定物理存储、带副作用的硬件指针寻址形态的桥梁正是MemRef内存引用类型。而MemRef Lowering内存降级则是整个编译器将代码推向物理硬件执行前的最关键一跃。1. 什么是 MemRef 及其 Lowering 的本质在高级语言中我们认为c a b是创建了一个全新的、独立的张量c符合 SSA 静态单赋值规则。这种数学表达对图优化极为友好但在运行时会导致显存急剧膨胀。MemRefMemory Reference是 MLIR 中内建的物理存储抽象。它代表一个拥有明确维数和数据排布、指向连续内存块Buffer的引用句柄。MemRef Lowering 的本质Bufferization它是编译器将没有物理实体的“张量值Tensor Value”转化为必须在内存中开辟真实空间的“内存块生存周期Buffer Lifecycle”的过程。它的核心任务是完成显存/内存空间的跨算子复用In-place 覆盖写并将多维索引公式在编译期彻底展平Flatten为一维的机器寻址表达式。2. 文本 IR 视角多维内存空间向一维指针寻址的蜕变为了让你对 MemRef Lowering 产生最直观的工程体感我们来看看一个2×42 \times 42×4的二维内存矩阵在执行MemRef-to-LLVM Lowering Pass的瞬间IR 发生的剧烈蜕变。Lowering 前带有高级几何维度的MemRef形态此时的代码虽然已经有了物理空间的底座但其加载算子memref.load仍然携带着多维坐标的信息。// 定义一个输入输出参数均为物理内存引目的函数 func.func matrix_access(%arg0: memref2x4xf32, %i: index, %j: index) - f32 { // 此时编译器依然能看懂 %i 和 %j 分别代表行和列 %res memref.load %arg0[%i, %j] : memref2x4xf32 return %res : f32 }Lowering 后完全展平、直面硬件的LLVM IR形态在经过 Lowering 后原本高大上的多维矩阵完全消失了。取而代之的是一个标准的LLVM 结构体Descriptor以及纯粹的一维硬件寻址公式LinearIndexOffset(i×Stride0)(j×Stride1)\text{LinearIndex} \text{Offset} (i \times \text{Stride}_0) (j \times \text{Stride}_1)LinearIndexOffset(i×Stride0)(j×Stride1)// 矩阵被彻底降级为一个 LLVM 描述符结构体 (包含基础指针、对齐指针、偏置和步长) llvm.func matrix_access(%arg0: !llvm.struct(ptr, ptr, i64, array2 x i64, array2 x i64), %i: i64, %j: i64) - f32 { // 1. 提取出真正的内存首地址裸指针 (Raw Pointer) %base_ptr llvm.extractvalue %arg0[1] : !llvm.struct... // 2. 编译期展平二维寻址计算一维偏移量 // 步长在编译期被确定Row-major 排布下第二维 Stride 轴长为 4 %stride_row llvm.mlir.constant(4 : i64) %tmp llvm.mul %i, %stride_row : i64 %linear_offset llvm.add %tmp, %j : i64 // 3. 硬件裸指针指针偏移 (Gep) 与直接访存加载 (Load) %pixel_ptr llvm.getelementptr %base_ptr[%linear_offset] : (ptr, i64) - ptr, f32 %res llvm.load %pixel_ptr : ptr - f32 llvm.return %res : f32 }3. C 源码视角编写一个 MemRef 降级模式在定制自研 AI 芯片的底层后端如对接你家芯片的 DMA 控制器或者特定的物理寄存器时工程师需要纯手写相应的 C Conversion Pattern将memref相关的操作转换为硬件直接认可的底层方言如LLVM或厂商私有方言。下面展示了如何捕获一个memref.alloc开辟物理显存并将其转换为底层硬件内存分配逻辑的经典范式#includemlir/Conversion/LLVMCommon/ConversionTarget.h#includemlir/Dialect/MemRef/IR/MemRef.h#includemlir/Dialect/LLVMIR/LLVMDialect.husingnamespacemlir;structCustomAllocLowering:publicOpConversionPatternmemref::AllocOp{usingOpConversionPatternmemref::AllocOp::OpConversionPattern;LogicalResultmatchAndRewrite(memref::AllocOp op,OpAdaptor adaptor,ConversionPatternRewriterrewriter)constoverride{Location locop.getLoc();// 1. 获取该内存分配块的强类型信息automemrefTypeop.getType().castMemRefType();// 2. 依据几何大小与元素位宽在编译期精确计算所需开辟的物理字节数 (Bytes)int64_telementBitCountmemrefType.getElementTypeBitWidth();// 例如 FP32 返回 32int64_ttotalElements1;for(int64_tdim:memrefType.getShape()){totalElements*dim;}int64_ttotalBytestotalElements*(elementBitCount/8);// 3. 翻译并替换为底层底层的硬件指针分配或者标准的全局 llvm.malloc 算子autosizeConstantrewriter.createLLVM::ConstantOp(loc,rewriter.getI64Type(),rewriter.getI64IntegerAttr(totalBytes));// 呼叫底层的裸分配指令彻底交出内存控制权autoallocatedPtrrewriter.createLLVM::MallocOp(loc,LLVM::LLVMPointerType::get(getContext()),sizeConstant);// 4. 利用生成的裸指针填充并组装成一个符合上层预期的 MemRef 描述符结构体 (此处略去结构体填空细节)rewriter.replaceOp(op,allocatedPtr.getResult());returnsuccess();}};4. MemRef Lowering 阶段的核心硬核性能优化当编译管线进入MemRef层级时编译器专家可以施展通用的通用图优化彻底无法触碰的微观物理存储优化① 显存/内存空间原地复用In-place Bufferization大模型推理例如 Transformer 的大矩阵相加中有很多属于Element-wise逐元素算子。在 MemRef 级别编译器会执行一种名为以局部性为核心的生存周期分析Liveness Analysis。一旦发现前驱张量输入 A 在当前算子执行完毕后就彻底死掉Dead ValueMemRef Lowering Pass 会直接把输出结果写入输入 A 的物理地址中In-place修改从而在编译期完美擦除了数个吉字节GB的临时显存开辟。② 多级异构存储搬运优化DMA Allocation在现代专用 AI 芯片如 TPU 或专用大模型 ASIC 算力卡中通常拥有多级异构存储空间片外的高带宽显存HBM和片内极度昂贵的高速缓存SRAM。在 MemRef Lowering 阶段编译器可以显式地将内存属性打上不同的存储标签// 明确指示将数据由远端的 HBM通过 DMA 异步控制器搬运View至片内最快的 SRAM 空间 %sram_ref memref.alloc() : memref128x128xf32, SRAM配合memref.dma_start与memref.dma_wait算子编译器能做到在片内计算矩阵乘法的同时在后台隐式利用 DMA 预取下一轮迭代所需的权重数据从而完美掩盖异构计算中的访存延迟将硬件的算力吞吐压榨到物理极限。总结一句话概括MemRef Lowering 是 AI 编译器卸下所有空灵的数学面具、赤裸裸直面底层计算机物理存储结构的惊险一跃。它用严谨的描述符、一维化的寻址展开式以及带副作用的物理指针空间分配将宏观的几何阵列彻底降维打击成晶体管可执行的地址空间。深刻理解 MemRef 的类型边界与指针生成时机是每一个试图精通大模型底层显存复用、片上 SRAM 极限调度以及最终通往 LLVM IR 机器码生成领域的 AI 基础设施工程师AI Infra跨越软硬件鸿沟的成人礼。
