本篇是学习 MLIR 中的 Transform Dialect 的第一篇，主要讲解 MLIR 为何及如何保持计算结构。内容来自 MLIR 官方教程 《Chapter 0: A Primer on “Structured” Linalg Operations》。
相关链接： LLVM Project ，MLIR 官方文档，MLIR 官网教程， 《Chapter 0: A Primer on “Structured” Linalg Operations》。
作为初学者，错误在所难免，还望不吝赐教。
本章节主要内容是在阐述 ： 为什么编译器不直接把高级计算（如矩阵乘）转成低级循环，而是要先用 “结构化” 操作（如 linalg.generic）来表示它？
答案是：为了保留 “结构” 信息，给编译器更大的优化空间和灵活性。低级循环（如 scf.for）是一盘散沙，而结构化操作是一块块清晰的乐高积木，方便我们进行各种重组和优化。
这里讲的结构化，指的是：它们拥有清晰、固定的语义，这与 “非结构化” 的、包含任意控制流的通用循环（如 scf.for）形成对比。具体体现在：
操作有明确名称：如 linalg.matmul（矩阵乘）、linalg.conv（卷积）、linalg.generic（通用形式）。
数据访问模式可预测：操作明确声明哪些张量是输入、哪些是输出，以及如何访问它们（如读取、写入、归约）。
便于分析和转换：因为结构规整，编译器能更容易地进行循环重排、分块（tiling）、向量化等优化，而无需进行复杂的数据流分析。
接下来教程从简单到复杂，一步步阐述了保持结构的思想是如何体现的。
考虑 MLIR 中一个简单的标量算术加法操作，该操作在大多数支持浮点运算的架构上直接映射为机器指令：
%2 = arith.addf %0, %1 : f32 |
这个结构非常稳固，直接对应一条硬件指令。它本身就是一个最小、最清晰的 “结构单元”。编译器无需任何额外信息，就能完美地处理它。
现在将此操作拓展，均匀地应用于一维向量的各个元素：
%2 = arith.addf %0, %1 : vector<8xf32> |
当前操作的核心语义（元素对元素的加法）被完整地保留下来了。它没有被拆解成 8 个独立的标量加法，而是作为一个整体操作存在。编译器看到了这个清晰的 “元素级并行” 结构后，就有了选择的自由：如果目标硬件有支持 vector<8xf32 > 加法的 SIMD 指令，编译器可以直接映射过去，获得高性能；如果硬件不支持，编译器也可以选择把它 “降级” 成循环或标量指令；重点在于，这个决策权（何时、如何实现）被保留在了编译流程中，而不是在编写代码时就固定死。
目前仅有少数现代指令集提供了对二维或更高维向量的指令支持。然而，在 MLIR 中，可以透明地将统一的逐元素操作扩展到任意秩的向量。
%2 = arith.addf %0, %1 : vector<8x4xf32> | |
%5 = arith.addf %3, %4 : vector<2x2x2x2x2x2x2xf32> |
如您所见，MLIR 对向量的算术操作保留了元素逐个统一应用的结构。编译器可以利用这一结构，例如生成目标硬件上可用的低秩操作，或在支持融合指令的情况下将乘法和加法合并（当连续一百次乘法后接一百次加法时，这种合并会变得复杂）。
MLIR 的 Vector 方言显式的表达了在一个 vector 中发生的规约操作。
%1 = vector.reduction <add>, %0 : vector<8xf32> into f32 |
这个操作清晰地表达了 “归约” 的意图。编译器可以灵活选择实现方式：是使用专用硬件指令，还是展开成一个循环，甚至是用相邻加法指令组合实现。
在没有特别支持的情况下，该操作可以转换为循环：
%c0 = arith.constant 0 : index | |
%c1 = arith.constant 1 : index | |
%c8 = arith.constant 8 : index | |
%init = arith.constant 0.0 : f32 | |
%result = scf.for %i = %c0 to %c8 step %c1 iter_args(%partial = %init) -> (f32) { | |
%element = vector.extract %0[%i] : f32 into vector<8xf32> | |
%updated = arith.addf %partial, %element : f32 | |
scf.yield %updated : f32 | |
} |
Contraction（收缩）是一种 “带归约的乘法”。它描述了一个非常通用的计算模式：从两个数组中取出元素相乘，然后将乘积结果累加起来。
MLIR 专门提供了收缩操作的表示：
%init = arith.constant 0.0 : f32 | |
%ones = arith.constant dense<1.0> : vector<8xf32> | |
%result = vector.contract { | |
indexing_maps = [affine_map<(i) -> (i)>, | |
affine_map<(i) -> (i)>, | |
affine_map<(i) -> ()>], | |
iterator_types = ["reduction"] | |
} %0, %ones, %init : vector<8xf32>, vector<8xf32> into f32 |
上述收缩表示等同于下方伪代码：
for i in 0 to 8: | |
init += p0[i] * ones[i] |
不仅一维数据，教程又给了个经典的二维例子：二维矩阵乘法 (Matrix Multiplication)
%result = vector.contract { | |
indexing_maps = [affine_map<(i, j, k) -> (i, k)>, | |
affine_map<(i, j, k) -> (k, j)>, | |
affine_map<(i, j, k) -> (i, j)>], | |
iterator_types = ["parallel", "parallel", "reduction"] | |
} %lhs, %rhs, %init: vector<8x10xf32>, vector<10x16xf32> into vector<8x16xf32> |
iterator_types ：明确告诉编译器， i 和 j 是并行维度（结果矩阵的每个元素可以独立计算），而 k 是归约维度（需要累加）
indexing_maps ：精确描述了数据如何被索引。比如 (i, j, k) -> (i, k) 表示在访问左矩阵时，行索引由 i 决定，列索引由 k 决定。
它的伪代码等同于：
for i in 0 to 8: | |
for j in 0 to 16: | |
for k in 0 to 10: | |
init[i, j] += lhs[i, k] * rhs[k, j] |
明确的并行性：编译器一眼就能看出 i 和 j 维度是完全并行的，这为后续的自动并行化（如多线程、GPU 映射）提供了最直接的信息。
可识别的模式：这个结构清晰地表明这是一个矩阵乘法。编译器可以利用这个信息，去调用经过高度优化的、手写的矩阵乘库（如 BLAS 库），或者使用特定的硬件加速器指令。这比让编译器从一堆循环中去 “猜” 这是矩阵乘要可靠和高效得多。
前面讲了 vector 方言在虚拟寄存器上表示的缩进操作，其实再内存上也有类似缩进的操作：
linalg.generic { | |
indexing_maps = [affine_map<(i, j, k) -> (i, k)>, | |
affine_map<(i, j, k) -> (k, j)>, | |
affine_map<(i, j, k) -> (i, j)>], | |
iterator_types = ["parallel", "parallel", "reduction"] | |
} ins(%lhs, %rhs : memref<8x10xf32>, memref<10x16xf32>) | |
outs(%init : memref<8x16xf32>) { | |
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32): | |
%0 = arith.mulf %lhs_one, %rhs_one : f32 | |
%1 = arith.addf %init_one, %0 : f32 | |
linalg.yield %1 : f32 | |
} |
之前 (vector.contract)：操作在虚拟寄存器中的向量上，数据是只读的，操作产生新向量，像函数式编程。
现在 (linalg.generic)：操作在内存中的缓冲区（memref）上，数据是 “就地” 读写的，更接近底层硬件的工作方式。
ins (inputs)：是只读的输入缓冲区，就像仓库 A 和 B，里面的原料（lhs, rhs）在整个计算过程中不会被修改。
outs (outputs)：是既读取又更新的输出缓冲区，就像成品仓库 init。计算开始时，它里面可能有初始值（比如全 0），计算过程中会不断从中读取旧值、写入新值。
{ | |
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32): | |
%0 = arith.mulf %lhs_one, %rhs_one : f32 | |
%1 = arith.addf %init_one, %0 : f32 | |
linalg.yield %1 : f32 | |
} |
以上是 linalg.generic 最核心、也最 “结构化” 的部分。它提供了一个代码区域（region）
linalg.generic 的 “布局图” 规定了它要遍历所有工作点，在每个点上，它会从 ins 和 outs 指定的缓冲区中，根据 indexing_maps 取出对应的一个元素，作为参数传给 ^bb0 。
参数顺序：前两个 (%lhs_one, %rhs_one) 对应 ins 列表，最后一个 (%init_one) 对应 outs 列表。
操作体拿到这三个元素（两个输入，一个初始 / 旧输出）后，按照手册操作：先乘后加。
linalg.yield %1 ：完成操作后，必须通过 linalg.yield 把结果 %1 交出来。这个结果会被写回到 outs 缓冲区中对应的位置（在这里就是 %init 的 (i, j) 位置）。
从这里能够看到 linalg.generic 的强大之处：
1. 它声明了 “是什么”：通过 indexing_maps 和 iterator_types，它精确无误地声明了这是一个带归约的、三维的、数据访问模式明确的计算。这比散落的循环嵌套更具信息量。
2. 它分离了 “控制流” 和 “计算”：循环的遍历方式（控制流）由操作属性控制，而每个点上具体的乘加计算（计算）被封装在 region 里。
3. 它为优化打开了大门：因为结构如此清晰，编译器可以轻松地对它进行分块（Tiling）、融合（Fusion）、并行化等高级转换，而不需要费力的分析。
这一小节讲的是 linalg.generic 用单个操作实现 “循环融合（Loop Fusion）。
传统循环中，如果做两个循环，一般是先算一个，再算另一个，那么就有两个独立的循环：
for (i = 0; i < N; i++) { | |
temp[i] = input[i] * 2.0; | |
} | |
for (i = 0; i < N; i++) { | |
output[i] = temp[i] + 1.0; | |
} |
循环融合就是把这两个循环融合成一个：
for (i = 0; i < N; i++) { | |
temp = input[i] * 2.0; | |
output[i] = temp + 1.0; | |
} |
好处是：消除临时缓冲区，减少内存访问，提高缓存利用率。传统编译器做循环融合很困难，因为要分析两个循环的迭代空间是否兼容；分析数据依赖；处理迭代空间不匹配的情况。这些分析复杂且容易出错。
linalg.generic 能够天然地实现融合，因为 linalg.generic 的 region 里可以写任意多条操作，所以 “融合” 就变成了 “在同一个 region 里多写几行代码” 这么简单。
比如下面地 Relu 操作：
linalg.generic { | |
indexing_maps [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], | |
iterator_types = ["parallel"] | |
} ins(%in : memref<?xf32>) outs(%out : memref<?xf32>) { | |
^bb0(%in_one : f32, %out_one : f32): | |
%c0 = arith.constant 0.0 : f32 | |
%0 = arith.cmpf ogt %in_one, %c0 : f32 | |
%1 = arith.select %0, %in_one, %c0 : f32 | |
linalg.yield %1 : f32 | |
} |
如果没有融合，ReLU 可能需要两个操作：一个操作做比较，产生一个布尔结果（需要一个临时缓冲区）。
另一个操作根据布尔结果做选择。
举个例子：
传统循环：就像给厨师两张分开的菜谱 ——“先切菜，切完放盘子里；再炒菜，从盘子里拿切好的菜”。中间那个 “盘子” 就是临时缓冲区。
linalg.generic 的融合：就像给厨师一张合并的菜谱 ——“切完直接炒，不用放盘子”。厨师也可以选择还是分两步做，但菜谱本身给了他一气呵成的可能。
该小节将 linalg.generic 从内存缓冲区（memref）提升到了张量（tensor）。
Buffer 版本：outs (% init) 就像一个 “草稿纸”，操作会直接在上面涂改，最后这张草稿纸本身就是结果。
Tensor 版本：outs (% init) 就像一份 “参考模板”，操作不会修改它，而是照着它生成一份全新的文档作为结果返回。
这就是 值语义（Value Semantics） 和 引用语义（Reference Semantics） 的区别。
%result = linalg.generic { | |
indexing_maps = [affine_map<(i, j, k) -> (i, k)>, | |
affine_map<(i, j, k) -> (k, j)>, | |
affine_map<(i, j, k) -> (i, j)>], | |
iterator_types = ["parallel", "parallel", "reduction"] | |
} ins(%lhs, %rhs : tensor<8x10xf32>,tensor<10x16xf32>) | |
outs(%init :tensor<8x16xf32>) { | |
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32): | |
%0 = arith.mulf %lhs_one, %rhs_one : f32 | |
%1 = arith.addf %init_one, %0 : f32 | |
linalg.yield %1 : f32 | |
} -> tensor<8x16xf32> |
在 Tensor 上的操作和上述在 Buffer 上的操作很相似，除了操作数类型外，主要区别在于该操作现在会生成一个新结果，而不是更新输出缓冲区。
分块就是把一个大计算，切成许多小块来计算，让每一块的数据都能塞进缓存（Cache），从而避免频繁访问慢速的主存。
linalg.generic 分块非常简单。
1. 迭代空间是隐式的： linalg.generic 没有显式地写出 for i... for j... for k... 。它的迭代空间完全由操作数的形状（如 8x10、10x16）推导出来。
2. 分块 = 对数据切片做同样的操作：既然迭代空间是隐式的，那么 “分块” 就不需要去修改循环边界。你只需要从原数据中切出一小块（Slice），然后对这一小块应用完全相同的 linalg.generic 操作即可。
3. 无需依赖分析：因为 linalg.generic 明确规定 “对每个元素元组的执行顺序未指定”，所以各个分块之间天然就是可并行、可乱序的。编译器不需要做复杂的依赖分析，就能安全地重排分块顺序。
%0 = scf.forall (%i, %j) in (4, 2) shared_outs(%shared = %init) -> (tensor<8x16xf32>) { | |
%3 = affine.apply affine_map<(d0) -> (d0 * 2)>(%i) | |
%4 = affine.apply affine_map<(d0) -> (d0 * 8)>(%j) | |
%lhs_slice = tensor.extract_slice %lhs[%3, 0] [2, 10] [1, 1] | |
: tensor<8x10xf32> to tensor<2x10xf32> | |
%rhs_slice = tensor.extract_slice %rhs[0, %4] [10, 8] [1, 1] | |
: tensor<10x16xf32> to tensor<10x8xf32> | |
%result_slice = tensor.extract_slice %shared[%3, %4] [2, 8] [1, 1] | |
: tensor<8x16xf32> to tensor<2x8xf32> | |
%partial = linalg.generic { | |
indexing_maps = [affine_map<(i, j, k) -> (i, k)>, | |
affine_map<(i, j, k) -> (k, j)>, | |
affine_map<(i, j, k) -> (i, j)>], | |
iterator_types = ["parallel", "parallel", "reduction"] | |
} ins(%lhs_slice, %rhs_slice : tensor<2x10xf32>, tensor<10x8xf32>) | |
outs(%result_slice : tensor<2x8xf32>) -> tensor<2x8xf32> { | |
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32): | |
%0 = arith.mulf %lhs_one, %rhs_one : f32 | |
%1 = arith.addf %init_one, %0 : f32 | |
linalg.yield %1 : f32 | |
} : tensor<2x8xf32> | |
scf.forall.in_parallel { | |
tensor.parallel_insert_slice %partial into %shared[%3, %4] [2, 8] [1, 1] | |
: tensor<2x8xf32> into tensor<8x16xf32> | |
} | |
} |
之前： linalg.generic 的循环是隐式的，你看不到 for 循环，只能从 indexing_maps 和 iterator_types 推导出来。
分块后：为了控制各个分块的执行顺序（比如先算哪一块、后算哪一块），编译器生成了显式的循环（这里是 scf.forall ）
分块的过程，同时也是把隐式循环 “物化”（Materialize）成显式循环的过程。
生产者和消费者，其实就是数据从 A 流向 B。比如矩阵乘 linalg.matmul，它产生了一个结果张量，它是生产者，之后一个逐元素平方操作，读取矩阵乘的结果，它是消费者。
传统循环融合（Loop Fusion）要求两个循环的迭代空间完全匹配才能融合。但结构化操作允许迭代空间不匹配也能融合。
%0 = scf.forall (%i, %j) in (4, 2) | |
shared_outs(%shared = %init) -> (tensor<8x16xf32>, tensor<8x16xf32>) { | |
%1 = affine.apply affine_map<(d0) -> (d0 * 2)>(%i) | |
%2 = affine.apply affine_map<(d0) -> (d0 * 8)>(%j) | |
%lhs_slice = tensor.extract_slice %lhs[%1, 0] [2, 10] [1, 1] | |
: tensor<8x10xf32> to tensor<2x10xf32> | |
%rhs_slice = tensor.extract_slice %rhs[0, %2] [10, 8] [1, 1] | |
: tensor<10x16xf32> to tensor<10x8xf32> | |
%result_slice = tensor.extract_slice %result[%1, %2] [2, 8] [1, 1] | |
: tensor<8x16xf32> to tensor<2x8xf32> | |
%partial = linalg.generic { | |
indexing_maps = [affine_map<(i, j, k) -> (i, k)>, | |
affine_map<(i, j, k) -> (k, j)>, | |
affine_map<(i, j, k) -> (i, j)>], | |
iterator_types = ["parallel", "parallel", "reduction"] | |
} ins(%lhs_slice, %rhs_slice : tensor<2x10xf32>, tensor<10x8xf32>) | |
outs(%result_slice : tensor<2x8xf32>) { | |
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32): | |
%5 = arith.mulf %lhs_one, %rhs_one : f32 | |
%6 = arith.addf %init_one, %5 : f32 | |
linalg.yield %6 : f32 | |
} -> tensor<2x8xf32> | |
%shared_slice = tensor.extract_slice %shared[%1, %2] [2, 8] [1, 1] | |
: tensor<8x16xf32> to tensor<2x8xf32> | |
%elemwise = linalg.generic { | |
indexing_maps = [affine_map<(i, j) -> (i, j)>, | |
affine_map<(i, j) -> (i, j)>], | |
iterator_types = ["parallel", "parallel"] | |
} ins(%partial : tensor<2x8xf32>) | |
outs(%shared_slice : tensor<2x8xf32>) { | |
^bb0(%in: f32, %out: f32): | |
%5 = arith.mulf %in, %in : f32 | |
linalg.yield %5 : f32 | |
} -> tensor<2x8xf32> | |
scf.forall.in_parallel { | |
tensor.parallel_insert_slice %elemwise into %shared[%1, %2] [2, 8] [1, 1] | |
: tensor<2x8xf32> into tensor<8x16xf32> | |
} | |
} |
教程给的例子：Producer：矩阵乘，3D 迭代空间 (i, j, k)。Consumer：逐元素平方，2D 迭代空间 (i, j)
教程之所以能够 Tile 切块还能融合：是因为逆推 Shape。
分块是从 Consumer 开始的。我们先把 Consumer 的迭代空间切成 2x8 的小块，要计算出这个 2x8 的消费者小块，生产者需要提供哪一部分数据，逆推出生产者 Producer 的迭代空间。
前面我们一直在手写 linalg.generic，需要自己填 indexing_maps、iterator_types、region 里的乘加逻辑，比较麻烦，但对于常见计算（矩阵乘、点积、卷积等），这些 “模板” 是固定的，所以 Linalg 直接提供了一批命名操作（如 linalg.matmul、linalg.dot、linalg.conv），把那些固定内容预先填好了。它们在语义上完全等价于对应的 linalg.generic 。
之前的通用版矩阵乘：
%matmul = linalg.generic { | |
indexing_maps = [affine_map<(i, j, k) -> (i, k)>, | |
affine_map<(i, j, k) -> (k, j)>, | |
affine_map<(i, j, k) -> (i, j)>], | |
iterator_types = ["parallel", "parallel", "reduction"] | |
} ins(%lhs, %rhs : tensor<8x10xf32>, tensor<10x16xf32>) | |
outs(%init : tensor<8x16xf32>) { | |
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32): | |
%0 = arith.mulf %lhs_one, %rhs_one : f32 | |
%1 = arith.addf %init_one, %0 : f32 | |
linalg.yield %1 : f32 | |
} -> tensor<8x16xf32> |
等价于当前的：
%matmul = linalg.matmul ins(%lhs, %rhs: tensor<8x10xf32>, tensor<10x16xf32>) | |
outs(%init: tensor<8x10xf32xf32>) -> tensor<8x16xf32> |
一个在应对环境方面有异常困难的人在挣扎着，尘土飞扬。 | |
我曾使用过这样一个形象：一条被鱼钩钩住的鱼。 | |
在其他不了解这种情况的鱼看来，它旋转的样子一定很奇特；但它激起的水花并不是它的痛苦，而是摆脱痛苦的努力。 | |
每个渔民都知道，这种努力很可能会成功。 | |
—— 卡尔.A.梅宁格(Karl A .Menninger),1945 |
本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目 或随便一个项目下提出 issue，并指明哪一篇博客，看到一定及时回复！
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。