MLIR学习笔记

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 重写思想。

两种定义方式

  1. 命令式(C++):继承 RewritePattern,手动实现 matchAndRewrite

  2. 声明式(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 &registry) 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 框架

核心组件

  1. ConversionTarget:定义哪些操作/Dialect 是合法的

  2. RewritePatternSet:转换模式集合(ConversionPattern)

  3. TypeConverter:类型转换规则

  4. 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 学习建议

  1. 先理解概念:掌握 Dialect、Operation、Type 等核心概念再动手

  2. 跟着 Toy 走:逐章实现 Toy 编译器,理解每个阶段的作用

  3. 阅读源码:MLIR 源码是最好的学习资料,多看内置 Dialect 实现

  4. 实践项目:尝试为自己的 DSL 定义 Dialect

  5. 关注社区: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 编译器

十、学习资源


说明:由于 B 站安全风控策略(错误 412),无法直接访问指定视频内容。本文档基于 MLIR 官方标准知识体系整理,涵盖了 MLIR 完整的核心概念、Dialect 机制、Pattern Rewriting、Pass 管理、Lowering 转换等主要内容,可作为系统学习 MLIR 的参考文档。

(注:部分内容可能由 AI 生成)

相关推荐
从零开始的代码生活_1 小时前
C++ string 详解:常用接口、字符串算法与深拷贝实现
开发语言·c++·后端·学习·算法
阿哟阿哟3 小时前
Ansys Electronics Desktop(hfss)仿真PCB(AD)
笔记
~kiss~3 小时前
LLM 的 层归一化(稳定训练) - Layer Normalization & RMSNorm
学习
MartinYeung53 小时前
[论文学习]LLM-based AI Agent 安全威胁与防御系统性综述
人工智能·学习·安全
阿米亚波3 小时前
【C++ STL】std::deque
数据结构·c++·笔记·算法·stl
再玩一会儿看代码3 小时前
JUnit 测试框架详解:从实际开发、业务测试到 Java 面试高频问题
java·经验分享·笔记·junit·面试
巴巴媛6664 小时前
STM32学习笔记【36.CAN收发实验】
笔记·stm32·学习
光学补偿4 小时前
文件IO的前世今生
开发语言·学习·面试·java-ee
xian_wwq4 小时前
【学习笔记】分布式推理:TP / PP / EP / CP 并行策略详解(20/35)
笔记·分布式·学习·ai