张量中端：Tensor 与 Linalg 方言

从这一篇开始进入 MLIR 在 AI 编译中的核心应用。tensor 和 linalg 两个方言构成了 MLIR 的”张量中端”——它们是 PyTorch、TensorFlow、JAX 等框架的计算图降阶为 MLIR 后首先接触的抽象层，也是 tiling、fusion、vectorization 等关键优化的主战场。

一、Tensor 方言：不可变的数学张量

tensor 方言表示不可变的多维数组。它的语义来自函数式编程——没有副作用、没有别名、没有内存位置的概念。一个 tensor 值就是一个纯数学对象。

1.1 核心操作

Op	语义
`tensor.empty`	创建一个未初始化的 tensor（用于 `outs` 参数）
`tensor.empty_like`	创建与已有 tensor 形状相同的未初始化 tensor
`tensor.dim`	获取 tensor 某个维度的运行时大小
`tensor.extract`	从 tensor 中读取单个元素
`tensor.insert`	将元素写入 tensor——返回一个新的 tensor（不修改原 tensor）
`tensor.extract_slice`	从 tensor 中提取一个子视图（slice）
`tensor.insert_slice`	将子视图插入 tensor——返回新 tensor
`tensor.pad`	对 tensor 进行填充（padding）
`tensor.collapse_shape`	将多个维度折叠为一个（reshape 变体）
`tensor.expand_shape`	将一个维度展开为多个（reshape 变体）
`tensor.reshape`	改变形状（静态参数）
`tensor.generate`	通过一个 Region 按元素生成 tensor 内容

1.2 设计意图：纯函数式、无副作用

// tensor.insert 不修改原 tensor——它返回一个新 tensor
%new_tensor = tensor.insert %value into %old_tensor[%i, %j]
             : tensor<256x256xf32>

// %old_tensor 在 insert 后仍然有效且不变
// 后续使用 %old_tensor 仍然得到原始值

这种不可变性使得数据流分析可以精确做依赖判断——没有通过内存的隐式副作用。两个 linalg.matmul 如果它们的 outs tensor 不同，就确定是独立的（可以并行或重排）。如果 outs 相同（RAW 依赖），就必须顺序执行。

1.3 Tensor 的生命周期终点

Tensor 是不可变的抽象——但在真实硬件上，计算必然操作内存。解决方案是 bufferization——将 tensor 方言的 IR 转换为 memref 方言：

tensor (不可变、无别名) ──→ (bufferization) ──→ memref (可变、有内存地址)

bufferization 在 tensor 的每个定义点插入 memref.alloc，将 tensor.insert 转为 memref.store，tensor.extract 转为 memref.load。关键决策是：哪些 tensor 需要独立的 memref（即哪些操作需要自己的输出缓冲区，哪些可以复用输入缓冲区）。

MLIR 的 One-Shot Bufferization 实现了这个分析——它通过别名分析确定 buffer 的生命周期，并自动解决”输入 tensor 的 buffer 能否复用为输出 tensor 的 buffer”（in-place bufferization）的问题。

二、Linalg 方言：结构化操作代数

linalg（Linear Algebra）是 MLIR 中表达结构化数值计算的方言。它的核心抽象是 structured operations on tensors/memrefs。

2.1 结构化操作的三个层级

flowchart TB
  N["Named Ops<br/>matmul / conv_2d / fill …"]
  G["linalg.generic<br/>逐元素 + 归约 + 收缩"]
  M["indexing_maps（affine_map）<br/>操作数到迭代空间的映射"]
  N --> G
  G --> M

2.2 Named Ops：语义明确的常见操作

// 矩阵乘法：C[m,n] += A[m,k] * B[k,n]
%result = linalg.matmul
    ins(%A, %B : tensor<MxKxf32>, tensor<KxNxf32>)
    outs(%C : tensor<MxNxf32>) -> tensor<MxNxf32>

// 2D 卷积：O[n,oh,ow,oc] += I[n,ih,iw,ic] * F[fh,fw,ic,oc]
%result = linalg.conv_2d
    ins(%input, %filter : tensor<NxIHxIWxICxf32>, tensor<FHxFWxICxOCxf32>)
    outs(%init : tensor<NxOHxOWxOCxf32>) -> tensor<NxOHxOWxOCxf32>

