从图循环到寄存器发射：深入浅出 MLIR 中的 Loop Lowering 从图循环到寄存器发射深入浅出 MLIR 中的 Loop Lowering在完成了内存生命周期的锚定alloc/dealloc与局部转换Partial Conversion后编译管线已经成功为计算图穿上了物理内存的骨架。然而大模型中最庞大的矩阵计算如MatMul、Attention本质上依然被包裹在宏观的图级循环或多维度的结构化操作中。在真正的芯片CPU、GPU、NPU内没有任何一个硬件执行单元能直接吞下一个“多维抽象循环”。硬件渴望的是一维化的迭代、基础的分支条件跳转以及最终被展平的串行/并行指令流。在 MLIR 的多级方言谱系中Loop Lowering循环降级扮演着将高级几何循环转化为最朴素的机器控制流的关键角色。本文将带你深度剖析循环变换如何一步步逼近晶体管。1. MLIR 循环方言的阶梯架构与传统编译器一刀切的控制流不同MLIR 内部构建了一个优雅的循环方言阶梯。Loop Lowering 的本质就是代码在这座阶梯上由高到低、由抽象到具象的垂直降落过程[ 阶梯顶层Linalg / TOSA 算子 ] -- 隐藏了循环只表达宏观矩阵计算 │ ▼ (Linalg-to-Loops Lowering) [ 阶梯中层Affine 方言循环 ] -- 严格的多面体数学模型边界高度可预测 │ ▼ (Affine-to-SCF Lowering) [ 阶梯下层SCF (标准控制流) 方言 ] -- 最纯粹的显式结构化循环 (scf.for / scf.while) │ ▼ (SCF-to-ControlFlow/LLVM Lowering) [ 阶梯底层ControlFlow / LLVM IR ] -- 彻底打碎只剩 Basic Block、条件跳转与裸指针在大模型推理编译中让循环在不同层级停留是为了在最适合的阶段做最适合的优化如在Affine层做平铺在LLVM层做发射。2. 文本 IR 视角scf.for向底层跳转的冷酷蜕变为了让你对 Loop Lowering 产生最直观的工程体感我们来看看一个大模型中极其常见的scf.for标准结构化循环在执行SCF-to-ControlFlow Lowering Pass的瞬间IR 发生的颠覆性剧变。Lowering 前高级、规整的scf.for结构此时循环变量、步长Stride、上界Upper Bound以高层操作数的形式清晰并存。// 这是一个简单的单层结构化循环 func.func scf_loop(%lb: index, %ub: index, %step: index) { scf.for %iv %lb to %ub step %step { // 循环体Loop Body被安全包裹在当前算子的专属 Region 内部 my_npu.do_work(%iv) : (index) - () } return }Lowering After彻底平铺、直面硬件的“控制流图CFG”形态经历 Lowering 后原本高大上的scf.for算子彻底解体。Region 被无情撕裂代码被平铺成了数个互相关联的基础块Basic Blocks通过底层的条件跳转cf.cond_br构成了物理世界里的经典循环环路func.func scf_loop(%lb: index, %ub: index, %step: index) { // 1. 直接越入循环头Header并以形参Block Arguments的形式传递迭代变量 cf.br ^loop_header(%lb : index) ^loop_header(%iv: index): // 2. 硬件级的边界判定检查当前的迭代变量 %iv 是否小于上界 %ub %loop_cond arith.cmpi slt, %iv, %ub : index // 条件跳转如果为 true 进循环体如果为 false 彻底跳出 cf.cond_br %loop_cond, ^loop_body, ^loop_exit ^loop_body: // 3. 执行核心业务逻辑 my_npu.do_work(%iv) : (index) - () // 4. 显式步长递增计算下一次迭代的索引值 %next_iv %iv %step %next_iv arith.addi %iv, %step : index // 闭环跳转带着更新后的索引跳回 Header 重新质检 cf.br ^loop_header(%next_iv : index) ^loop_exit: return }3. C 源码视角利用转换框架实现循环定向爆破在为自研硬件编写编译后端时我们经常需要把定制的循环算子下推到通用控制流。MLIR 官方通过SCFToControlFlow库提供了标准的转换模式Conversion Patterns。下面展示了如何在 C 后端 Pass 中激活这一降级通道#includemlir/Conversion/SCFToControlFlow/SCFToControlFlow.h#includemlir/Dialect/ControlFlow/IR/ControlFlowOps.h#includemlir/Dialect/SCF/IR/SCF.husingnamespacemlir;structLowerLoopPass:publicPassWrapperLowerLoopPass,OperationPassModuleOp{voidrunOnOperation()override{ModuleOpmodulegetOperation();MLIRContext*contextgetContext();// 1. 设定转换目标ConversionTargettarget(*context);// 宣告底层控制流方言是合法的target.addLegalDialectcf::ControlFlowDialect();target.addLegalDialectarith::ArithDialect();// 强力通缉全图所有的 scf.for / scf.while 必须在此处被消灭非法target.addIllegalDialectscf::SCFDialect();// 2. 手机官方提供的标准循环降级模式RewritePatternSetpatterns(context);populateSCFToControlFlowConversionPatterns(patterns);// 3. 启动全转换Full Conversion一击将结构化循环打碎为 CFG 状态机if(failed(applyFullConversion(module,target,std::move(patterns)))){signalPassFailure();}}};4. 大模型时代 Loop Lowering 阶段的极致性能压榨当循环被编译器一步步降级并展平时大模型底层的算力压榨进入了最硬核的向量化与多核发射阶段① 终极向量化Vectorization Inlining在高级循环如affine.for或scf.for完全被打碎成 CFG 之前的最后一公里编译器会对其执行向量化 Pass。如果循环的边界是固定位宽如 Transformer 的多头注意力维度dk64d_k64dk​64Loop Lowering 会直接配合Vector方言把原本需要循环 64 次的标量加法塌陷并转换为一条单指令多数据SIMD的 512 位向量加法指令循环次数直接骤降为原来的几分之一甚至彻底消除。② 异步控制流与流水线掩盖Pipeline Hiding大模型专用 NPU 在执行循环时往往伴随着高频的片外显存HBM到片内缓存SRAM的数据搬运。在 Loop Lowering 阶段编译器如果探测到循环体内含有异步搬运指令会执行循环展开与流水线重排Loop Unrolling and Pipelining。它会打破单次迭代的边界让硬件在计算第NNN次循环矩阵乘法的同时隐式发射第N1N1N1次迭代所需的权重 DMA 搬运指令。这种控制流级别的微调能完美用计算时间掩盖访存延迟让芯片的执行单元永远处于满载暴击状态。总结一句话概括Loop Lowering 是 AI 编译器将宏观的高维几何时空彻底揉碎并降维组织成微观的基础块跳转与流水线指令流的蜕变之路。它通过将层级分明的结构化循环scf.for精准定向爆破为硬件原生认可的条件跳转网络CFG打通了高级数学算法走向物理芯片寄存器的最后一道控制门槛。深刻精通循环降级的流转机理与向量化时机是每一个致力于大模型推理极致加速、异构算力潜能榨干的 AI 基础设施工程师AI Infra无可动摇的底层核心功底。