阅读和理解 MLIR(多层中间表示)是任何在编译器层面优化机器学习工作负载的人所需的基本能力。MLIR 的结构及其方言和运算,为表示不同抽象层级的计算提供了一个框架。对常见 MLIR 模式的分析提供了对其结构和功能的实践性理解。我们假设您拥有可以显示或生成 MLIR 的工具。许多机器学习编译器框架(例如启用了 XLA 的 TensorFlow、IREE,或直接使用 LLVM/MLIR 的项目)可以在不同阶段导出其 MLIR 表示形式。示例 1:逐元素相加我们从一个基本的张量运算开始:将两个浮点数的二维张量相加。在像 linalg 这样的方言中(常用于张量上的线性代数运算),这可能看起来像这样:#map0 = affine_map<(d0, d1) -> (d0, d1)> module { func.func @add_tensors(%arg0: tensor<128x256xf32>, %arg1: tensor<128x256xf32>) -> tensor<128x256xf32> { // 为结果分配内存 (简化表示) %init_result = linalg.init_tensor [128, 256] : tensor<128x256xf32> // 使用 linalg.generic 执行逐元素相加 %result = linalg.generic { indexing_maps = [#map0, #map0, #map0], // 输入 A, 输入 B, 输出 iterator_types = ["parallel", "parallel"] // 两个维度都是并行的 } ins(%arg0, %arg1 : tensor<128x256xf32>, tensor<128x256xf32>) outs(%init_result : tensor<128x256xf32>) { ^bb0(%in0: f32, %in1: f32, %out_init: f32): // 内部计算的基本块 %sum = arith.addf %in0, %in1 : f32 linalg.yield %sum : f32 // 为此元素返回结果 } -> tensor<128x256xf32> return %result : tensor<128x256xf32> } }分析要点:方言: 这里的主要运算是 linalg.generic。linalg 方言旨在处理张量上的结构化运算,通常作为降级为循环或特定硬件内部函数之前的中间层抽象。内部计算使用 arith 方言中的 arith.addf 运算。类型: 请注意显式张量类型:tensor<128x256xf32>。这立刻告知我们形状(128x256)和元素类型(32 位浮点数)。此静态形状信息对优化很重要。运算: func.func 定义一个函数。linalg.init_tensor 分配空间(实际上,内存管理更为复杂)。linalg.generic 是一个强大的运算,定义对张量元素进行的计算。其主体(^bb0)指定对每个元素执行的标量计算。arith.addf 执行浮点加法。linalg.yield 返回为此元素计算的值。结构: 运算以 SSA(静态单赋值)值作为操作数(例如 %arg0、%arg1),并产生 SSA 值作为结果(例如 %result)。仿射映射(#map0)定义循环迭代器如何映射到张量索引,这对于理解数据访问模式很重要。iterator_types 指定循环的性质(此处,“并行”表示这些维度上没有跨迭代依赖)。示例 2:矩阵乘法矩阵乘法是许多机器学习模型的运算基础。一个高级表示,可能在领域特定方言中或使用 linalg.matmul,可能看起来像:module { func.func @matmul(%A: tensor<32x64xf32>, %B: tensor<64x128xf32>) -> tensor<32x128xf32> { %C = linalg.matmul ins(%A, %B : tensor<32x64xf32>, tensor<64x128xf32>) outs(/* init tensor omitted */) -> tensor<32x128xf32> // 假设 %C 在矩阵乘法之前已正确初始化 return %C : tensor<32x128xf32> } }如果我们查看此代码在降级为 affine 和 scf(结构化控制流)等方言之后,为 CPU 执行做准备,我们可能会看到这样的代码(高度简化):#map_A = affine_map<(d0, d1, d2) -> (d0, d2)> // 访问 A[i, k] #map_B = affine_map<(d0, d1, d2) -> (d2, d1)> // 访问 B[k, j] #map_C = affine_map<(d0, d1, d2) -> (d0, d1)> // 访问 C[i, j] module { func.func @matmul_lowered(%argA: memref<32x64xf32>, %argB: memref<64x128xf32>, %argC: memref<32x128xf32>) { // 结果维度 (i, j) 的外层循环 scf.for %i = 0 to 32 step 1 { scf.for %j = 0 to 128 step 1 { // 内层归约循环 (k) %init_acc = arith.constant 0.0 : f32 %acc = scf.for %k = 0 to 64 step 1 iter_args(%iter_acc = %init_acc) -> (f32) { // 从输入张量加载元素 %a_val = affine.load %argA[%i, %k] : memref<32x64xf32> // 简化访问 %b_val = affine.load %argB[%k, %j] : memref<64x128xf32> // 简化访问 // 计算乘积并累加 %prod = arith.mulf %a_val, %b_val : f32 %new_acc = arith.addf %iter_acc, %prod : f32 scf.yield %new_acc : f32 } // 存储最终累加值 affine.store %acc, %argC[%i, %j] : memref<32x128xf32> // 简化访问 } } return } }分析要点(降级示例):方言转换: linalg.matmul 运算消失了。我们现在看到用于循环的 scf.for、用于内存访问的 affine.load/affine.store(使用 memref 类型而非 tensor),以及用于核心计算的 arith 运算。这代表一个 降级 步骤。抽象层级: 这更接近传统命令式代码。循环结构、内存访问和算术运算都是显式的。此层级适用于平铺、向量化以及最终生成机器指令等转换。类型: 请注意使用 memref(内存引用)而非 tensor。memref 通常意味着已分配的、具有特定内存布局的内存,而 tensor 则更抽象。仿射映射(尽管在此处已简化)对于分析内存访问模式以进行缓存优化变得非常重要。分析指南在检查不同工具或不同编译阶段生成的 MLIR 代码时,请考虑以下问题:存在哪些方言? 这会告知您当前的抽象层级(例如,mhlo 用于高级图运算,linalg 用于结构化张量运算,vector 用于 SIMD,affine/scf 用于循环/内存,gpu 用于 GPU 特有内容,llvm 用于最终降级)。主要运算是什么? 寻找表示计算、控制流或类型转换的运算。理解它们的操作数和结果。张量形状和类型如何表示? 形状是静态的还是动态的?元素类型是什么?这严重影响优化选择。数据如何流动? 追踪从生产者到消费者的 SSA 值。能否发现潜在的高级优化? 寻找可以融合的运算序列(linalg.generic 通常是融合的结果)。如果已降级,结构与原始运算有何关联? 能否看出 matmul 如何变为嵌套循环?这可能如何映射到硬件(例如,并行循环到 GPU 线程/块)?是否存在定义运算行为的属性? 寻找运算上的属性,这些属性指定了步长、填充、数据布局或精度细节。这种动手分析非常重要。通过阅读和理解 MLIR,您将获得对机器学习编译器如何组织计算并应用优化的认识。它使您能够理解不同编译策略的影响,并找出可以提高性能的方面,摆脱机器学习框架的黑盒视图。