# 前言

本篇是学习 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 对向量的算术操作保留了元素逐个统一应用的结构。编译器可以利用这一结构，例如生成目标硬件上可用的低秩操作，或在支持融合指令的情况下将乘法和加法合并（当连续一百次乘法后接一百次加法时，这种合并会变得复杂）。

# 归约 (Reduction)

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)

Contraction（收缩）是一种 “带归约的乘法”。它描述了一个非常通用的计算模式：从两个数组中取出元素相乘，然后将乘积结果累加起来。

MLIR 专门提供了收缩操作的表示：

	// 1. 初始化累加器为 0.0
	%init = arith.constant 0.0 : f32
	// 2. 创建一个全为 1.0 的向量，作为乘法的中性元素（这里是为了演示，实际点积会用两个数据向量）
	%ones = arith.constant dense<1.0> : vector<8xf32>
	// 3. 核心收缩操作
	%result = vector.contract {
	indexing_maps = [affine_map<(i) -> (i)>, // 第一个输入 %0 的索引方式
	affine_map<(i) -> (i)>, // 第二个输入 % ones 的索引方式
	affine_map<(i) -> ()>], // 输出 % init 的索引方式（空括号表示标量）
	iterator_types = ["reduction"] // 迭代维度 i 是一个归约维度
	} %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)>, // 左矩阵 A [i, k]
	affine_map<(i, j, k) -> (k, j)>, // 右矩阵 B [k, j]
	affine_map<(i, j, k) -> (i, j)>], // 结果矩阵 C [i, j]
	iterator_types = ["parallel", "parallel", "reduction"] //i 并行，j 并行，k 归约
	} %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 库），或者使用特定的硬件加速器指令。这比让编译器从一堆循环中去 “猜” 这是矩阵乘要可靠和高效得多。

# 内存上的通用操作 (Generic Operation on Memory)

前面讲了 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）、并行化等高级转换，而不需要费力的分析。

# 循环融合 (“Loop” Fusion)

这一小节讲的是 linalg.generic 用单个操作实现 “循环融合（Loop Fusion）。
传统循环中，如果做两个循环，一般是先算一个，再算另一个，那么就有两个独立的循环：

	// 循环 1：计算中间结果
	for (i = 0; i < N; i++) {
	temp[i] = input[i] * 2.0;
	}
	// 循环 2：使用中间结果
	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 // 比较：in > 0 ?
	%1 = arith.select %0, %in_one, %c0 : f32 // 选择：如果是，取 in；否则取 0
	linalg.yield %1 : f32
	}

如果没有融合，ReLU 可能需要两个操作：一个操作做比较，产生一个布尔结果（需要一个临时缓冲区）。

另一个操作根据布尔结果做选择。
举个例子：
传统循环：就像给厨师两张分开的菜谱 ——“先切菜，切完放盘子里；再炒菜，从盘子里拿切好的菜”。中间那个 “盘子” 就是临时缓冲区。

linalg.generic 的融合：就像给厨师一张合并的菜谱 ——“切完直接炒，不用放盘子”。厨师也可以选择还是分两步做，但菜谱本身给了他一气呵成的可能。

# 张量的通用操作 (Generic Operation on Tensors)

该小节将 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 上的操作很相似，除了操作数类型外，主要区别在于该操作现在会生成一个新结果，而不是更新输出缓冲区。

# 分块与循环物化 (Tiling and Loop Materialization)

分块就是把一个大计算，切成许多小块来计算，让每一块的数据都能塞进缓存（Cache），从而避免频繁访问慢速的主存。
linalg.generic 分块非常简单。
1. 迭代空间是隐式的： linalg.generic 没有显式地写出 for i... for j... for k... 。它的迭代空间完全由操作数的形状（如 8x10、10x16）推导出来。

2. 分块 = 对数据切片做同样的操作：既然迭代空间是隐式的，那么 “分块” 就不需要去修改循环边界。你只需要从原数据中切出一小块（Slice），然后对这一小块应用完全相同的 linalg.generic 操作即可。

3. 无需依赖分析：因为 linalg.generic 明确规定 “对每个元素元组的执行顺序未指定”，所以各个分块之间天然就是可并行、可乱序的。编译器不需要做复杂的依赖分析，就能安全地重排分块顺序。

	// A special "multi-for" loop that supports tensor-insertion semantics
	// as opposed to implicit updates. The resulting 8x16 tensor will be produced
	// by this loop.
	// The trip count of iterators is computed dividing the original tensor size,
	// 8x16, by the tile size, 2x8, to obtain 4x2.
	// When tensor sizes are dynamic, the trip count computation is emitted as IR
	// and is being computed at runtime.
	%0 = scf.forall (%i, %j) in (4, 2) shared_outs(%shared = %init) -> (tensor<8x16xf32>) { // 外层循环 (4,2) 循环的迭代次数

	// Scale the loop induction variables by the tile sizes.
	%3 = affine.apply affine_map<(d0) -> (d0 * 2)>(%i) // 计算切片偏移
	%4 = affine.apply affine_map<(d0) -> (d0 * 8)>(%j) // 循环变量 % i 和 % j 是分块的索引 % i*2：第 i 块在行方向的起始位置

	// Take slices of inputs and outputs. Only the "i" and "j" dimensions are sliced.
	%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>

	// This is exactly the same operation as before, but now operating on smaller
	// slices of data.
	%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>

	// Terminator for the loop with tensor-insertion semantics. Inserts a slice
	// into a larger tensor, potentially in parallel.
	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）成显式循环的过程。

# 生产者 - 消费者融合和重物化 (Producer/Consumer Fusion and Rematerialization)

生产者和消费者，其实就是数据从 A 流向 B。比如矩阵乘 linalg.matmul，它产生了一个结果张量，它是生产者，之后一个逐元素平方操作，读取矩阵乘的结果，它是消费者。
传统循环融合（Loop Fusion）要求两个循环的迭代空间完全匹配才能融合。但结构化操作允许迭代空间不匹配也能融合。

	// Same loop as before.
	%0 = scf.forall (%i, %j) in (4, 2) // 和前述同样的分块 (2*8)
	shared_outs(%shared = %init) -> (tensor<8x16xf32>, tensor<8x16xf32>) { // 教程这里标记了两个输出，实际这段代码只有 % elemwise 一个输出，可能教程代码做了简化之类的
	// Scale the loop induction variables by the tile sizes.
	%1 = affine.apply affine_map<(d0) -> (d0 * 2)>(%i)
	%2 = affine.apply affine_map<(d0) -> (d0 * 8)>(%j)

	// Take slices of inputs and outputs. Only the "i" and "j" dimensions are sliced.
	%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>

	// This is exactly the same matmul slice as before. It replaces the slice
	// extraction for the generic operation below.
	%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>

	// Take the slice of the final result. Note that we don't need to take
	// the slice of the operand because the matmul operation above computes
	// it in-place.
	%shared_slice = tensor.extract_slice %shared[%1, %2] [2, 8] [1, 1]
	: tensor<8x16xf32> to tensor<2x8xf32>

	// The elementwise operation that we tiled.
	%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>

	// Terminator for the loop with tensor-insertion semantics. Inserts a slice
	// into a larger tensor, potentially in parallel.
	scf.forall.in_parallel { // 最终输出 % elemwise (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 的迭代空间。

# Shorthand “Named” Forms of Linalg Ops

前面我们一直在手写 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，并指明哪一篇博客，看到一定及时回复！