// 填充（fill）：将 tensor 的所有元素设置为常量
%filled = linalg.fill ins(%value : f32)
    outs(%tensor : tensor<256x256xf32>) -> tensor<256x256xf32>

Named Ops 的价值：后续 Pass（tiling、vectorization、lowering to loops）可以根据 Op 名称选择最优策略。linalg.matmul 知道中间维度 k 是归约维度——tiling 时可以沿着 m 和 n 分块，同时在 k 维度累加。

2.3 linalg.generic：统一表示

linalg.generic 是 Linalg 的最通用表示——用三样东西完全描述一个结构化计算：

Indexing Maps：每个操作数到迭代空间的映射。
Iterator Types：每个维度的遍历方式（parallel、reduction）。
Body Region：循环体的逐元素计算。

以矩阵乘法为例：

#map_A = affine_map<(m, n, k) -> (m, k)>   // A[m,k] ——不依赖 n
#map_B = affine_map<(m, n, k) -> (k, n)>   // B[k,n] ——不依赖 m
#map_C = affine_map<(m, n, k) -> (m, n)>   // C[m,n] ——不依赖 k

%result = linalg.generic {
    indexing_maps = [#map_A, #map_B, #map_C],
    iterator_types = ["parallel", "parallel", "reduction"]
}
ins(%A, %B : tensor<MxKxf32>, tensor<KxNxf32>)
outs(%C : tensor<MxNxf32>) {
^bb0(%a: f32, %b: f32, %c: f32):
  %mul = arith.mulf %a, %b : f32
  %add = arith.addf %c, %mul : f32
  linalg.yield %add : f32
} -> tensor<MxNxf32>

拆解：

iterator_types = ["parallel", "parallel", "reduction"]：前两个维度 m 和 n 是并行的（无数据依赖），第三个维度 k 是归约的（需要累加）。
indexing_maps 描述了每个操作数如何映射到迭代变元 (m, n, k)：
- A[m,k] 不依赖 n——当沿着 n 遍历时，A 的访问索引不变。
- C[m,n] 不依赖 k——归约维度对 C 来说不是索引而是累加器。
body Region 定义了内层循环的计算——接收三个 f32 标量，做乘法和加法。

linalg.matmul 本质上就是上述 linalg.generic 的语法糖。

三、linalg.generic 的 iterator_types 三分类

这是理解 Linalg 表示能力的关键。每种分类对应一种数据依赖模式：

iterator_type	含义
`parallel`	迭代之间无依赖，可并行或重排
`reduction`	归约维度，需要初始值，迭代间有累加依赖
`window`	滑动窗口（卷积等），迭代间有重叠数据访问

window 在较新 Linalg 版本中用于卷积类 Op；具体可用 iterator 集合以你所用 MLIR 版本的 LinalgOps.td 为准。

分类	示例	说明
全 `parallel`	`linalg.add`（逐元素加法）	所有维度无依赖，完全并行
`parallel` + `reduction`	`linalg.matmul`	归约维度有累加依赖
`parallel` + `reduction` + `window`	`linalg.conv_2d`	窗口维度有重叠数据访问

常见的具体示例

逐元素操作（纯 parallel）——每个输出元素只依赖对应位置的输入元素：

// C[i,j] = A[i,j] + B[i,j]
linalg.generic {
  indexing_maps = [affine_map<(i, j) -> (i, j)>,   // A 和 B 与迭代变量一一映射
                   affine_map<(i, j) -> (i, j)>,
                   affine_map<(i, j) -> (i, j)>],
  iterator_types = ["parallel", "parallel"]
}
ins(%A, %B : tensor<256x512xf32>, tensor<256x512xf32>)
outs(%C : tensor<256x512xf32>) {
^bb0(%a: f32, %b: f32, %c: f32):
  %add = arith.addf %a, %b : f32
  linalg.yield %add : f32
}

归约（reduction 维度上做累加）：

// C[i] += A[i,j]  —— 沿 j 维度归约
linalg.generic {
  indexing_maps = [affine_map<(i, j) -> (i, j)>,   // A 与迭代变量一一映射
                   affine_map<(i, j) -> (i)>],       // C 不依赖 j（归约维度）
  iterator_types = ["parallel", "reduction"]
}

四、从 Linalg 到 Loops 的降阶

Linalg 结构化操作最终需要降阶为实际的循环。-convert-linalg-to-loops Pass 将 linalg.generic 展开为嵌套的 scf.for 循环：

// 输入：linalg.matmul
linalg.matmul ins(%A, %B : ...) outs(%C : ...) -> ...

// 降阶后（概念层面）
scf.for %m = %c0 to %M step %c1 {
  scf.for %n = %c0 to %N step %c1 {
    %acc = memref.load %C[%m, %n] : memref<MxNxf32>
    scf.for %k = %c0 to %K step %c1 {
      %a = memref.load %A[%m, %k] : memref<MxKxf32>
      %b = memref.load %B[%k, %n] : memref<KxNxf32>
      %mul = arith.mulf %a, %b : f32
      %acc_new = arith.addf %acc, %mul : f32
      memref.store %acc_new, %C[%m, %n] : memref<MxNxf32>
    }
  }
}

关键点：循环降阶通常发生在 bufferization 之后——此时 tensor 已转为 memref，循环中的 load/store 操作真实内存地址。部分教学管线（如第 03 章）为便于观察会把 -linalg-bufferize 放在 -convert-linalg-to-loops 之前；生产管线（IREE、StableHLO 路径）多在 linalg 层完成 tiling/fusion 后再做 bufferization。Pass 顺序因项目而异，以目标编译器为准。

五、Linalg 的优化策略：Tiling、Fusion、Vectorization

Linalg 方言的设计初衷就是为这三项优化提供清晰的 IR 接口：

Tiling（分块）：将大 linalg.matmul 拆成多个更小的 linalg.matmul（或 linalg.generic），每个 tile 负责一部分计算。tiling 由 MLIR 的 Transform 方言或 Linalg 的 tiling Pass 自动完成：

// tiled：256x256 matmul → 四个 128x128 matmul
scf.for %i = %c0 to %c256 step %c128 {
  scf.for %j = %c0 to %c256 step %c128 {
    %tile = linalg.matmul ins(%A_slice, %B_slice : ...)
                          outs(%C_slice : ...) -> ...
  }
}

Fusion（融合）：将两个连续的 linalg.generic（如 matmul → relu）合并为一个——减少中间 tensor 的内存分配和数据搬运。

Vectorization（向量化）：将最内层 linalg.generic 转为 vector 方言的 SIMD 操作——利用 CPU 的 AVX/NEON/SVE 或 GPU 的 SIMT lane。

六、Linalg 命名操作总览

类别	Op	对应 DL 操作
逐元素	`linalg.add`, `linalg.sub`, `linalg.mul`, `linalg.div`	张量的逐元素运算
逐元素（带广播）	`linalg.broadcast`	`torch.broadcast_to`
归约	`linalg.reduce`	`torch.sum(dim=d)`
收缩	`linalg.matmul`	`torch.matmul`
收缩	`linalg.batch_matmul`	`torch.bmm`
收缩	`linalg.conv_2d`, `linalg.conv_2d_nhwc_hwcf`	`torch.nn.functional.conv2d`
收缩	`linalg.depthwise_conv_2d`	`torch.nn.functional.conv2d(groups=C)`
填充	`linalg.fill`	`torch.full`
复制	`linalg.copy`	`torch.clone`
通用	`linalg.generic`	所有自定义逐元素/归约/收缩操作

七、本篇后续

Tensor 和 Linalg 是 AI 编译的”中间语言”——Tensor 提供函数式语义，Linalg 提供结构化操作表示。下一章进循环层——Affine 和 SCF 方言，看结构化操作如何展开为可调度、可并行的循环 IR。

参考资料

官方文档（A 级）

MLIR Tensor Dialect — https://mlir.llvm.org/docs/Dialects/Tensor/
MLIR Linalg Dialect — https://mlir.llvm.org/docs/Dialects/Linalg/
MLIR Bufferization — https://mlir.llvm.org/docs/Bufferization/

论文（A 级）

Lattner, C. et al. MLIR: A Compiler Infrastructure for the End of Moore’s Law. arXiv:2002.11054, 2020.
MLIR Linalg Design Document — MLIR 社区设计文档

源码（A 级）

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

同主题继续阅读

把当前热点继续串成多页阅读，而不是停在单篇消费。

2026-06-09 · compiler / architecture

【编译器与 MLIR】张量中端：Tensor 与 Linalg 方言

文章导航

目录