如果你翻过 MLIR 的 dialect 列表大概会注意到 Tensor 和 MemRef 这两个名字总是被摆在一起。很多入门的朋友一开始不太明白既然 tensor 这么好用为什么编译器还要折腾出一个 memref 来直到我第一次跑mlir-opt --one-shot-bufferize看到中间 IR 里冒出一堆memref.load和memref.store才算真正想通——Tensor 描述的是“计算”MemRef 描述的是“内存”而 Bufferization 就是把前者翻译成后者的过程。这篇笔记是“MLIR 入门笔记”的第二篇。第一篇我讲了 dialect、pass 和基本 IR 结构这次直接切入整个 MLIR 体系里最容易劝退新人的一道坎从 Tensor 到 MemRef。我会尽量用一个实际例子讲清楚 Bufferization 到底在干什么MemRef 类型里的 offset、sizes、strides 是怎么来的以及那些在网上搜半天也搜不到的布局不匹配报错背后到底是什么原因。适合正在学 MLIR、想搞懂one-shot-bufferize输出、或者被memref类型定义折磨过的朋友。1. 为什么说 Bufferization 是通往底层的中转站1.1 Tensor 是“值”MemRef 是“存储”MLIR 里的tensor类型本质上是一个不可变的抽象值。你在 IR 里写%0 arith.addf %lhs, %rhs : tensor4xf32其含义是“这两个 tensor 相加产生一个新的 tensor”。编译器并不关心这个 tensor 在内存里怎么摆也不关心数据到底有没有真实存在过。它只是一条 SSA 值的定义跟数学里的向量运算很像。而memref类型不同。memref4xf32表示“一块可以被读写、可以被取地址、可以被别名分析插件跟踪的内存区域”。它是有副作用的读它可能是memref.load写它可能是memref.store而且多次访问同一个memref可能拿到不同的数据。这种类型才是真正靠近机器模型的表示。所以你可以把 Bufferization 理解为一次“语义降级”从“值世界”掉进“存储世界”。掉进去以后你才开始真正面对内存分配、数据复用、生命周期这些底层问题。这也是为什么大家常说 MLIR 适合做多级编译——它允许你在高层次的 tensor 语义上做优化完了再一次性降到 memref 上做代码生成。1.2 一次转换解决的核心问题Bufferization 要解决的核心问题其实很朴素一个计算图上到处都是张量操作而这些张量最终一定得落到内存里。谁来分配内存什么时候可以复用旧 buffer哪些写操作可以原地进行这些如果不提前规划好代码生成的效率会非常难看。举个例子你要写一个linalg.generic把一个tensor4xf32和一个tensor4xf32相加结果放在一个新的tensor4xf32里。如果在底层每次加法都重新分配一块内存性能会惨到没法看。所以编译器需要知道这个结果 tensor 能不能复用输入之一的 buffer如果两个输入都不再被使用完全可以把结果直接写到第一个或第二个输入里。这听起来像是一个简单的“死代码分析”但在复杂 IR 里它涉及到别名、逃逸分析、以及跨函数边界的参数传递。MLIR 里提供的one-shot-bufferizepass就是要用一套统一的算法把所有 tensor 操作一次性转换为 memref 操作并在这过程中尽量选择就地更新。这个“一次性”非常关键因为它避免了分成多个 pass 逐步替换时前面 pass 做的内存决策干扰后面 pass 的优化空间。1.3 在 MLIR 分层架构里的位置MLIR 的多级策略是不同抽象级别的 dialect 共存于同一个 IR 里。tensor和memref通常不是互斥的一个模块里可以既有 tensor 表示的计算也有 memref 表示的访存操作。Bufferization 之后tensor和memref出现在同一个模块里是很正常的尤其在函数边界处。这种混合状态并不是 bug而是设计。很多 pass 只需要处理 tensor 部分处理完以后再用bufferization.to_memref做显式转换另一些 pass 只关心 memref 的 layout 和内存空间根本不需要管 tensor。所以理解 Bufferization也就理解了整个 MLIR 代码生成路径里“抽象优化”和“底层优化”的交接点。2. 理解 MemRef 类型与内存布局2.1 MemRefType 的五个组成部分第一次在错误信息里看到memref1x2x?xf32, affine_map(d0, d1, d2)[s0] - (d0 * s0 d1 * 2 d2), 3这样一串东西很多人第一反应是头大。但实际上 MemRefType 的组成是有规律的拆开看就五块元素类型这里是f32。形状1x2x?表示第一维静态长度为 1第二维静态长度为 2第三维运行时才知道。布局映射affine_map...描述逻辑下标如何映射到线性内存地址。内存空间数字3或者 GPU shared memory、全局内存等命名空间。位宽和格式相关的额外属性比如memref... { ... }里的属性字典。记忆方式很简单MemRef 就是“一块内存 如何访问它的说明”。前两部分决定数据长什么样后两部分决定数据在哪、怎么找。2.2 offset、sizes、strides 怎么读在 MLIR 的降级过程中你会经常看到带 strides 的写法比如memref?x?xf32, strided[?, ?], offset: ?这里strided[?, ?]表示两个维度的步长都是运行时值offset也是运行时值。stride的意思是沿着某个逻辑维度下标加 1在内存地址上要跳过多少个元素。比如一个二维memref4x4xf32如果按行主序存放那么strides就是[4, 1]因为第一维每加一格要跳过一整行的 4 个元素第二维每加一格只要跳过 1 个元素。很多初学者会把“形状”和“步长”混在一起。其实形状只是逻辑上的尺寸步长规定了实际内存里的排布方式。也就是说两个memref4x4xf32完全可能有不同的 strides一个是行主序一个是列主序甚至一个是经过 padding 的第二个维度不是 4 而是 5因此步长是[5, 1]。2.3 layout map 与仿射映射的作用除了strided...这种简写MemRef 还可以用完整的仿射映射来描述布局。比如memref2x4xf32, affine_map(d0, d1) - (d0 * 8 d1 * 2)这个映射表示第一维的下标d0乘 8第二维的下标d1乘 2。也就是说这个 buffer 的物理布局并不是紧凑的 2x4而是一个每行间隔 8 个元素、每列间隔 2 个元素的结构。这种映射最典型的用途是描述“padding 后的矩阵”或者“多个小矩阵打包在一起”的布局。无论哪种写法布局信息对编译器的影响都是决定性的。后续的向量化、循环展开、数据搬移都要根据 strides 来算地址。如果你的 pass 生成了错误的 strides后面的代码生成阶段就会“带着错误的地图跑路”最后跑到很深的层面才爆炸这也是为什么布局不匹配的错误往往是整个 MLIR 开发里最常见的头疼问题。3. 从 Tensor 到 MemRefBufferization 的核心机制3.1 一个最简单的 linalg 加法示例先写一个非常朴素的 linalg 加法。假设我有两个tensor4xf32想逐元素相加返回结果func.func add_two(%lhs: tensor4xf32, %rhs: tensor4xf32) - tensor4xf32 { %init linalg.init_tensor [4] : tensor4xf32 %0 linalg.generic { indexing_maps [ affine_map(i) - (i), affine_map(i) - (i), affine_map(i) - (i) ], iterator_types [parallel] } ins(%lhs, %rhs : tensor4xf32, tensor4xf32) outs(%init : tensor4xf32) { ^bb0(%lhs_scalar: f32, %rhs_scalar: f32, %init_scalar: f32): %sum arith.addf %lhs_scalar, %rhs_scalar : f32 linalg.yield %sum : f32 } - tensor4xf32 return %0 : tensor4xf32 }这个 IR 里的所有张量都还没有物理内存。linalg.init_tensor本质上是“一个尚未分配内存的抽象占位”它告诉编译器这里需要一个 4 元素张量作为输出缓冲区的雏形。只有当你开始做 bufferization这个占位才会真正变成一次分配。3.2 one-shot-bufferize 做了什么把上面的 IR 保存成add.mlir然后执行mlir-opt --one-shot-bufferizebufferize-function-boundaries1 add.mlir你会看到输出变成类似下面这样的东西func.func add_two(%lhs: memref4xf32, %rhs: memref4xf32) - memref4xf32 { %alloc memref.alloc() : memref4xf32 %0 linalg.generic { indexing_maps [ affine_map(i) - (i), affine_map(i) - (i), affine_map(i) - (i) ], iterator_types [parallel] } ins(%lhs, %rhs : memref4xf32, memref4xf32) outs(%alloc : memref4xf32) { ^bb0(%lhs_scalar: f32, %rhs_scalar: f32, %init_scalar: f32): %sum arith.addf %lhs_scalar, %rhs_scalar : f32 linalg.yield %sum : f32 } - memref4xf32 return %alloc : memref4xf32 }注意几个变化。第一函数参数从tensor4xf32变成了memref4xf32这就是bufferize-function-boundaries1的效果它把函数边界也纳入了 bufferization。第二linalg.init_tensor变成了memref.alloc说明这里真的分配了一块能放 4 个f32的内存。第三linalg.generic的输入输出类型都变成了memref但内部标量加法逻辑没有任何变化。如果你不加bufferize-function-boundaries1结果会不太一样函数签名可能保持tensor参数函数内部则插入bufferization.to_memref与bufferization.to_tensor用来在“值世界”和“内存世界”之间做桥接。这种混合模式对中间优化有用但如果你最终要做底层代码生成还是尽快把边界一起 bufferize 掉比较省心。3.3 就地更新、拷贝与 buffer 生命周期Bufferization 最关键的部分不是改类型而是决定“哪些操作可以原地写”。还是用加法举例假设结果可以直接写到%lhs里那就不需要memref.alloc。但是能不能这么干取决于%lhs是否还被后续代码使用。如果%lhs和%rhs传进来以后都再也不会被用到那答案是可以如果其中一个还要在函数外被读取那就不能随便覆盖。这就是别名分析和读后写分析要做的事。one-shot-bufferize会对整个 IR 做一次 SSA 扫描判断哪些张量的 “buffer 版本” 是最后一次使用哪些不是。最后一次使用的尽量就地写不是的就分配新内存并做bufferization.copy。你会发现这和 PyTorch 里那个经典报错有异曲同工之妙在推理模式下对 inference tensor 做 in-place 更新会直接报错提示你 clone 以后再改。PyTorch 之所以不让你原地改是因为那个 tensor 背后可能被其它计算引用改了会破坏整个计算图。MLIR 的 Bufferization 也一样它不是在语法上禁止你写 buffer而是通过分析告诉你这块 buffer 能不能原地写是由它的“使用者”决定的不是由你想不想决定的。一旦思考方式转到“谁是这块 buffer 的最后使用者”很多 pass 的决策逻辑就变得好懂了。4. 实操记录在 mlir-opt 里跑通一次转换4.1 准备环境与输入 IR我假设你本地已经装好了 MLIR 构建环境至少有mlir-opt这个可执行文件版本不用太新近两年的版本都行。没有的话用 LLVM 项目里的 MLIR 构建目标编译一下或者直接用官方 release 包里的 one-shot bufferization 相关工具。新建一个文件叫add.mlir把上面那一节里的 linalg 加法代码放进去。注意我刻意选了linalg.generic而不是纯arith操作因为linalg操作对 tensor 和 memref 都支持bufferize 之后结果比较直观。如果你用纯arith.addf加 tensorbufferization 也会降低成一个scf.for循环多一层循环结构对初学者来说有点干扰。4.2 执行 bufferization pass先直接用最基础命令跑一遍mlir-opt --one-shot-bufferize add.mlir这个命令不会 bufferize 函数边界。你大概率会看到输出里函数签名还是tensor但函数体里多了memref.alloc和linalg.generic的 memref 版本同时会有bufferization.to_tensor把结果转回 tensor 输出。这种结果对学习很有用我会建议你多跑几次对比不同 flag 的输出差异。现在把边界一起打开mlir-opt --one-shot-bufferizebufferize-function-boundaries1 add.mlir结果就是我上一节展示的那样函数整体都变成了 memref 版本。这时候再叠加--canonicalize和--cse很多冗余的memref.cast和重复分配会被清理掉mlir-opt --one-shot-bufferizebufferize-function-boundaries1 --canonicalize --cse add.mlir我个人习惯在写测试的时候用这一条命令作为“标准动作”它能在保留语义的同时把 IR 整理得更容易读也更容易发现 bufferization 阶段引入的冗余操作。4.3 解读生成的 MemRef IR跑完上面的命令重点关注几个位置。第一个是memref.alloc出现在哪。如果函数内部有alloc意味着这个输出的 buffer 是函数自己分配的调用者需要负责释放。另一个是函数参数里有没有多出额外的输出 buffer比如func.func add_two(%lhs: memref4xf32, %rhs: memref4xf32, %out: memref4xf32)这种签名表示调用者提前把输出 buffer 准备好了函数直接往里写。两种方式各有利弊自分配的好处是 API 简洁坏处是没人管释放就会内存泄漏外部传 buffer 的好处是内存管理更清晰适合嵌入在大循环里反复使用同一块缓冲区坏处是函数签名变复杂。MLIR 的 bufferization 默认会采用返回 alloc 的方式如果你想要外部 buffer 风格通常需要自己写bufferization.alloc_tensor的 copy 语义或者在 pass 中加入buffer-deallocation来处理显式释放。实际操作中我建议新手先接受“返回 alloc”的临时结果理解清楚以后再看外部 buffer 的版本否则两个 memory 管理模型混在一起很容易晕。5. 我踩过的坑布局、动态形状与类型转换5.1 总觉得 memref 类型对不上多半是 strides 没包含刚开始手写 MLIR 时我经常遇到类似这样的报错error: memref.load op expected a memref type with same layout as operand最开始我以为是形状不一致查了半天发现两个都是memref4x4xf32。后来才明白MLIR 里两个 memref 即使形状相同只要 strides 不同就被认为是不同类型。一个显式写了strided[4, 1], offset: 0另一个简写成memref4x4xf32两者在某些 pass 的眼里并不等价。解决办法也简单尽量用--canonicalize让 MLIR 把 layout map 规范化。如果需要手动构造 IR就用显式affine_map或者strided写法别靠默认方案蒙混过关。另一个常见操作是memref.reinterpret_cast或者memref.cast但要注意 cast 只能做无损的类型转换如果你想把 strides 从非连续改成连续那必须真的把数据搬一遍cast 是做不到的。5.2 动态形状与符号维度tensor?x?xf32 并不等于 memref?x?xf32这是新手最容易误解的地方。tensor?x?xf32和memref?x?xf32看起来只是关键字不同但在 bufferization 之后memref?x?xf32通常会被展开成带有运行时 sizes 和 strides 的结构。也就是说它的内部需要一个memref.descriptor里面记录每一维的大小和步长。这意味着函数签名里传递的不再是一块裸内存而是一个“描述符”。在实际生成代码时如果要把memref?x?xf32传给一个只接受固定memref4x4xf32的函数没法直接 cast必须重新生成一个带静态形状的 buffer 再来回拷数据。这也是为什么很多性能敏感的代码会把动态维度在早期尽量约束成静态维度。MLIR 的shapedialect 和tensor操作可以帮你做 shape 推断但最终 bufferization 阶段的动态维度处理仍然是整个 pipeline 里最容易出问题的地方。5.3 to_memref / to_tensor 混用的教训bufferization 之后你可能会看到代码里混着bufferization.to_memref和bufferization.to_tensor。这两个操作是用来在两种世界之间“切换”的to_memref把一个 tensor 值变成 memref 视图to_tensor把一个 memref 包回 tensor 值。它们听起来很方便但滥用会导致 pass 分析失效。我印象最深的一次是在一个自定义 pass 里为了省事在函数入口把 tensor 转成 memref然后在函数尾部用to_tensor转回去。结果那个函数的 bufferization 结果特别糟糕很多本应复用的内存被强制拷贝因为编译器看不到to_memref和to_tensor背后对应的是同一块数据别名分析只能悲观地认为它们可能不同。后来我把这个“桥接”尽量控制在小范围内能不用就不用。真正需要跨边界时明确用memref.tensor_store或bufferization.to_tensor并确保在同一个作用域内最好不要让to_tensor跨越复杂控制流。6. 从“能用”到“高效”内存布局优化6.1 strides 与向量化、Tensor Core 的关系为什么内存布局这么重要因为现代硬件的向量指令和矩阵加速单元对内存连续性非常敏感。比如在 x86 上做 AVX 向量化最好一次性加载 256 位连续数据在 NVIDIA GPU 上做 Tensor Core 矩阵乘操作数需要满足特定的 tile 布局通常是 16x16x16 的块而且每次读取的ldmatrix指令也要求数据在共享内存里按特定 stride 摆放。如果在 bufferization 阶段只生成memref...而不关心 strides后续向量化 pass 就会看到 “不友好的布局”要么退化成标量循环要么额外插入vector.transpose或vector.shape_cast来重新组织数据。这些额外操作本质上都是拷贝效率会大打折扣。所以真正追求性能的 MLIR pipeline往往在linalg到vector的转换之前就把memref的 layout map 显式设置成目标硬件友好的格式。比如你可以在 memref 类型里直接写memref16x16xf32, affine_map(d0, d1) - (d0 * 32 d1)表示每行实际上占 32 个元素但只用前 16 个。这种 padding 布局在某些场景里能避免 cache line 伪共享或者 bank conflict。提前在类型系统层面把这些信息定下来后面的vector.contract才能生成更干净的代码。6.2 应该自己写 layout map还是交给 pass这是个很实际的问题。如果只是做原型验证我会直接让 one-shot-bufferize 生成默认布局保持 IR 简单。可一旦要上真正的性能 benchmark就必须认真选布局。MLIR 里常见的做法是先用默认布局把功能跑通然后用--linalg-generalize-named-ops、--linalg-fuse-elementwise-ops这类优化 pass 把操作融合再在--convert-linalg-to-loops之前检查 vectorizer 是否满意当前的 strides。如果发现它生成了大量vector.extract和vector.insert多半就是布局没对齐需要回头改affine_map或者插入memref.expand_shape/memref.collapse_shape来重新解释 buffer 维度。另外很多硬件专用方言比如 GPU dialect 里的 shared memory会强制要求 memref 的内存空间。比如一个memref16x16xf32, 3可能被解释成 shared memory 空间而memref16x16xf32, 1是全局内存空间。要是没刻意设置内存空间布局优化跑得再好访存指令也落不到正确的硬件单元上。做 GPU 相关的 bufferization 时这一步尤其不能偷懒。6.3 给刚入门的朋友的建议如果让我给刚开始碰 MLIR 的人一个建议我会说先不要急着写复杂 layout map也不要一上来就追求 “零拷贝” 的极致 bufferization。先拿最简单的linalg算子跑一遍one-shot-bufferize把输出 IR 每一行都读明白特别是memref.alloc、memref.dealloc、memref.copy这三个操作的位置和数量。等你能解释清楚“为什么这里要分配”“为什么这里要拷贝”之后再去看布局优化就顺理成章了。我自己到现在都保留着一个习惯每次调试 MLIR pipeline第一件事是打开--mlir-print-ir-after-all把每个 pass 之后的 IR 打印出来重点看memref类型有没有突变。很多时候性能问题不是算法问题而是某一个不起眼的 pass 把一个本来连续布局的 memref 悄悄换成了带奇怪 strides 的版本。这种问题如果不在 IR 阶段抓住落到汇编层面就非常难查了。Bufferization 和内存布局这块确实要花不少时间琢磨。它不像 dialect 语法那样背一背就能上手而是需要你在 “抽象值” 和 “物理存储” 之间反复切换视角。但一旦想通了后面看 MLIR 里所有的代码生成 pass都会有豁然开朗的感觉。下一篇我打算接着写一写 bufferization 之后最常见的linalg到vector的转换过程希望能继续帮到正在这条路上摸索的朋友。