MLIR学习笔记
- MLIR
- [一、MLIR 概述](#一、MLIR 概述)
-
- [1\.1 什么是 MLIR](#1.1 什么是 MLIR)
- [1\.2 为什么需要 MLIR](#1.2 为什么需要 MLIR)
- [1\.3 MLIR 的设计哲学](#1.3 MLIR 的设计哲学)
- 二、核心概念体系
-
- [2\.1 Operation(操作)](#2.1 Operation(操作))
-
- [Operation 的组成](#Operation 的组成)
- [2\.2 Dialect(方言)](#2.2 Dialect(方言))
-
- [Dialect 的作用](#Dialect 的作用)
- [常用内置 Dialect](#常用内置 Dialect)
- [2\.3 Type(类型)](#2.3 Type(类型))
- [2\.4 Attribute(属性)](#2.4 Attribute(属性))
-
- [常用 Attribute 类型](#常用 Attribute 类型)
- [2\.5 Region 与 Block](#2.5 Region 与 Block)
- [2\.6 Trait(特性)](#2.6 Trait(特性))
-
- [常用 Trait](#常用 Trait)
- [2\.7 Interface(接口)](#2.7 Interface(接口))
- [三、ODS 声明式定义](#三、ODS 声明式定义)
-
- [3\.1 什么是 ODS](#3.1 什么是 ODS)
-
- [ODS 的优势](#ODS 的优势)
- [3\.2 Dialect 定义](#3.2 Dialect 定义)
- [3\.3 如何定义一个 Dialect(完整示例)](#3.3 如何定义一个 Dialect(完整示例))
-
- [步骤一:TableGen ODS 声明](#步骤一:TableGen ODS 声明)
- [步骤二:定义 Operation(ODS)](#步骤二:定义 Operation(ODS))
- [步骤三:生成 C\+\+ 代码](#步骤三:生成 C++ 代码)
- [步骤四:C\+\+ 头文件实现](#步骤四:C++ 头文件实现)
- [步骤五:C\+\+ 源文件实现](#步骤五:C++ 源文件实现)
- [步骤六:注册 Dialect 到 MLIRContext](#步骤六:注册 Dialect 到 MLIRContext)
- [步骤七:CMake 构建配置](#步骤七:CMake 构建配置)
- 完整文件结构
- [3\.4 Operation 定义](#3.4 Operation 定义)
- [3\.5 Type 定义](#3.5 Type 定义)
- [四、Pattern Rewriting 模式重写框架](#四、Pattern Rewriting 模式重写框架)
-
- [4\.1 概述](#4.1 概述)
- [4\.2 RewritePattern 类](#4.2 RewritePattern 类)
- [4\.3 PatternRewriter 常用 API](#4.3 PatternRewriter 常用 API)
- [4\.4 Greedy Pattern Rewrite Driver](#4.4 Greedy Pattern Rewrite Driver)
- [4\.5 DRR 声明式重写规则](#4.5 DRR 声明式重写规则)
- [4\.6 Canonicalization(规范化)](#4.6 Canonicalization(规范化))
- [五、Pass 管理体系](#五、Pass 管理体系)
-
- [5\.1 Pass 分类](#5.1 Pass 分类)
- [5\.2 Pass 定义](#5.2 Pass 定义)
- [5\.3 PassManager 管线编排](#5.3 PassManager 管线编排)
- [5\.4 Pass 调度机制](#5.4 Pass 调度机制)
- [六、Dialect Conversion 与 Lowering](#六、Dialect Conversion 与 Lowering)
-
- [6\.1 Lowering 概念](#6.1 Lowering 概念)
- [6\.2 Dialect Conversion 框架](#6.2 Dialect Conversion 框架)
- [6\.3 ConversionPattern](#6.3 ConversionPattern)
- [6\.4 ConversionTarget 合法性定义](#6.4 ConversionTarget 合法性定义)
- [6\.5 TypeConverter 类型转换](#6.5 TypeConverter 类型转换)
- [6\.6 应用转换](#6.6 应用转换)
- [6\.7 常见 Lowering 路径](#6.7 常见 Lowering 路径)
- [七、Toy Tutorial 学习路径](#七、Toy Tutorial 学习路径)
-
- [7\.1 章节概览](#7.1 章节概览)
- [7\.2 学习建议](#7.2 学习建议)
- 八、开发环境与工具
-
- [8\.1 编译 MLIR](#8.1 编译 MLIR)
- [8\.2 常用工具](#8.2 常用工具)
- [8\.3 mlir\-opt 常用选项](#8.3 mlir-opt 常用选项)
- 九、进阶主题
-
- [9\.1 数据分析与优化](#9.1 数据分析与优化)
- [9\.2 多线程与并行](#9.2 多线程与并行)
- [9\.3 TOSA 与 AI 编译](#9.3 TOSA 与 AI 编译)
- 十、学习资源
MLIR
部分内容由豆包生成
一、MLIR 概述
1.1 什么是 MLIR
MLIR(Multi-Level Intermediate Representation,多级中间表示)是 LLVM 项目下的一个编译器基础设施框架,由 Chris Lattner 主导开发。它的核心目标是解决传统编译器在面对异构计算、领域特定语言(DSL)和深度学习时的扩展性问题。
核心理念:通过可扩展的 Dialect(方言)机制,在统一的 IR 框架下支持多层次的抽象表示,从高层的语言语义到底层的硬件指令,实现渐进式 lowering(降级)。
1.2 为什么需要 MLIR
-
传统编译器的局限:LLVM IR 只有单一抽象层级,难以表达高层语义信息
-
异构计算需求:CPU、GPU、TPU、NPU 等多种硬件需要不同的抽象层次
-
DSL 开发成本高:每种新语言都要从零构建完整编译器栈
-
AI 编译器爆发:深度学习框架需要从计算图到硬件指令的多级转换
1.3 MLIR 的设计哲学
| 设计原则 | 说明 |
|---|---|
| 渐进式降级 | 从高层 Dialect 逐步 lowering 到底层 Dialect,每一步保留语义 |
| Dialect 可扩展 | 每个领域可以定义自己的 Dialect,共享统一的基础设施 |
| SSA 形式 | 基于静态单赋值形式,便于分析和优化 |
| 声明式定义 | 通过 TableGen(ODS)声明式定义 Operation、Type、Attribute |
二、核心概念体系
2.1 Operation(操作)
Operation 是 MLIR 的核心抽象单元,类似于 LLVM 中的 Instruction,但功能更强大。
Operation 的组成
-
名称(Name) :格式为
dialect.op-name,如arith.addi -
操作数(Operands):输入值,SSA 值的引用
-
结果(Results):输出值,可多个
-
属性(Attributes):编译期常量元数据
-
区域(Regions):嵌套的子控制流区域
-
后继块(Successors):控制流跳转目标
mlir
// arith.addi 操作:两个输入,一个输出
%result = arith.addi %a, %b : i32
// 带属性的操作
%c42 = arith.constant 42 : i32
// 带区域的操作(func.func 包含一个 Region)
func.func @add(%a: i32, %b: i32) -> i32 {
%sum = arith.addi %a, %b : i32
return %sum : i32
}
2.2 Dialect(方言)
Dialect 是 MLIR 的扩展机制,是 Operation、Type、Attribute 的逻辑分组,通过唯一命名空间隔离。
Dialect 的作用
-
提供命名空间隔离,避免命名冲突
-
组织相关的 Operation、Type、Attribute
-
支持 Dialect 级别的常量折叠、验证等钩子
-
实现模块化的编译器设计
常用内置 Dialect
| Dialect | 用途 | 典型 Operation |
|---|---|---|
builtin |
内建基础类型和模块 | module, unrealized_conversion_cast |
func |
函数定义与调用 | func.func, func.call, func.return |
arith |
算术运算 | arith.addi, arith.mulf, arith.constant |
affine |
仿射循环与内存访问 | affine.for, affine.load, affine.store |
memref |
多维内存引用 | memref.load, memref.store, memref.alloc |
tensor |
张量类型操作 | tensor.extract, tensor.insert |
vector |
SIMD 向量化操作 | vector.fma, vector.load, vector.store |
llvm |
LLVM IR 兼容层 | llvm.add, llvm.br, llvm.call |
scf |
结构化控制流 | scf.for, scf.if, scf.while |
tosa |
TOSA 神经网络标准 | tosa.add, tosa.conv2d |
2.3 Type(类型)
MLIR 的类型系统是可扩展的,每个 Dialect 可以定义自己的类型。
内置基础类型
-
整数类型 :
i1,i8,i32,i64等 -
浮点类型 :
f32,f64,f16,bf16 -
索引类型 :
index(平台相关指针大小) -
函数类型 :
(i32, f32) -> i64 -
张量类型 :
tensor<2x3xf32> -
MemRef 类型 :
memref<10x20xi32> -
向量类型 :
vector<4xf32>
2.4 Attribute(属性)
Attribute 是附加在 Operation 上的编译期常量元数据,不同于运行时的 Operand。
常用 Attribute 类型
-
整数属性 :
IntegerAttr- 如常量值 42 -
浮点属性 :
FloatAttr -
字符串属性 :
StringAttr -
数组属性 :
ArrayAttr- 属性的数组 -
字典属性 :
DictionaryAttr- 键值对集合 -
类型属性 :
TypeAttr- 包装类型 -
符号引用属性 :
SymbolRefAttr- 引用函数等符号
2.5 Region 与 Block
Region(区域)
Region 是 Operation 内部包含的 CFG(控制流图)容器,一个 Operation 可以包含多个 Region。
-
SSACFG Region:标准 SSA 控制流图,包含多个 BasicBlock
-
Graph Region:图区域(如 ModuleOp),不要求 SSA 支配关系
Block(基本块)
Block 是 Region 内的基本块,包含一系列 Operation 和终止符操作(terminator)。
-
Block 可以有参数(Block Arguments),用于接收前驱块传来的值
-
每个 Block 必须以 terminator op 结尾(如
func.return,cf.br)
mlir
func.func @foo(%arg0: i32) -> i32 {
// ^entry 是第一个 Block,%arg0 是 Block Argument
%0 = arith.addi %arg0, %arg0 : i32
cf.br ^bb1(%0 : i32)
^bb1(%x: i32): // 带参数的 Block
func.return %x : i32
}
2.6 Trait(特性)
Trait 是附加在 Operation 上的语义标签,用于描述 Operation 的属性,无需额外实现代码。
常用 Trait
| Trait | 语义 |
|---|---|
OneResult |
操作只有一个结果 |
NoSideEffect |
无副作用,可被 DCE 消除 |
Commutative |
交换律,操作数顺序不影响结果 |
SameOperandsAndResultType |
所有操作数和结果类型相同 |
Terminator |
是基本块的终止符 |
SingleBlock |
Region 内只有一个 Block |
2.7 Interface(接口)
Interface 定义了一组操作需要实现的方法,使得泛化 Pass 可以跨 Dialect 工作。
-
OpInterface :操作接口,如
ShapeInferenceOpInterface -
DialectInterface:方言接口
-
TypeInterface:类型接口
-
AttributeInterface:属性接口
Trait vs Interface:Trait 是静态标签,无虚函数开销;Interface 是动态多态,需要实现虚函数。Trait 用于简单属性标记,Interface 用于复杂的泛化交互。
三、ODS 声明式定义
3.1 什么是 ODS
ODS(Operation Definition Specification)是 MLIR 提供的声明式 Operation 定义框架,基于 TableGen 语法。
ODS 的优势
-
声明式定义,自动生成 C++ 代码
-
统一的验证、解析、打印逻辑
-
支持 Trait、Interface 的自动接入
-
便于维护和扩展
3.2 Dialect 定义
3.3 如何定义一个 Dialect(完整示例)
定义一个自定义 Dialect 需要以下几个步骤:TableGen ODS 声明、C++ 类实现、Dialect 注册。下面以一个名为 mydialect 的自定义方言为例,展示完整流程。
步骤一:TableGen ODS 声明
首先创建 MyDialect.td 文件,用 ODS 声明 Dialect 的基本信息。
tablegen
#ifndef MY_DIALECT_TD
#define MY_DIALECT_TD
include "mlir/IR/DialectBase.td"
// 定义 Dialect
def MyDialect : Dialect {
// Dialect 命名空间前缀(操作名格式:mydialect.xxx)
let name = "mydialect";
// C++ 命名空间
let cppNamespace = "::mlir::mydialect";
// Dialect 描述文档
let summary = "My custom MLIR dialect";
let description = [{
This is an example custom dialect for demonstration purposes.
It contains simple arithmetic operations and custom types.
}];
// 使用的属性前缀(默认与 name 相同)
let cppClassName = "MyDialect";
// 注册该 Dialect 包含的 Operation
let operations = [
My_AddOp,
My_ConstantOp,
My_ReturnOp
];
// 注册该 Dialect 包含的 Type
let types = [
My_TensorType
];
// 注册该 Dialect 包含的 Attribute
let attributes = [
My_ShapeAttr
];
// 额外的 C++ 类声明(会插入到头文件中)
let extraClassDeclaration = [{
/// 注册 Dialect 时的初始化钩子
void initialize();
/// 自定义常量折叠钩子
Optional<SmallVector<OpFoldResult, 2>>
tryFold(Operation *op, ArrayRef<Attribute> operands);
/// 自定义操作解析
Operation *parseOperation(OpAsmParser &parser,
OperationState &state) const;
/// 自定义操作打印
void printOperation(Operation *op,
OpAsmPrinter &p) const;
}];
// Dialect 级别的 Trait
let hasConstantMaterializer = 1;
let hasOperationVerifier = 1;
}
#endif // MY_DIALECT_TD
步骤二:定义 Operation(ODS)
在同一个 .td 文件或单独的文件中定义该 Dialect 的 Operation。
tablegen
#ifndef MY_OPS_TD
#define MY_OPS_TD
include "mlir/IR/OpBase.td"
include "MyDialect.td"
// 加法操作
def My_AddOp : My_Op<"add">,
Arguments<(ins AnyType:$lhs, AnyType:$rhs)>,
Results<(outs AnyType:$result)>,
Traits<[
NoSideEffect,
SameOperandsAndResultType,
Commutative
]> {
let summary = "Addition operation";
let description = "Element-wise addition of two values";
// 自定义汇编格式
let assemblyFormat = [{
$lhs `,` $rhs attr-dict `:` type($result)
}];
// 启用 verify 方法
let hasVerifier = 1;
// 启用 folder 方法
let hasFolder = 1;
}
// 常量操作
def My_ConstantOp : My_Op<"constant">,
Arguments<(ins AnyAttr:$value)>,
Results<(outs AnyType:$output)>,
Traits<[
NoSideEffect,
ConstantLike
]> {
let summary = "Constant operation";
let assemblyFormat = "attr-dict $value `:` type($output)";
}
// 返回操作(终止符)
def My_ReturnOp : My_Op<"return">,
Arguments<(ins Variadic<AnyType>:$inputs)>,
Results<(outs)>,
Traits<[
Terminator,
NoSideEffect,
ReturnLike
]> {
let summary = "Return operation";
let assemblyFormat = "($inputs `:` type($inputs))? attr-dict";
}
#endif // MY_OPS_TD
步骤三:生成 C++ 代码
使用 mlir-tblgen 工具从 TableGen 生成 C++ 头文件和实现文件。
bash
# 生成 Dialect 声明头文件
mlir-tblgen -gen-dialect-decls \
-I /path/to/mlir/include \
MyDialect.td \
-o MyDialect.h.inc
# 生成 Dialect 实现文件
mlir-tblgen -gen-dialect-defs \
-I /path/to/mlir/include \
MyDialect.td \
-o MyDialect.cpp.inc
# 生成 Operation 声明
mlir-tblgen -gen-op-decls \
-I /path/to/mlir/include \
MyOps.td \
-o MyOps.h.inc
# 生成 Operation 实现
mlir-tblgen -gen-op-defs \
-I /path/to/mlir/include \
MyOps.td \
-o MyOps.cpp.inc
步骤四:C++ 头文件实现
创建 MyDialect.h 头文件,引入生成的代码并声明自定义方法。
cpp
#ifndef MY_DIALECT_H
#define MY_DIALECT_H
#include "mlir/IR/Dialect.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/DialectImplementation.h"
// 引入生成的 Dialect 类声明
#define GET_TYPEDEF_CLASSES
#include "MyTypes.h.inc"
#define GET_ATTRDEF_CLASSES
#include "MyAttributes.h.inc"
namespace mlir {
namespace mydialect {
/// 自定义 Dialect 类
class MyDialect : public Dialect {
public:
explicit MyDialect(MLIRContext *context);
~MyDialect() override = default;
/// Dialect 唯一名称
static constexpr StringLiteral getDialectName() {
return StringLiteral("mydialect");
}
/// 初始化方法
void initialize();
/// 常量折叠
Optional<SmallVector<OpFoldResult, 2>>
tryFold(Operation *op, ArrayRef<Attribute> operands);
/// 解析自定义类型
Type parseType(DialectAsmParser &parser) const;
/// 打印自定义类型
void printType(Type type, DialectAsmPrinter &printer) const;
/// 解析自定义属性
Attribute parseAttribute(DialectAsmParser &parser, Type type) const;
/// 打印自定义属性
void printAttribute(Attribute attr, DialectAsmPrinter &printer) const;
};
} // namespace mydialect
} // namespace mlir
// 引入生成的 Operation 类声明
#define GET_OP_CLASSES
#include "MyOps.h.inc"
#endif // MY_DIALECT_H
步骤五:C++ 源文件实现
创建 MyDialect.cpp 实现文件,实现 Dialect 的初始化和自定义方法。
cpp
#include "MyDialect.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Support/LogicalResult.h"
namespace mlir {
namespace mydialect {
// 构造函数:注册 Dialect
MyDialect::MyDialect(MLIRContext *context)
: Dialect(getDialectName(), context, TypeID::get<MyDialect>()) {
initialize();
}
// 初始化:注册所有 Operation、Type、Attribute
void MyDialect::initialize() {
// 注册 Operation(由 ODS 自动生成)
addOperations<
#define GET_OP_LIST
#include "MyOps.cpp.inc"
>();
// 注册 Type
addTypes<
#define GET_TYPEDEF_LIST
#include "MyTypes.cpp.inc"
>();
// 注册 Attribute
addAttributes<
#define GET_ATTRDEF_LIST
#include "MyAttributes.cpp.inc"
>();
}
// 常量折叠实现
Optional<SmallVector<OpFoldResult, 2>>
MyDialect::tryFold(Operation *op,
ArrayRef<Attribute> operands) {
// 示例:常量折叠 add(const, const)
if (auto addOp = dyn_cast<AddOp>(op)) {
if (operands[0] && operands[1]) {
// 执行常量折叠逻辑...
return SmallVector<OpFoldResult, 2>{foldedResult};
}
}
return None;
}
// 解析自定义类型
Type MyDialect::parseType(DialectAsmParser &parser) const {
StringRef keyword;
if (parser.parseKeyword(&keyword)) return Type();
if (keyword == "tensor") {
// 解析 mydialect.tensor<...>
return parseTensorType(parser);
}
parser.emitError(parser.getNameLoc())
<< "unknown mydialect type: " << keyword;
return nullptr;
}
// 打印自定义类型
void MyDialect::printType(Type type,
DialectAsmPrinter &printer) const {
// 分发到具体类型的打印方法
if (auto tensorTy = dyn_cast<TensorType>(type)) {
printTensorType(tensorTy, printer);
return;
}
llvm_unreachable("unknown mydialect type");
}
} // namespace mydialect
} // namespace mlir
// 引入生成的 Operation 实现
#define GET_OP_CLASSES
#include "MyOps.cpp.inc"
步骤六:注册 Dialect 到 MLIRContext
使用时需要将自定义 Dialect 注册到 MLIR 上下文中。
cpp
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/DialectRegistry.h"
#include "MyDialect.h"
int main() {
mlir::DialectRegistry registry;
// 注册自定义 Dialect
registry.insert<mlir::mydialect::MyDialect>();
// 同时注册依赖的内置 Dialect
registry.insert<mlir::func::FuncDialect>();
registry.insert<mlir::arith::ArithDialect>();
// 创建 MLIRContext
mlir::MLIRContext context(registry);
// 加载所有已注册的 Dialect
context.loadAllAvailableDialects();
// ... 后续使用 context 构建和处理 IR
return 0;
}
步骤七:CMake 构建配置
在 CMakeLists.txt 中配置 TableGen 代码生成和库编译。
cmake
add_mlir_dialect_library(MLIRMyDialect
MyDialect.cpp
MyOps.cpp
MyTypes.cpp
MyAttributes.cpp
# 生成的 TableGen 文件
DEPENDS
MyDialectTableGen
MyOpsTableGen
MyTypesTableGen
MyAttributesTableGen
LINK_COMPONENTS
Support
LINK_LIBS PUBLIC
MLIRIR
MLIRSupport
MLIRDialect
)
# TableGen 生成规则
mlir_tablegen(MyDialect.h.inc -gen-dialect-decls)
mlir_tablegen(MyDialect.cpp.inc -gen-dialect-defs)
mlir_tablegen(MyOps.h.inc -gen-op-decls)
mlir_tablegen(MyOps.cpp.inc -gen-op-defs)
mlir_tablegen(MyTypes.h.inc -gen-typedef-decls)
mlir_tablegen(MyTypes.cpp.inc -gen-typedef-defs)
add_public_tablegen_target(MyDialectTableGen)
完整文件结构
| 文件 | 作用 |
|---|---|
MyDialect.td |
Dialect 的 TableGen 定义 |
MyOps.td |
Operation 的 TableGen 定义 |
MyTypes.td |
Type 的 TableGen 定义 |
MyAttributes.td |
Attribute 的 TableGen 定义 |
MyDialect.h |
Dialect C++ 头文件 |
MyDialect.cpp |
Dialect C++ 实现文件 |
*.h.inc / *.cpp.inc |
mlir-tblgen 自动生成的代码 |
CMakeLists.txt |
构建配置 |
最佳实践:优先使用 ODS 声明式定义,尽量少写手动 C++ 代码。只有需要自定义逻辑(如特殊验证、常量折叠、自定义解析打印)时才手动实现。
tablegen
def Toy_Dialect : Dialect {
let name = "toy";
let cppNamespace = "::mlir::toy";
// 描述
let description = [{
This dialect implements the Toy language operations.
}];
// 注册的 Type
let types = [
Toy_StructType
];
// 注册的 Attribute
let attributes = [
Toy_StructAttr
];
}
3.4 Operation 定义
tablegen
def Toy_AddOp : Toy_Op<"add">,
Arguments<(ins AnyType:$lhs, AnyType:$rhs)>,
Results<(outs AnyType:$output)>,
Traits<[
NoSideEffect,
SameOperandsAndResultType,
Commutative
]> {
let summary = "Addition operation";
let description = [{
Element-wise addition of two tensor values.
}];
// 自定义汇编格式
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($output)";
// 验证方法
let hasVerifier = 1;
}
3.5 Type 定义
tablegen
def Toy_StructType : TypeDef<Toy_Dialect, "StructType"> {
let mnemonic = "struct";
let summary = "Struct type";
// Type 的参数
let parameters = (ins
ArrayRefParameter<"Type">:$elementTypes
);
// 打印/解析
let printer = [{
$_printer << "<";
llvm::interleaveComma(getElementTypes(), $_printer);
$_printer << ">";
}];
let parser = [{
if ($_parser.parseLess()) return Type();
SmallVector<Type, 2> elementTypes;
do {
Type elementType;
if ($_parser.parseType(elementType)) return Type();
elementTypes.push_back(elementType);
} while (succeeded($_parser.parseOptionalComma()));
if ($_parser.parseGreater()) return Type();
return get($_parser.getContext(), elementTypes);
}];
}
四、Pattern Rewriting 模式重写框架
4.1 概述
Pattern Rewriting 是 MLIR 中实现 IR 转换的核心机制,基于 DAG-to-DAG 重写思想。
两种定义方式
-
命令式(C++):继承 RewritePattern,手动实现 matchAndRewrite
-
声明式(DRR):Declarative Rewrite Rules,TableGen 定义
4.2 RewritePattern 类
cpp
struct SimplifyRedundantTranspose : public OpRewritePattern<TransposeOp> {
using OpRewritePattern<TransposeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(
TransposeOp op,
PatternRewriter &rewriter) const override {
// 匹配:输入也是一个 transpose 操作
TransposeOp transposeInputOp =
op.getOperand().getDefiningOp<TransposeOp>();
if (!transposeInputOp) return failure();
// 重写:transpose(transpose(x)) = x
rewriter.replaceOp(op, transposeInputOp.getOperand());
return success();
}
};
4.3 PatternRewriter 常用 API
-
replaceOp(op, newValues):用新值替换操作的所有结果 -
replaceOpWithNewOp<OpTy>(op, args...):创建新操作并替换 -
eraseOp(op):删除操作 -
create<OpTy>(loc, args...):创建新操作 -
notifyRootUpdated(op):通知根操作被修改
4.4 Greedy Pattern Rewrite Driver
贪婪模式重写驱动,反复应用模式直到无法继续匹配。
-
按 benefit 值排序,优先应用收益高的模式
-
支持局部重应用(locally apply)
-
自动处理 IR 修改通知
cpp
RewritePatternSet patterns(&context);
patterns.add<Pattern1, Pattern2, Pattern3>(&context);
// 设置最大迭代次数
GreedyRewriteConfig config;
config.maxIterations = 10;
if (failed(applyPatternsAndFoldGreedily(
module, std::move(patterns), config))) {
signalPassFailure();
}
4.5 DRR 声明式重写规则
tablegen
// transpose(transpose(x)) -> x
def : Pat<
// 匹配模式
(Toy_TransposeOp (Toy_TransposeOp $input)),
// 替换结果
(Toy_TransposeOp $input),
// 可选的约束
[]>;
// constant folding: add(const x, const y) -> const(x+y)
def : Pat<
(Toy_AddOp (Toy_ConstantOp F64Attr:$lhs),
(Toy_ConstantOp F64Attr:$rhs)),
(Toy_ConstantOp (AddF64Attr $lhs, $rhs))>;
4.6 Canonicalization(规范化)
每个 Operation 可以注册自己的规范化模式,在 canonicalize Pass 中统一应用。
cpp
void AddOp::getCanonicalizationPatterns(
RewritePatternSet &patterns,
MLIRContext *context) {
patterns.add<SimplifyRedundantTranspose>(context);
}
五、Pass 管理体系
5.1 Pass 分类
| Pass 类型 | 模板参数 | 作用范围 |
|---|---|---|
| OperationPass | OperationPass<OpTy> |
特定类型的 Operation |
| FunctionPass | OperationPass<func::FuncOp> |
函数级别 |
| ModulePass | OperationPass<ModuleOp> |
模块级别 |
| RegionPass | RegionPass |
每个 Region |
5.2 Pass 定义
cpp
struct ShapeInferencePass
: public PassWrapper<ShapeInferencePass,
OperationPass<ModuleOp>> {
// 声明依赖的 Dialect
void getDependentDialects(
DialectRegistry ®istry) const override {
registry.insert<ArithDialect>();
}
// Pass 主逻辑
void runOnOperation() override {
ModuleOp module = getOperation();
// ... 执行形状推断 ...
}
// Pass 选项
Option<std::string> someOption{
this, "some-option",
llvm::cl::desc("An example option"),
llvm::cl::init("default")
};
};
5.3 PassManager 管线编排
cpp
PassManager pm(&context);
// 添加模块级 Pass
pm.addPass(createCanonicalizerPass());
pm.addPass(createCSEPass());
// 嵌套 Pass:在 FuncOp 上运行
OpPassManager &funcPM = pm.nest<func::FuncOp>();
funcPM.addPass(createShapeInferencePass());
funcPM.addPass(createCanonicalizerPass());
// 执行管线
if (failed(pm.run(module))) {
llvm::errs() << "Pass failed\n";
return 1;
}
5.4 Pass 调度机制
-
Pass 排序:自动按依赖关系调度
-
验证:每个 Pass 运行后可执行验证器
-
崩溃恢复:支持 Pass 失败时的 IR dump
-
Pass 计时:内置 PassTiming 统计
-
多线程:函数级 Pass 可并行执行
六、Dialect Conversion 与 Lowering
6.1 Lowering 概念
Lowering(降级)是将高层 Dialect 的操作转换为低层 Dialect 操作的过程,语义保持不变,仅表示形式变化。
典型 Lowering 路径:Toy Dialect → Affine Dialect → SCF + MemRef → LLVM Dialect → LLVM IR → 机器码
6.2 Dialect Conversion 框架
核心组件
-
ConversionTarget:定义哪些操作/Dialect 是合法的
-
RewritePatternSet:转换模式集合(ConversionPattern)
-
TypeConverter:类型转换规则
-
ConversionPatternRewriter:重写器,支持类型转换
6.3 ConversionPattern
cpp
struct AddOpLowering : public OpConversionPattern<AddOp> {
using OpConversionPattern<AddOp>::OpConversionPattern;
LogicalResult matchAndRewrite(
AddOp op,
AddOpAdaptor adaptor, // 转换后的操作数
ConversionPatternRewriter &rewriter) const override {
// 将 toy.add 替换为 arith.addf
rewriter.replaceOpWithNewOp<arith::AddFOp>(
op,
adaptor.getLhs(),
adaptor.getRhs()
);
return success();
}
};
6.4 ConversionTarget 合法性定义
cpp
ConversionTarget target(getContext());
// 标记合法的 Dialect
target.addLegalDialect<ArithDialect, FuncDialect>();
// 标记非法的 Dialect(需要被转换)
target.addIllegalDialect<ToyDialect>();
// 动态合法性判断
target.addDynamicallyLegalOp<func::FuncOp>(
[](func::FuncOp op) {
// 检查函数签名是否都已转换
return ...;
});
6.5 TypeConverter 类型转换
cpp
TypeConverter typeConverter;
// 基础类型直接转换
typeConverter.addConversion([](Type type) -> Type {
return type; // 原样返回
});
// Tensor 转 MemRef
typeConverter.addConversion(
[](TensorType type) -> Optional<Type> {
return MemRefType::get(type.getShape(),
type.getElementType());
});
// 函数类型转换
typeConverter.addConversion(
[&](FunctionType type) -> Type {
SmallVector<Type> inputs, results;
if (failed(typeConverter.convertTypes(
type.getInputs(), inputs)))
return nullptr;
if (failed(typeConverter.convertTypes(
type.getResults(), results)))
return nullptr;
return FunctionType::get(type.getContext(),
inputs, results);
});
6.6 应用转换
cpp
RewritePatternSet patterns(&context);
patterns.add<AddOpLowering, ConstantOpLowering>(
&typeConverter, &context);
// 完全转换:确保所有非法操作都被转换
if (failed(applyFullConversion(
module, target, std::move(patterns)))) {
signalPassFailure();
}
// 部分转换:允许部分操作保留
// applyPartialConversion(module, target, patterns);
6.7 常见 Lowering 路径
| 阶段 | 源 Dialect | 目标 Dialect | 说明 |
|---|---|---|---|
| 高层优化 | 自定义/TOSA | Affine + Linalg | 计算图转循环表示 |
| 循环优化 | Affine | SCF + MemRef | 仿射循环结构化 |
| 内存抽象 | MemRef | LLVM | 多维数组扁平化 |
| 代码生成 | LLVM Dialect | LLVM IR | 翻译到标准 LLVM |
七、Toy Tutorial 学习路径
7.1 章节概览
| 章节 | 主题 | 核心知识点 |
|---|---|---|
| Chapter 1 | Toy 语言与 AST | 语言定义、AST 设计、词法语法分析 |
| Chapter 2 | 生成基础 MLIR | 自定义 Dialect、ODS 定义 Operation、AST 遍历发 IR |
| Chapter 3 | 高层分析与转换 | Pattern Rewriting、Canonicalization、DRR |
| Chapter 4 | 接口与泛化转换 | ShapeInference Interface、内联、通用 Pass |
| Chapter 5 | 部分 Lowering 优化 | Dialect Conversion、Lowering to Affine |
| Chapter 6 | Lowering 到 LLVM | Full Conversion、LLVM Dialect、代码生成 |
| Chapter 7 | 自定义复合类型 | 自定义 Type、Type Conversion、结构体支持 |
7.2 学习建议
-
先理解概念:掌握 Dialect、Operation、Type 等核心概念再动手
-
跟着 Toy 走:逐章实现 Toy 编译器,理解每个阶段的作用
-
阅读源码:MLIR 源码是最好的学习资料,多看内置 Dialect 实现
-
实践项目:尝试为自己的 DSL 定义 Dialect
-
关注社区:MLIR 发展很快,新 Dialect 和优化不断加入
八、开发环境与工具
8.1 编译 MLIR
bash
git clone https://github.com/llvm/llvm-project.git
cd llvm-project
mkdir build && cd build
cmake -G Ninja ../llvm \
-DLLVM_ENABLE_PROJECTS=mlir \
-DLLVM_BUILD_EXAMPLES=ON \
-DLLVM_TARGETS_TO_BUILD="X86;NVPTX;AMDGPU" \
-DCMAKE_BUILD_TYPE=Release \
-DLLVM_ENABLE_ASSERTIONS=ON
ninja check-mlir
8.2 常用工具
-
mlir-opt:MLIR Pass 驱动,用于调试和运行 Pass 管线 -
mlir-translate:MLIR 与其他格式互译 -
mlir-tblgen:TableGen 代码生成器 -
mlir-reduce:测试用例缩减 -
mlir-lsp-server:LSP 语言服务器
8.3 mlir-opt 常用选项
bash
# 运行单个 Pass
mlir-opt --canonicalize input.mlir
# 运行 Pass 管线
mlir-opt --cse --canonicalize --inline input.mlir
# 查看所有可用 Pass
mlir-opt --help
# 打印 Pass 前后的 IR
mlir-opt --canonicalize
--print-ir-before-all
--print-ir-after-all input.mlir
九、进阶主题
9.1 数据分析与优化
-
数据流分析:DataFlowAnalysis 框架
-
别名分析:AliasAnalysis Interface
-
循环优化:Affine 循环变换(融合、平铺、向量化)
9.2 多线程与并行
-
PassManager 自动并行化函数级 Pass
-
GPU Dialect 支持 GPU 内核生成
-
Vector Dialect 支持 SIMD 向量化
9.3 TOSA 与 AI 编译
-
TOSA:神经网络算子标准 Dialect
-
Linalg:通用线性代数抽象
-
IREE、TPU-MLIR 等基于 MLIR 的 AI 编译器
十、学习资源
-
Toy Tutorial:https://mlir.llvm.org/docs/Tutorials/Toy/
-
论文:MLIR: A Compiler Infrastructure for the End of Moore's Law
-
Discourse:https://discourse.llvm.org/c/mlir/
说明:由于 B 站安全风控策略(错误 412),无法直接访问指定视频内容。本文档基于 MLIR 官方标准知识体系整理,涵盖了 MLIR 完整的核心概念、Dialect 机制、Pattern Rewriting、Pass 管理、Lowering 转换等主要内容,可作为系统学习 MLIR 的参考文档。
(注:部分内容可能由 AI 生成)