MLIR 入门完全指南

1. MLIR 是什么？

1.1 核心定义

定位：通用编译器基础设施（Unified Compiler Infrastructure）
核心思想：
- 混合IR：支持多级别抽象（从高层计算图到低层硬件指令）
- 可扩展性：允许自定义领域特定IR（如TensorFlow IR、HLO、GPU方言等）

1.2 设计目标

解决传统编译器的问题：
- 不同框架（TensorFlow/PyTorch）和硬件（CPU/GPU/TPU）需要重复开发编译器
- 优化过程碎片化（如XLA/LLVM各自为政）
统一基础设施：覆盖从计算图优化到芯片指令生成的完整流程

2. MLIR 核心功能

2.1 关键能力

数据流图表示
- 支持动态形状、用户自定义算子（Op）、变量（如TF Variables）
- 示例：TensorFlow的计算图可直接转换为MLIR方言
跨层级优化
- 高层优化（如算子融合、常量折叠）
- 低层优化（循环分块、内存布局变换、向量化）
硬件特定支持
- 加速器专用操作（如TPU指令、GPU Warp操作）
- 显式内存管理（DMA插入、缓存控制）
渐进式Lowering
- 从抽象操作逐步降级到硬件指令（如 linalg.matmul → loops → vector → llvm）

3. MLIR 核心概念

3.1 基础术语

术语	解释	示例
SSA	静态单赋值：每个变量只赋值一次	`%x = add %a, %b`
CSE	公共子表达式消除	消除重复计算的 `%x + %x`
DCE	死代码消除	删除未被引用的 `%unused = ...`
Invariant	循环不变量	可提到循环外的计算
Legalization	非法操作转换为合法操作	将 `sin(x)`展开为泰勒级数

3.2 核心数据结构

Operation (Op)
- 语义最小单元，可包含属性、操作数、结果和区域
```
%result = "dialect.op_name"(%arg1, %arg2) {attr = 42} : (i32, i32) -> i64
```

Block (^bb)

顺序执行的Op序列，必须以终止符（如 return）结尾

^bb0(%arg: i32):
  %1 = "add"(%arg, %arg) : (i32, i32) -> i32
  "return"(%1) : (i32) -> ()

Region
- 由多个Block组成的子图，用于表示控制流（如 if/while的then/else分支）
Module
- 顶级容器，包含所有Op（类似C++的命名空间）

4. MLIR 工作流程

4.1 典型编译流程

前端转换：框架计算图 → MLIR方言（如 mhlo）
中端优化：
- 通用优化（CSE/DCE）
- 领域优化（量化、算子融合 Fusion）
硬件Lowering：
- 逐步降级到低级方言（如 vector→llvm→nvvm）
代码生成：目标二进制（CPU/GPU机器码）

4.2 示例：矩阵乘法优化

