Skk's Blog|

🧐 MLIR基本结构

MLIR代码的基本结构

首先从一段示例的MLIR代码开始

!custom_struct = type !llvm.struct<(i32, i32)>

module {
	func.func private @printf(%arg: i32)

	func.func @function_add(%arg0: i32, %arg1: i32) -> i32 attributes { attr_name = "value" } {
		%c = arith.addi %arg0, %arg1 : i32
		return %c : i32
	}

	func.func @main() -> i32 {
		%a = arith.constant 1 : i32
		%b = arith.constant 2 : i32
		%result = func.call @function_add(%a, %b) : (i32: i32) -> i32
		return %result : i32
	}
}

在MLIR中, 大部分的代码通常都包含在一个顶层Module中, 这个Module中包含一个或者多个函数定义或者函数声明. 函数体内包含了各种操作(Operation). 除此之外还可以自定义一些复合类型, 例如上面的第1行代码.

LLVM的基本运作单位是指令(你可以将LLVM视作一个与特定硬件无关的虚拟机, 它接受指令(LLVM IR)并执行), 类似的在MLIR中, 基本运作单位是操作(Operation), 所有的MLIR代码可以视为由操作为节点构成的一棵树.

MLIR是树形结构, 每个节点都是Operation(或者叫Op), Op可以组成Block, Block可以组成Region, 而Region又可以嵌套在Op内部.

func.func @example(%arg0: i32) -> i32 {  // 函数操作包含一个 Region
  %0 = arith.constant 0 : i32           // Region 中的第一个 Block 开始
  cf.cond_br %condition, ^bb1, ^bb2
  
^bb1:  // 第二个 Block
  %1 = arith.addi %0, %arg0 : i32
  cf.br ^bb3(%1 : i32)
  
^bb2:  // 第三个 Block
  %2 = arith.subi %0, %arg0 : i32
  cf.br ^bb3(%2 : i32)
  
^bb3(%result: i32):  // 第四个 Block，带有参数
  return %result : i32
}

Block指的是基本块, 它由一个或者多个Op组成(有点类似于C里面goto的标签). Region指区域, 它由多个Block组成, 比如循环体和函数体(通常用{}表示). 每一个基本块都可以有块参数(Block Arguments), 这些参数可以在块内部访问. Block和Block之间可以通过分支跳转连接, 比如cf.cond_br. 而每个Block都必须以终结操作(Terminator)或者分支跳转结束. 一般来讲, Block是控制流的基本单位, 而Region则是结构化的范围.

Module一般来讲是MLIR代码的顶层结构, 可以被视作MLIR树的根节点.

操作(Operation)

一条操作的基本结构如下:

%0 = toy.tranpose(%1) { inplace = "true" } : (tensor<2x3xf32>) -> tensor<3x2xf32>
  1. 操作的返回值, 以%开头. 一个操作可以有零个或者多个返回值. 操作的返回值都是SSA(Static Single Assignment)形式的, 也就是每个变量只能被赋值一次. 因此可以很方便地进行数据流分析, 因为你可以找到定义它的唯一操作(除非变量是块参数).
  2. 操作的名称, 例如上面的toy.transpose. .前面的内容是操作所属的方言(Dialect, 即将提到, 暂且理解为namespace). .后面的内容则是操作的名称.
  3. 操作的操作数/参数(Operand), 包含在()中. 一个操作可以有零个或者多个操作数, 并且这些操作数都是SSA形式.
  4. 操作的属性(Attribute), 包含在{}中, 用键值对的形式表示. 可以简单地理解为Operand是运行时才知道的值, 而Attribute是编译时就知道的值.
  5. 操作的类型(Type), 放在:后面. 类型包括两者, 操作数类型和返回值类型. 这些类型可以是简单类型(例如i32), 复合类型(例如tensor<2x3xf32>), 也可以是自定义类型(例如!custom_struct).

方言(Dialect)

一些特定操作的集合被称为方言(Dialect), 例如上面已经出现的toy和cf. 方言可以被视作一个命名空间, 其中包含了一些特定的操作和类型. 如果你想把操作理解为机器的指令, 那方言可以理解成指令集.

MLIR设计的初衷就是为了高扩展性, 因此允许用户自定义自己的方言, 本身也提供了很多常用的方言, 比如cf(控制流), scf(结构化控制流), arith(算术)等等. 不同的方言之间可以互相转换(Conversion), 这是MLIR的一个重要特性. MLIR中的ML(Multi-Level)正是通过多种不同抽象级别的方言直接互相转换实现的.

MLIR的数据流结构

前面已经提到了MLIR适合用于构建数据流图, 在MLIR - Understanding The IR Structure中有所提及

DefUseChains

这幅图表明了每个Value的定义流(Def Chain), 每个Value要么来自另一个操作的结果(OpOperand), 要么来自块参数(BlockOperand), 并且来源是唯一的, 可以很轻松地根据这幅图知晓某个操作数是在哪被定义的.

Uselist

这幅图是每个Value的使用流图(Use Chain), 它是一个双向链表, 在替换某个Value的所有用法时非常有用(例如在优化时重写代码).

有几点需要注意的是:

  1. 操作的Operand成员是到Value的指针, 如果你只想修改当前操作的输入Operand, 需要通过OpOperand来访问, 否则会修改到原始的Value.
  2. 如果修改了OpOperand, 那么其对应Value的使用流图也会被修改掉, 但定义流图不会被改.
  3. 由于Value都是SSA形式的, 因此可以通过value->getDefiningOp()获取定义它的操作, 但如果Value是块参数, 意味着没有定义它的操作, getDefiningOp()会返回nullptr.

这些使用流图是可以遍历和修改的, MLIR提供了相关的接口, 暂时留到后续再说.

总之到现在应该大概清楚了MLIR的基本工作结构. 接下来我们会开始介绍怎么创建一个最简单的MLIR方言.

本页内容