触碰晶体管前的几何解构：深入浅出 MLIR 中的 Affine Lowering（仿射降级）

触碰晶体管前的几何解构深入浅出 MLIR 中的 Affine Lowering仿射降级在大模型与高性能 AI 算子编译的宏观世界里我们习惯了谈论高层张量RankedTensorType、形状推导inferShapes以及图级别的算子融合DRR Pattern。然而当计算图层面的优化做到了极致编译器最终必须直面一个冷酷的现实底层硬件芯片无论 CPU、GPU 还是专用 NPU并不认识什么是“张量”它们只认识指针、内存偏置Offset以及多层嵌套的循环Loops。在 MLIR 的多级方言谱系中Affine仿射方言扮演着多维张量数学向微观硬件执行跨越的桥头堡。而Affine Lowering仿射降级则是将宏观的算子计算彻底打碎、翻译成确定性多维多层循环以及精细化内存寻址的核心通道。本文将带你深入这场代码走向底层硬件的跃迁之旅。1. 什么是 Affine 方言与 Affine Lowering在传统编译器如老式 GCC中多层嵌套循环只是普通的控制流for循环里套for循环。编译器很难在编译期断定这些循环的访存轨迹。MLIR 的Affine方言另辟蹊径。它施加了一种强烈的数学约束循环的步长、边界以及内存索引Indexing必须是循环迭代变量Induction Variables和全局符号Symbols/Constants的线性组合即仿射函数Affine Functions。Affine Lowering 的本质它是将高层的结构化算子如Linalg方言的矩阵乘法或完全展平的高层算子降级Lowering表达为带有严格多面体数学模型约束Polyhedral Model的多层嵌套循环与底层内存块MemRef读写组合的过程。2. 文本 IR 视角多维张量到仿射循环的无缝蜕变为了让你对 Affine Lowering 产生最直观的工程体感我们来看一个标准的转译案例。假设经过前端翻译和初步图优化后我们得到了一个针对2×42 \times 42×4矩阵的元素级逐加算子属于Linalg结构化控制流方言。在执行Affine Lowering Pass的瞬间IR 发生了如下剧变Lowering 前结构化高层算子// 高层张量已经被 Bufferize 成了底层的物理内存引用 memref // 此时它仍然是一个宏观的通用 Generic 算子 linalg.generic { indexing_maps [#map0, #map0, #map0], iterator_types [parallel, parallel] } ins(%arg0, %arg1 : memref2x4xf32, memref2x4xf32) outs(%arg2 : memref2x4xf32) { ^bb0(%in1: f32, %in2: f32, %out: f32): %0 arith.addf %in1, %in2 : f32 linalg.yield %0 : f32 }Lowering 后高度规范的 Affine 仿射嵌套循环// 宏观算子被彻底平铺拆解为两层显式嵌套循环 // affine.for 规定了其边界在编译期高度可预测 affine.for %i 0 to 2 { affine.for %j 0 to 4 { // 精确的仿射寻址映射由多面体模型保证绝无内存越界和指针别名冲突 %a affine.load %arg0[%i, %j] : memref2x4xf32 %b affine.load %arg1[%i, %j] : memref2x4xf32 %sum arith.addf %a, %b : f32 affine.store %sum, %arg2[%i, %j] : memref2x4xf32 } }3. C 源码视角编写一个 Affine Lowering 通道在真实的编译器后端开发中如果你写了一个自研 NPU 的定制高层算子想要让它降级到Affine循环你需要在 C 的 Lowering Pass 中利用PatternRewriter纯手写对循环骨架的搭建。下面是触发多层affine.for生成的经典 C 代码范式#includemlir/Dialect/Affine/IR/AffineOps.h#includemlir/Dialect/Arith/IR/Arith.husingnamespacemlir;// 假设我们正在将自研的 my_npu.vector_add 算子转换为 Affine 循环structVectorAddLowering:publicOpRewritePatternmy_npu::VectorAddOp{usingOpRewritePatternmy_npu::VectorAddOp::OpRewritePattern;LogicalResultmatchAndRewrite(my_npu::VectorAddOp op,PatternRewriterrewriter)constoverride{Location locop.getLoc();Value memrefAop.getLhs();// 输入内存块Value memrefBop.getRhs();Value memrefCop.getOutput();// 1. 获取明确的几何形状边界 (例如 2x4)automemrefTypememrefC.getType().castMemRefType();int64_tdim0memrefType.getShape()[0];// 2int64_tdim1memrefType.getShape()[1];// 4// 2. 利用特殊的闭包构建器buildAffineLoopNest一击喷涌出多层嵌套循环// 它会隐式创建迭代变量 (Induction Variables) 塞给 lambda 表达式llvm::SmallVectorint64_t,2lowerBounds({0,0});llvm::SmallVectorint64_t,2upperBounds({dim0,dim1});llvm::SmallVectorint64_t,2steps({1,1});buildAffineLoopNest(rewriter,loc,lowerBounds,upperBounds,steps,[](OpBuildernestedBuilder,Location l,ValueRange ivs){// ivs[0] 代表外层循环变量 %iivs[1] 代表内层循环变量 %j// 3. 在最内层 Block 中发射标准的仿射加载与存储算子Value anestedBuilder.createaffine::AffineLoadOp(l,memrefA,ivs);Value bnestedBuilder.createaffine::AffineLoadOp(l,memrefB,ivs);Value sumnestedBuilder.createarith::AddFOp(l,a,b);nestedBuilder.createaffine::AffineStoreOp(l,sum,memrefC,ivs);});// 4. 完美擦除原有高层算子节点rewriter.eraseOp(op);returnsuccess();}};4. 为什么编译管线必须经过 Affine 层的严苛洗礼很多同学会问既然都要变循环为什么不直接降级到更通用的scf.for标准控制流循环而一定要在Affine停留因为Affine携带的线性数学约束能够激活编译器里最硬核的多面体循环优化空间Polyhedral Optimizations。一旦代码进入了标准Affine谱系以下四大顶级性能压榨 Pass 就可以直接无锁执行循环平铺Loop Tiling为了完美压榨芯片的片上缓存L1/L2 Cache 或 SRAM把大图切成硬件刚好能吞下的小方块Tiles。在Affine级别平铺 Pass 可以完全基于数学不等式运算自动计算出最佳的搬运尺寸Tile Size并重构循环结构。循环展开与重组Loop Unrolling Permutation为了填满底层向量寄存器的发射空闲Issue SlotsAffine可以极其安全地自动将内层循环展开。甚至当它发现行存和列存步长不匹配时可以直接将外层循环和内层循环对调位置Loop Interchange大幅提升合并访存Memory Coalescing的效率。绝对安全的自动向量化Vectorization由于仿射映射消除了由于指针指针偏移可能带来的内存混淆Aliasing编译器能够百分之百确信不同的循环迭代之间没有任何数据依赖冲突。这就允许底层后端 Pass 能够毫无顾忌地将affine.load一击提炼为最强悍的硬件向量SIMD加载指令。总结一句话概括Affine Lowering 是 AI 编译器把人类宏观的张量算子愿景降维打击成冷酷、高确定性的硬件循环时空的铁轨。它利用精妙的多面体数学模型给原本无序的循环访存空间拉上了防线从而赋予了中后期编译器 Pass 极致施展“空间变幻”的自由度。经历过 Affine 洗礼的代码不仅保留了严谨的逻辑语义更为最终降级到 LLVM IR 或专用硬件指令集发射 Tensor Core 暴击奠定了无可挑剔的底层骨架。