// 原始HLO级表示
%result = "mhlo.dot"(%A, %B) : (tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>

// 降级为循环嵌套
%result = "linalg.matmul"(%A, %B) : (memref<1024x1024xf32>, memref<1024x1024xf32>) -> memref<1024x1024xf32>

// 分块+向量化
"transform.vectorize"(%matmul) : !transform.any_op

以下是关于 MLIR Builtin Dialect 的详细介绍，适合编译器初学者系统学习：

MLIR Builtin Dialect 完全解析

Builtin Dialect 是 MLIR 的核心基础方言，定义了所有其他方言共享的基础数据类型、结构和操作。它类似于编程语言中的“标准库”，是 IR 的底层基石。

1. Builtin Dialect 的定位

1.1 核心作用

提供基础设施：定义所有方言共用的基础类型（如整数、浮点数）、属性（如字符串、数组）和操作（如模块、函数）。
容器功能：作为顶级容器（Module、FuncOp）的承载者。
IR 结构基础：定义 Block、Region、Operation 等核心结构的表示方式。

1.2 特点

隐式加载：无需显式声明，所有 MLIR 文件自动可用。
不可扩展：用户不能添加新的 Builtin 类型或操作（与自定义方言不同）。

2. 关键组成部分

2.1 基础类型（Types）

类型	语法示例	描述
整数	`i32`、`i64`	有符号整数（`si32`）、无符号（`ui8`）
浮点数	`f32`、`f64`	IEEE 浮点
索引类型	`index`	循环索引/内存偏移（平台相关）
张量	`tensor<2x3xf32>`	静态形状多维数组
MemRef	`memref<4x?xi32>`	带有内存布局的可变缓冲区
函数类型	`(i32, f32) -> i64`	输入输出类型签名

2.2 属性（Attributes）

属性	语法示例	描述
整数	`42 : i32`	带类型的常量
浮点数	`3.14 : f32`
字符串	`"hello"`
数组	`[1, 2, 3]`	同类型元素列表
字典	`{key = value}`	键值对集合
符号引用	`@func_name`	全局符号（如函数名）

2.3 核心操作（Ops）

(1) 模块操作（ModuleOp）

作用：IR 的顶级容器，包含所有全局定义（如函数、全局变量）。

示例：

module {
  func.func @main() -> i32 {
    %c1 = arith.constant 1 : i32
    return %c1 : i32
  }
}

(2) 函数操作（FuncOp）

作用：定义可调用的函数单元，包含参数、返回值和函数体。

示例：

func.func @add(%a: i32, %b: i32) -> i32 {
  %sum = arith.addi %a, %b : i32
  return %sum : i32
}

(3) 基本控制流操作

操作	示例	描述
`return`	`return %value : i32`	函数返回
`call`	`%r = call @add(%a, %b)`	调用函数

3. Builtin 的底层结构

3.1 Operation 的通用表示

所有 MLIR 操作（包括自定义方言）都隐式使用 Builtin 的结构：

// 通用操作语法
%result = "dialect.op_name"(%operand1, %operand2) {
  attr1 = value1,
  attr2 = value2
} : (input_type1, input_type2) -> result_type

3.2 Block 和 Region

Block：由一组顺序执行的 Operation 组成，以终止符（如 return）结尾。

^bb0(%arg: i32):  // 带参数的Block
  %1 = arith.addi %arg, %arg : i32
  "some_dialect.some_op"(%1) : (i32) -> ()
  return

Region：包含一组 Block 的子图，用于结构化控制流（如 scf.if 的 then/else 分支）。

Builtin Dialect 是理解 MLIR 设计哲学的起点，掌握它后能更轻松地学习其他方言（如 arith、func、scf）。

MLIR与LLVM IR的区别

MLIR 和 LLVM IR 都是编译器中间表示（IR），但设计目标和应用场景有显著区别。以下是它们的核心对比：

1. 设计哲学

维度	LLVM IR	MLIR
定位	低级静态单赋值（SSA）IR，面向通用硬件代码生成	多级混合IR，支持从计算图到硬件的全栈抽象
核心目标	生成高效的机器码	统一异构编译（深度学习/HPC/DSL）
扩展性	固定指令集（不可扩展）	可自定义方言（Dialect）

2. 关键特性对比

2.1 抽象级别

LLVM IR：
- 接近机器码的底层IR（类似RISC指令集）
- 强制要求显式内存管理、指针运算、基本块控制流
```
; LLVM IR示例（计算 %a + %b）
%result = add i32 %a, %b
```

MLIR：

支持多级抽象，可同时表示高层计算图和低层硬件操作
允许携带领域特定信息（如张量形状、量化参数）

// MLIR示例（高层张量操作 + 低层循环）
%result = "tensor.add"(%A, %B) : (tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
"scf.for"(%i = 0 to 1024) {
  %elem = "memref.load"(%A, %i) : (memref<1024xf32>, index) -> f32
  ...
}

2.2 类型系统

特性	LLVM IR	MLIR
基础类型	仅支持机器兼容类型（i32, f64, ptr）	扩展类型系统（张量、索引、符号引用等）
类型安全	弱类型（指针可任意转换）	强类型 + 自定义类型约束
动态形状	不支持	支持（如 `tensor`）

2.3 控制流表示

LLVM IR：
- 基于基本块（BasicBlock）和分支指令（br/switch）
- 所有控制流必须显式展开
```
; LLVM IR条件分支
%cond = icmp eq i32 %a, 0
br i1 %cond, label %true_block, label %false_block
```