MLIR编译器手册，Dialect及Operation详解

Posted 2023-04-28 吴建明

MLIR编译器手册，Dialect及Operation详解

https://mlir.llvm.org/docs/LangRef/#symbol-reference-attribute

https://zhuanlan.zhihu.com/p/582517107

论文链接： MLIR: A Compiler Infrastructure for the End of Moore\'s Law

MLIR Language Reference

MLIR语言参考

 

MLIR（Multi-Level IR）是一种编译器中间表示，与传统的三地址SSA表示（如LLVM IR或SIL）相似，但它引入了多面体循环优化的概念作为一级概念。这种混合设计经过优化，可以表示、分析和转换高级数据流图以及为高性能数据并行系统生成的特定目标代码。除了它的代表性功能之外，它的单一连续设计提供了一个框架，可以从数据流图降低到高性能的目标特定代码。

 

本文档定义并描述了MLIR中的关键概念，旨在成为一份枯燥的参考文档——基本原理文档、术语表和其他内容托管在其他地方。

 

MLIR被设计为以三种不同的形式使用：一种适合调试的人类可读文本形式，一种适合编程转换和分析的内存形式，以及一种适合存储和传输的紧凑序列化形式。不同的形式都描述了相同的语义内容。本文档描述了人类可读的文本形式。

MLIR (Multi-Level IR) is a compiler intermediate representation with similarities to traditional three-address SSA representations (like LLVM IR or SIL), but which introduces notions from polyhedral loop optimization as first-class concepts. This hybrid design is optimized to represent, analyze, and transform high level dataflow graphs as well as target-specific code generated for high performance data parallel systems. Beyond its representational capabilities, its single continuous design provides a framework to lower from dataflow graphs to high-performance target-specific code.

This document defines and describes the key concepts in MLIR, and is intended to be a dry reference document - the rationale documentation, glossary, and other content are hosted elsewhere.

MLIR is designed to be used in three different forms: a human-readable textual form suitable for debugging, an in-memory form suitable for programmatic transformations and analysis, and a compact serialized form suitable for storage and transport. The different forms all describe the same semantic content. This document describes the human-readable textual form.

• High-Level Structure
• Notation
• Common syntax
• Top level Productions
• Identifiers and keywords
• Dialects
• Target specific operations
• Operations
• Builtin Operations
• Blocks
• Regions
• Definition
• Value Scoping
• Control Flow and SSACFG Regions
• Graph Regions
• Arguments and Results
• Type System
• Type Aliases
• Dialect Types
• Builtin Types
• Attributes
• Attribute Value Aliases
• Dialect Attribute Values
• Builtin Attribute Values
• IR Versioning

High-Level Structure

高层结构

MLIR基本上是基于节点（称为运算）和边（称为值）的类似图形的数据结构。每个值都是一个操作或块参数的结果，并且具有由类型系统定义的值类型。操作包含在块中，块包含在区域中。操作也在其包含块中排序，块在其包含区域中排序，尽管这种顺序在给定类型的区域中可能有语义意义，也可能没有语义意义）。操作还可以包含区域，从而能够表示层次结构。

操作可以代表许多不同的概念，从更高级别的概念，如函数定义、函数调用、缓冲区分配、缓冲区视图或切片以及进程创建，到更低级别的概念（如目标无关算术、目标特定指令、配置寄存器和逻辑门）。这些不同的概念由MLIR中的不同操作来表示，并且MLIR中可用的操作集可以任意扩展。

MLIR还使用熟悉的编译器过程概念，为操作转换提供了一个可扩展的框架。对任意一组操作启用任意一组传递会带来显著的扩展挑战，因为每次转换都可能考虑到任何操作的语义。MLIR通过允许使用特征和接口抽象地描述操作语义来解决这种复杂性，从而使转换能够更通用地操作操作。特征通常描述有效IR上的验证约束，使复杂的不变量能够被捕获和检查。

 

MLIR的一个明显应用是表示基于SSA的IR，如LLVM核心IR，通过适当的操作类型选择来定义模块、函数、分支、内存分配和验证约束，以确保SSA优势属性。MLIR包括一组方言，这些方言定义了这样的结构。然而，MLIR旨在足够通用，以表示其他类似编译器的数据结构，例如语言前端中的抽象语法树、目标特定后端中生成的指令或高级合成工具中的电路。

以下是MLIR模块的示例：

MLIR is fundamentally based on a graph-like data structure of nodes, called Operations, and edges, called Values. Each Value is the result of exactly one Operation or Block Argument, and has a Value Type defined by the type system. Operations are contained in Blocks and Blocks are contained in Regions. Operations are also ordered within their containing block and Blocks are ordered in their containing region, although this order may or may not be semantically meaningful in a given kind of region). Operations may also contain regions, enabling hierarchical structures to be represented.

Operations can represent many different concepts, from higher-level concepts like function definitions, function calls, buffer allocations, view or slices of buffers, and process creation, to lower-level concepts like target-independent arithmetic, target-specific instructions, configuration registers, and logic gates. These different concepts are represented by different operations in MLIR and the set of operations usable in MLIR can be arbitrarily extended.

MLIR also provides an extensible framework for transformations on operations, using familiar concepts of compiler Passes. Enabling an arbitrary set of passes on an arbitrary set of operations results in a significant scaling challenge, since each transformation must potentially take into account the semantics of any operation. MLIR addresses this complexity by allowing operation semantics to be described abstractly using Traits and Interfaces, enabling transformations to operate on operations more generically. Traits often describe verification constraints on valid IR, enabling complex invariants to be captured and checked. (see Op vs Operation)

One obvious application of MLIR is to represent an SSA-based IR, like the LLVM core IR, with appropriate choice of operation types to define Modules, Functions, Branches, Memory Allocation, and verification constraints to ensure the SSA Dominance property. MLIR includes a collection of dialects which defines just such structures. However, MLIR is intended to be general enough to represent other compiler-like data structures, such as Abstract Syntax Trees in a language frontend, generated instructions in a target-specific backend, or circuits in a High-Level Synthesis tool.

Here’s an example of an MLIR module:

// Compute A*B using an implementation of multiply kernel and print the // result using a TensorFlow op. The dimensions of A and B are partially // known. The shapes are assumed to match. func.func @mul(%A: tensor<100x?xf32>, %B: tensor<?x50xf32>) -> (tensor<100x50xf32>) // Compute the inner dimension of %A using the dim operation. %n = memref.dim %A, 1 : tensor<100x?xf32> // Allocate addressable "buffers" and copy tensors %A and %B into them. %A_m = memref.alloc(%n) : memref<100x?xf32> memref.tensor_store %A to %A_m : memref<100x?xf32> %B_m = memref.alloc(%n) : memref<?x50xf32> memref.tensor_store %B to %B_m : memref<?x50xf32> // Call function @multiply passing memrefs as arguments, // and getting returned the result of the multiplication. %C_m = call @multiply(%A_m, %B_m) : (memref<100x?xf32>, memref<?x50xf32>) -> (memref<100x50xf32>) memref.dealloc %A_m : memref<100x?xf32> memref.dealloc %B_m : memref<?x50xf32> // Load the buffer data into a higher level "tensor" value. %C = memref.tensor_load %C_m : memref<100x50xf32> memref.dealloc %C_m : memref<100x50xf32> // Call TensorFlow built-in function to print the result tensor. "tf.Print"(%C)message: "mul result" : (tensor<100x50xf32>) -> (tensor<100x50xf32>) return %C : tensor<100x50xf32> // A function that multiplies two memrefs and returns the result. func.func @multiply(%A: memref<100x?xf32>, %B: memref<?x50xf32>) -> (memref<100x50xf32>) // Compute the inner dimension of %A. %n = memref.dim %A, 1 : memref<100x?xf32> // Allocate memory for the multiplication result. %C = memref.alloc() : memref<100x50xf32> // Multiplication loop nest. affine.for %i = 0 to 100 affine.for %j = 0 to 50 memref.store 0 to %C[%i, %j] : memref<100x50xf32> affine.for %k = 0 to %n %a_v = memref.load %A[%i, %k] : memref<100x?xf32> %b_v = memref.load %B[%k, %j] : memref<?x50xf32> %prod = arith.mulf %a_v, %b_v : f32 %c_v = memref.load %C[%i, %j] : memref<100x50xf32> %sum = arith.addf %c_v, %prod : f32 memref.store %sum, %C[%i, %j] : memref<100x50xf32> return %C : memref<100x50xf32>

Notation

符号

MLIR有一个简单而明确的语法，允许它可靠地往返于文本形式。这对编译器的开发很重要，例如，对于理解正在转换的代码的状态和编写测试用例。

本文档描述了使用扩展Backus-Naur巴克斯-诺尔形式（EBNF）的语法。

这是本文档中使用的EBNF语法，用黄色方框表示。

MLIR has a simple and unambiguous grammar, allowing it to reliably round-trip through a textual form. This is important for development of the compiler - e.g. for understanding the state of code as it is being transformed and writing test cases.

This document describes the grammar using Extended Backus-Naur Form (EBNF).

This is the EBNF grammar used in this document, presented in yellow boxes.

alternation ::= expr0 | expr1 | expr2 // Either expr0 or expr1 or expr2. sequence ::= expr0 expr1 expr2 // Sequence of expr0 expr1 expr2. repetition0 ::= expr* // 0 or more occurrences. repetition1 ::= expr+ // 1 or more occurrences. optionality ::= expr? // 0 or 1 occurrence. grouping ::= (expr) // Everything inside parens is grouped together. literal ::= `abcd` // Matches the literal `abcd`.

Code examples are presented in blue boxes.

// This is an example use of the grammar above: // This matches things like: ba, bana, boma, banana, banoma, bomana... example ::= `b` (`an` | `om`)* `a`

Common syntax

The following core grammar productions are used in this document:

// TODO: Clarify the split between lexing (tokens) and parsing (grammar). digit ::= [0-9] hex_digit ::= [0-9a-fA-F] letter ::= [a-zA-Z] id-punct ::= [$._-] integer-literal ::= decimal-literal | hexadecimal-literal decimal-literal ::= digit+ hexadecimal-literal ::= `0x` hex_digit+ float-literal ::= [-+]?[0-9]+[.][0-9]*([eE][-+]?[0-9]+)? string-literal ::= `"` [^"\\n\\f\\v\\r]* `"` TODO: define escaping rules

Not listed here, but MLIR does support comments. They use standard BCPL syntax, starting with a // and going until the end of the line.

Top level Productions

// Top level production toplevel := (operation | attribute-alias-def | type-alias-def)*

The production toplevel is the top level production that is parsed by any parsing consuming the MLIR syntax. Operations, Attribute aliases, and Type aliases can be declared on the toplevel.

Identifiers and keywords

Syntax:

// Identifiers bare-id ::= (letter|[_]) (letter|digit|[_$.])* bare-id-list ::= bare-id (`,` bare-id)* value-id ::= `%` suffix-id alias-name :: = bare-id suffix-id ::= (digit+ | ((letter|id-punct) (letter|id-punct|digit)*)) symbol-ref-id ::= `@` (suffix-id | string-literal) (`::` symbol-ref-id)? value-id-list ::= value-id (`,` value-id)* // Uses of value, e.g. in an operand list to an operation. value-use ::= value-id value-use-list ::= value-use (`,` value-use)*

标识符命名实体，如值、类型和函数，并由MLIR代码的编写者选择。标识符可以是描述性的（例如%batch_size、@matmul），也可以在自动生成时是非描述性的（如%23、@func42）。值的标识符名称可以在MLIR文本文件中使用，但不会作为IR的一部分保留-打印机会给它们提供匿名名称，如%42。

MLIR通过在标识符前面加一个sigil（例如%、#、@、^、！）来保证标识符永远不会与关键字冲突。在某些明确的上下文（例如仿射表达式）中，为了简洁起见，标识符没有前缀。可以将新的关键字添加到MLIR的未来版本中，而不会有与现有标识符冲突的危险。

值标识符只在定义它们的（嵌套）区域的范围内，不能在该区域之外访问或引用。映射函数中的参数标识符在映射主体的作用域中。特定的操作可能会进一步限制哪些标识符在其区域的范围内。例如，具有SSA控制流语义的区域中的值的范围是根据SSA优势的标准定义来约束的。另一个例子是IsolatedFromAbove特性，它限制直接访问包含区域中定义的值。

函数标识符和映射标识符与符号相关联，并且具有依赖于符号属性的作用域规则

Identifiers name entities such as values, types and functions, and are chosen by the writer of MLIR code. Identifiers may be descriptive (e.g. %batch_size, @matmul), or may be non-descriptive when they are auto-generated (e.g. %23, @func42). Identifier names for values may be used in an MLIR text file but are not persisted as part of the IR - the printer will give them anonymous names like %42.

MLIR guarantees identifiers never collide with keywords by prefixing identifiers with a sigil (e.g. %, #, @, ^, !). In certain unambiguous contexts (e.g. affine expressions), identifiers are not prefixed, for brevity. New keywords may be added to future versions of MLIR without danger of collision with existing identifiers.

Value identifiers are only in scope for the (nested) region in which they are defined and cannot be accessed or referenced outside of that region. Argument identifiers in mapping functions are in scope for the mapping body. Particular operations may further limit which identifiers are in scope in their regions. For instance, the scope of values in a region with SSA control flow semantics is constrained according to the standard definition of SSA dominance. Another example is the IsolatedFromAbove trait, which restricts directly accessing values defined in containing regions.

Function identifiers and mapping identifiers are associated with Symbols and have scoping rules dependent on symbol attributes.

Dialects

方言

方言是参与和扩展MLIR生态系统的机制。它们允许定义新的操作以及属性和类型。每个方言都有一个唯一的名称空间，该名称空间以每个定义的属性/操作/类型为前缀。例如，仿射方言定义了名称空间：Affine。

 

MLIR允许多种方言，甚至是主树之外的方言，在一个模块内共存。方言是由某些通行证产生和使用的。MLIR提供了一个在不同方言之间和内部转换的框架。

MLIR支持的几种方言：

• Affine dialect
• Func dialect
• GPU dialect
• LLVM dialect
• SPIR-V dialect
• Vector dialect

 

Dialects are the mechanism by which to engage with and extend the MLIR ecosystem. They allow for defining new operations, as well as attributes and types. Each dialect is given a unique namespace that is prefixed to each defined attribute/operation/type. For example, the Affine dialect defines the namespace: affine.

MLIR allows for multiple dialects, even those outside of the main tree, to co-exist together within one module. Dialects are produced and consumed by certain passes. MLIR provides a framework to convert between, and within, different dialects.

A few of the dialects supported by MLIR:

• Affine dialect
• Func dialect
• GPU dialect
• LLVM dialect
• SPIR-V dialect
• Vector dialect

Target specific operations

目标特定操作

方言提供了一种模块化的方式，通过这种方式，目标可以直接向MLIR公开特定于目标的操作。例如，一些目标通过LLVM。LLVM具有一组丰富的内部函数，用于某些与目标无关的操作（例如，带溢出检查的加法），并为其支持的目标提供对目标特定操作的访问（例如，向量置换操作）。MLIR中的LLVM内部函数通过以“LLVM.”名称开头的操作来表示。

例子：

Dialects provide a modular way in which targets can expose target-specific operations directly through to MLIR. As an example, some targets go through LLVM. LLVM has a rich set of intrinsics for certain target-independent operations (e.g. addition with overflow check) as well as providing access to target-specific operations for the targets it supports (e.g. vector permutation operations). LLVM intrinsics in MLIR are represented via operations that start with an “llvm.” name.

Example:

// LLVM: %x = call i16, i1 @llvm.sadd.with.overflow.i16(i16 %a, i16 %b) %x:2 = "llvm.sadd.with.overflow.i16"(%a, %b) : (i16, i16) -> (i16, i1)

这些操作仅在将LLVM作为后端（例如CPU和GPU）时有效，并且需要与这些内部的LLVM定义保持一致。

These operations only work when targeting LLVM as a backend (e.g. for CPUs and GPUs), and are required to align with the LLVM definition of these intrinsics.

Operations

操作

语法：

Syntax:

operation ::= op-result-list? (generic-operation | custom-operation) trailing-location? generic-operation ::= string-literal `(` value-use-list? `)` successor-list? region-list? dictionary-attribute? `:` function-type custom-operation ::= bare-id custom-operation-format op-result-list ::= op-result (`,` op-result)* `=` op-result ::= value-id (`:` integer-literal) successor-list ::= `[` successor (`,` successor)* `]` successor ::= caret-id (`:` block-arg-list)? region-list ::= `(` region (`,` region)* `)` dictionary-attribute ::= `` (attribute-entry (`,` attribute-entry)*)? `` trailing-location ::= (`loc` `(` location `)`)?

MLIR引入了一个称为操作的统一概念，从而能够描述许多不同级别的抽象和计算。MLIR中的操作是完全可扩展的（没有固定的操作列表），并且具有特定于应用程序的语义。例如，MLIR支持独立于目标的操作、仿射操作和特定于目标的机器操作。

操作的内部表示很简单：一个操作由一个唯一的字符串（例如dim、tf.Conv2d、x86.repmovsb、ppc.eieio等）标识，可以返回零个或多个结果，接受零个或更多操作数，具有属性字典，具有零个或更多后继项，以及零个或更少封闭区域。通用打印表单从字面上包括所有这些元素，并带有一个函数类型来指示结果和操作数的类型。

示例：

MLIR introduces a uniform concept called operations to enable describing many different levels of abstractions and computations. Operations in MLIR are fully extensible (there is no fixed list of operations) and have application-specific semantics. For example, MLIR supports target-independent operations, affine operations, and target-specific machine operations.

The internal representation of an operation is simple: an operation is identified by a unique string (e.g. dim, tf.Conv2d, x86.repmovsb, ppc.eieio, etc), can return zero or more results, take zero or more operands, has a dictionary of attributes, has zero or more successors, and zero or more enclosed regions. The generic printing form includes all these elements literally, with a function type to indicate the types of the results

MLIR编译器调度与优化点滴

MLIR编译器调度与优化点滴

MLIR编译框架下软硬协同设计的思考

自从AI芯片成为热门的研究课题，众多关于AI芯片架构探索的学术文章不断涌现，大家从不同的角度对AI芯片进行架构分析及性能优化。MLIR是谷歌团队推出的开源编译器框架，颇受瞩目，灵活的编译器架构提升了其在众多领域应用的潜力。通过自定义IR的衔接，可以在架构探索和MLIR之间架起一座桥梁，在编译的过程中，自动进行硬件架构的探索和软件的优化编译，甚至生成硬件的代码，实现软硬协同设计。

架构探索方法的介绍

近十年，AI领域专用芯片的演进极大地促进了架构探索(指架构定义及性能分析)的发展，先后出现了众多的分析方法，这些分析方法针对AI计算过程中关键算子以及网络模型进行建模分析，从PPA(Power-Performance-Area)三个角度评估硬件性能。与此同时，伴随着AI编译框架的发展，尤其受益于MLIR编译器框架的可复用及可扩展性(详见MLIR多层编译框架实现全同态加密的讨论），将这些分析方法融入到MLIR框架中也变得十分可能，从而使用编译器对硬件架构进行探索。

架构分析中关注三个方面的表达，分别是计算架构(Computation Element)，存储结构(Memory Hierarchy )和互联结构(Interconnect)。对硬件架构进行性能分析时，数据流是搭建分析方法的基础，根据数据流的表达，将workload的计算过程映射到硬件架构的三类实现中。在学术研究中，Eyeriss [1]是较早将数据流引入到AI芯片的性能分析中，根据定义，AI的数据流可以分为三类，输出静止(Output Stationary)，权重静止(Weight Stationary)和行静止(Row Stationary)。随后的研究中，MAGNet[2]将其扩种为更多的描述方式，如图1所示，但还是围绕OS，WS和RS展开。根据数据流的划分，AI架构既可以分为这三类，比如NVDLA属于WS，Shi-dinanao属于OS，Eyeriss属于RS。相同的数据流架构可以采用类似的方法进行分析。

 

 图1 不同数据流对应的for-loop表示[2]

围绕数据流表示和硬件映射的表达上，可以归为三类，分别是以计算为中心 (computation-centric)的Timeloop[3], 以数据流为中心(data-centric)的MAESTRO[4]和以关系为中心(relation-centric)的TENET[5]。以计算为中心的表示方法关注的是for-loop表达在时间维度上映射到硬件架构；以数据流为中心的表达关注的是数据映射(data mapping)和复用(reuse)；以关系为中心的表达关注循环表达和计算单元及调度之间的关系。将对第二种data-centric的表达方式展开。

在MAESTRO的工作中，将data mapping和reuse作为一等公民，关注的是数据在时间和空间两个维度的复用。对于WS的计算架构，weight在时间维度上复用(相当于保持不变)，中间计算结果是在空间维度上复用，其复用如图2所示。

 

 图2 2x2kernel的卷积在WS类型加速器数据复用的表示[4] 关于时间和空间数据复用的表达，文中提出了一种IR的表示方式，我们称之为时域映射(Temporal Map)和空域映射(Spatial Map)。时域映射表示特定的维度与单个PE之间的映射关系，空域映射表示的是特定的维度与多个PE之间的映射关系，具体的表示如下：1.T(size, offset)α：α表示的特定的维度，比如权重的weight，width及channel等，size表示单个时间步长(time step)下α所在维度的index映射到单个PE的尺寸，offset表示的是相邻的时间步长的index偏移。对应的for循环表达如图3所示。2.S(size, offset)：α表示特定的维度，size表示维度α映射到每个PE的index的尺寸，offset表示映射到相邻PE的index偏移。

 

 图3 时域和空域映射与循环表达之间的对应关系[4]

假设一个计算架构有3个PE，卷积的权重大小为6，输入元素个数为17，步进为1，计算过程可以通过图4表示。在图中，标签1表示for循环的表达，标签2表示在时域和空域的IR表达，标签3表示数据在PE的分布及时间上的计算过程，图中可以看出cycle1到cycle4复用S中的index(1)，也就是weight保持静止。标签4表示空域映射、时域映射以及计算顺序，其中t表示按照所示的箭头方向依次计算。基于这样的IR表达及时间上的计算过程，就可以表示出一个WS架构的计算过程。

 

 图4 1D卷积操作在时域和空域的表示演示图

基于IR的性能分析方法

Aladdin[6]是较早开展基于编译的方式进行硬件的性能分析，将性能分析提前到RTL代码之前，避免了RTL代码及C-model大量的开发工作，基本的思路是将计算任务lowering到动态数据依赖图（DDDG：Dynamic Data Dependence Graph）级别，DDDG是针对特定架构的中间表达（Intermediate Representation）的表示，如图5所示。针对特定的硬件架构，分析DDDG的动态执行过程，即可评估出性能和功耗的数据，他们基于ILDJIT compiler IR[7]。

 

 图5 DDDG的计算表示[6] 

基于GEM5的工作，他们将其扩展为GEM5-Aladdin，用于对加速器系统级的性能分析，涵盖了SoC的接口通信开销，从而实现加速器架构和通信的协同设计。GEM5负责CPU和内存系统的性能分析，Aladdin负责加速器的性能分析。DDDG的表示从ILDJIT IR迁移到LLVMIR。

Interstellar[8]是将Halide语言用于AI架构的性能分析，数据流表达的方式属于computation-centric，核心工作是将和计算及数据流相关的for-loop转换到Halide的scheduling language，同时显性表达存储和计算。其中，关于架构和数据流是在Halide编译过程中的IR表达中引入，同时和Halide语言中的hardware primitive对应起来，将整个计算过程拆解到IR级别，然后映射到硬件结构，最后根据数据流的计算过程评估硬件的性能，整体过程如图6所示。最终采用调用硬件语言代码库的方式生成硬件设计。

 

 图6 标签1为Halide语言描述conv操作；标签2表示Halide Lowering过程中对in, compute_at, split及reorder调度原语(scheduling primitives)的IR表示；标签3表示调度原语和硬件架构的对应关系[8]

架构级别的IR

Micro-IR[9]文章的核心思想是将加速器的架构表示为一个并发的结构图(Concurrent Structural Graph)，每个组件就是一个架构级别的硬件单元，比如计算单元、网络或者存储器。结构图中显性地表达了加速器的构成组件，以及不同组件之间的数据流动，最终回归到数据流的表达和实现上。定义架构级别IR的好处在于1)将算法的表达和硬件架构解耦，2)将硬件的优化和RTL的代码实现解耦。这样一来，硬件架构IR层的优化工作可以单独展开。

整个编译的架构基于LLVM的实现，前端接入为AI framework，然后编译到LLVM IR，LLVM IR再对接到Micro-IR，在Micro-IR优化的PASS中聚焦就是前文提及到的关于数据流的映射、调度，tiling以及映射到硬件的intrinsic。最后对接到chisel的IR FIRTL，生成可综合的硬件语言。

 

 

 图7 Micro-IR的编译流程[9]

对于架构的表达，也是围绕数据流、存储和互联的展开，如图8所示，将一个简单的奇偶乘法翻译到IR图层，再翻译到IR的具体表达。

 

 

 图8 Micro-IR的编译表示[9]

MLIR中引入架构探索的可能性和挑战
可能性：1.经过上述章节的分析发现现有的性能分析方法的研究工作都有IR表示的思想，而且基于数据流的表示思想具有较好的理论基础，从时域和空域两个维度展开，也有很好的IR具体实现。2.基于IR性能分析的方法也处于不断演进的过程中，从ILDJIT到LLVM再到Halide，都证实了基于IR进行架构探索的可行性。同时不同的表示方式具有不同的有点，比如Halide中突出调度的思想，可以将该思想引入到MLIR中，形成schedule IR。3.关于硬件架构IR表示的文章也较多，比如Spatial[10]，文中举例的micro-IR 是比较典型的标准，与MLIR都基于LLVM的编译流程，将其引入到MLIR中作为硬件架构存在可能性。4.Union[11]是将MAESTRO性能分析的工作引入到MLIR的框架中，但是MAESTRO是作为架构探索的工具使用，没有接入到MLIR的编译流程中。
挑战：1.目前的架构探索都是基于相对规则的架构展开，没有涉及到复杂的工业界的芯片，存有一定的局限性，将其方法应用到工业界还有很大的隔阂。2.定义一个通用型的架构IR比较困难。架构是比较分散的，不同的任务需求有不同的架构设计，虽然架构设计从大的层面分为计算、存储和互联，但通过IR精准地刻画架构充满挑战，比如对于架构IR控制流的表示，Micro-IR中关于控制流的表达没有进行详细的阐述。3.在编译过程中，如何将软件任务能够自动翻译到架构IR上，同时能够对硬件架构进行自动调整和优化，这也是很大的挑战。目前是针对特定的已知架构，将计算任务映射到硬件。

总结了现有的针对AI架构的数据流分析方法，以及基于数据流分析方法构建的架构探索工具，同时介绍了现有的硬件架构的IR。这些丰富的分析方法和IR表示为架构探索引入到MLIR提供了可能性，也让我们看到了基于MLIR的编译器框架开展软硬协同设计的巨大潜力。 

基于MLIR实现GEMM编译优化

GEMM(General Matrix Multiplication)即通用矩阵乘法运算，由于其计算行为具有一定的复杂性以及规律性，是编译算法研究的绝佳场景。MLIR是近期非常热门的一个编译器软件框架，是工业界及科研界研究的一个热点，其提供了一套灵活的软件基础设施，对中间表达式（IR）及其相互之间的转换进行规范的管理，是一个非常友好的编译器开发平台[1][2]。即是分析在MLIR框架下，实现GEMM优化的内容，以及对MLIR在这一方面的实现优势的讨论。

GEMM优化策略介绍

矩阵乘法运算，由于其过程会包含大量的乘加操作，并且伴随大量的数据读写，因而如何充分利用好底层硬件的存储资源以及计算资源，是编译器对其性能优化的关键。目前，已有的一些优化策略主要包括：

1.矩阵分块（Tile）

当前的处理器性能主要受限于内存墙，即计算速度要大于数据存储的速度。为了打破内存墙的约束，各类硬件包括CPU及其他专用处理器，会设置不同层次的存储单元，而这些不同层级的存储单元往往大小以及读写速度不同，一般越靠近计算单元的存储其存储容量也越小但访问的速度也越快。如果可以将计算过程的数据局部化分块，而这些分块的数据可以独立完成计算，那么分块的数据就可以放在层次化的存储中，然后通过不同存储间建立Ping-Pong的数据传输方式，将数据存储与计算解耦，从而可以有效得隐藏存储墙的问题，提高计算效率。矩阵运算就有这种特点，因而可以通过矩阵分块来加速运算，如下图1所示，假设有两层存储，将输入矩阵A和B，以及输出矩阵C，根据存储大小划分成相应的小块，即m->mc，n->nc，k->kc，每次将Ac(mc, kc), Bc(kc,nc), Cc(mc, nc)送入到离计算单元更近的存储模块内，完成局部的计算后再进行下一次的计算。

 

 图1 矩阵运算的Tile操作示意图

在不同的底层硬件中，由于存储的层次以及不同层次的存储的容量大小不一样，分块的大小也会不一样。比如，文章[3]中对CPU而言，（Ac, Bc, Cc）划块的大小与cache大小一致，而为了充分利用register的资源，还会对（Ac, Bc, Cc）再进一步细划块成（Ar, Br, Cr），其尺寸大小与寄存器的数量一致。

2.向量化（Vectorize）

向量化的操作，主要是利用硬件的向量化指令或者SIMD（单指令多数）指令的特性，实现一个指令周期对多个值操作的能力。如下图2所示，通过将4个数据组成向量，利用处理器可以处理4个元素的新向量的计算能力，可以将4个指令周期的处理时间，压缩成1个指令周期的处理时间，从而极大提高运算处理能力。

 

 图2 vectorize操作示意图

3.循环展开（Unroll）

由于矩阵乘法有多层循环构成，如果底层硬件有一定的并行化能力，包括多线程多进程处理能力等，那么可以对循环进行适当展开，从而提高计算的并行度，增加并发执行的机会。如下图3所示，将一个次数为1024的循环，展开成256次循环，新的循环内又包含4条可以并行执行的展开计算，如果处理器能够并行处理循环内部的展开计算，那么通过对原来的循环展开，可以获得接近4倍的性能提升。

 

 图3  循环展开操作示意图

矩阵乘法的运算也包括其他的优化策略，比如数据重排等，但总体而言，各类编译器都是利用这些策略，充分利用硬件的存储及计算资源，达到最佳的运算效率。一些传统的计算库，如：OpenBLAS, BLIS, MKL等，开发时间长，性能也有比较优秀的表现。
MLIR实现GEMM优化

MLIR基于多层中间表示的方言（Dialect）思想，提供了一整套完善的编译器基础框架，可以帮助开发者快速实现编译策略想法的编译器。主要参考论文[4]，分析GEMM运算在MLIR中的实现，对应的硬件Target是因特尔i7-8700K处理器，每个核包含有32/256KB L1/L2 Cache以及多核共享的12MB L3 Cache，处理器支持AVX-2指令（256bit),优化目标是一个2088x2048xf64与2048x2048xf64的矩阵乘。

首先，其在高层次的Dialect上定义了一个矩阵运算的算子，这个算子的参数包含了输入矩阵（A，B）以及输出矩阵（C），同时为这个算子添加了tile/unroll 的尺寸等属性。如下图4所示，其中（M_C, K_C, M_R, N_R）属于Tile尺寸，K_U属于Unroll的大小。这里面（M_C, K_C）的选择是使得M_CxK_C大小的A矩阵块能够在L2 cache中复用，（K_C, N_R）的选择是使得K_CxN_R大小的B矩阵块能够在L1 cache中复用，（M_R, N_R）的选择是使得M_RxN_R大小的输出矩阵块能够在CPU Register中复用，这些值是根据硬件计算或者tunning出来的，在这里面的测试取了一个经验值。这些属性可以协助转换到更低一层的算子的策略实现，而选择哪些属性，则是跟编译算法以及编译的底层硬件对象有关，这些属性也是协助转换成下一层跟硬件更贴近的中间表示的实现，因而可以根据实际需要，灵活使用。

 

图4 GEMM算子的高层次定义

其次，MLIR的特点就是通过统一的多层中间表示，来实现对算子的层层低层化（lower）到具体的硬件目标上。针对上述高层次上定义的矩阵乘法算子，通过利用其所携带的优化属性，以及底层硬件的特点，设计了多条转换的路径（Pass），从而进一步把该算子lower到MLIR框架提供的中间辅助层（此中选择了Affine, Linalg，和Standard Dialect）。在这一层的转换过程中，基本包含了所有的策略，如：Tile，定制化复制，unroll，vectorize等。然后再将中间的辅组层的Dialect，进一步lower到LLVM Dialect上，如下图5所示。

 

 

 图5 GEMM算子Lowing的层次化Dialect示意图

最后，通过mlir提供的mlir-cpu-runner工具，可以运行最后生成的LLVM Dialect的结果。总体优化及运行测试的命令，如下图6所示。其中，“-hopt”，“-hopt-vect”等，是从高层的算子（hop.matmul）到中间辅组层的转换路径，每一条路径都包含有相应的编译策略，可以根据需要，灵活添加以及改变，“-convert-linalg-to-loops”, “-lower-affine”等时中间辅助层之间的转换，最后转换成LLVM Dialect。

 

 图6 MLIR运行GEMM的命令示意图

总体上，一个GEMM运算通过在MLIR框架下实现，并重写优化策略的路径，可以得到如图7所示的结果，其中箭头1对应包含了所有重写优化策略的MLIR实现，可以看到其能达到的计算速率为61.94GFLOPS，离理论上的峰值计算速率（75.2GFLOPS）比较接近，跟传统的计算库相比（如：BLIS，OpenBLAS，MKL等），也有着可以媲美的结果，其中的差距可能是传统的计算库有tunning的机制以及在编译器后端生成汇编指令及机器码有更成熟且高效的优化，因而可以得到更好的优化结果。总体而言，用MLIR重写的GEMM优化算法有着非常良好的表现。

 

 图7  MLIR编译运行结果与其他计算库的对比示意图

另一方面，MLIR框架提供了非常完善的C++以及Python接口，因而可以很方便接入已有的计算库，进行联合优化。在[4]文中尝试了用MLIR+BLIS的方法，将MLIR放在外侧（提供手动优化功能），BLIS则作为micro-kernel放在内侧（提供auto tunning功能），最终的结果如图7中箭头2所示。可以看出，对于DGEMM（双精度），通过MLIR与BLIS的联合优化，也可以达到接近峰值的性能，而其性能要比单独的MLIR或者BLIS优化要差一点。但其实在SGEMM（单精度）的测试中，MLIR+BLIS的优化又要比单独的MLIR或者BLIS优化要好一些，因而其中的性能在差异还需要进一步分析。总体而言，MLIR提供了非常完善的支持，可以融合已有的计算库的优化算法，去实现联合的编译优化。
MLIR实现GEMM优化的优势

通过上面对MLIR实现GEMM优化算法的编译的介绍，可以看出MLIR在其中有着非常突出的优势。

首先，MLIR框架提供了非常完善的编译器基础设施，可以让开发者不需要花费太多精力在编译器周边的实现，从而更加专注于编译算法的开发。同时，其基于多层中间表达的方式，可以让编译器更加模块化，可以让编译算法利用不同层次的中间表达的抽象信息，在不同的层次中逐步具体化，从而使得算法实现更加层次化，更加易于实现及管理。

其次，MLIR框架提供了一直到最底层硬件的表示支持，由于其可以层次化在不同的中间表示层实现编译算法，可以在高层次的中间表示中实现不依赖于底层硬件的通用算法，而在接近硬件底层中，开发针对性的路径实现相应的编译算法，因而可以很方便地针对不同硬件目标开发统一的编译环境。本人认为，这也是MLIR相对于一些现有的AI编译器，如：TVM等，最有优势的地方之一，由于其框架可以根据需要自行扩展Dialect，同时这些Dialect又在系统中遵循一套统一的范式进行管理，因而对不同的编译目标（硬件target）会有很强的扩展性，同时编译器的工程管理又可以做到非常好的统一性。

另外，MLIR框架提供了完善的C++/Python接口，可以很方便地接入已有的优化算法，快速实现算法迁移。

主要介绍了矩阵乘法运算在MLIR编译器框架实现的主要过程及内容，以及其在优化算法的多层次实现，以及接入已有优化算法的能力等方面的优势。MLIR编译框架，为编译器的开发者，在多层中间表达转换以及算法优化等方面提供强大的基础设施支持，降低开发编译器的门槛。

 

 

参考文献链接

https://mp.weixin.qq.com/s/s5_tA28L94arLdm5UijkZg

https://mp.weixin.qq.com/s/A1h4pJSJ8VF97DrZksNULg

[1] Y. Chen, J. Emer, and V. Sze, “Eyeriss: A spatial architecture for energy efficient dataflow for convolutional neural networks,” in Proc. ISCA,2016.

[2] R. Venkatesan, Y. S. Shao, M. Wang, J. Clemons, S. Dai, M. Fojtik, B. Keller, A. Klinefelter, N. R. Pinckney, P. Raina et al., “MAGNet: A Modular Accelerator Generator for Neural Networks,” in ICCAD, 2019

[3] A. Parashar, P. Raina, Y. S. Shao, Y. Chen, V. A. Ying, A. Mukkara, R. Venkatesan, B. Khailany, S. W. Keckler, and J. Emer, “Timeloop: A Systematic Approach to DNN Accelerator Evaluation,” in 2019 IEEE International Symposium on Performance Analysis of Systems and

Software, 2019

[4] H. Kwon, P. Chatarasi, V. Sarkar, T. Krishna, M. Pellauer, and A. Parashar, “Maestro: A data-centric approach to understand reuse, performance, and hardware cost of dnn mappings,” IEEE Micro, 2020.

[5] L. Lu, N. Guan, Y. Wang, L. Jia, Z. Luo, J. Yin, J. Cong, and Y. Liang, “TENET: A Framework for Modeling Tensor Dataflow Based on Relation-centric Notation,” in 2021 ACM/IEEE 48rd Annual International Symposium on Computer Architecture, 2021.

[6] S. Shao, B. Reagen, G.-Y. Wei, and D. Brooks, “Aladdin: A Pre-RTL, Power-Performance Accelerator Simulator Enabling Large Design Space Exploration of Customized Architectures,” in ISCA, 2014.

[7] S. Campanoni, G. Agosta, S. Crespi-Reghizzi, and A. D. Biagio, “A highly flexible, parallel virtual machine: Design and experience of ildjit,” Software Practice Expererience, 2010.

[8] X. Yang, M. Gao, Q. Liu, J. Setter, J. Pu, A. Nayak, S. Bell, K. Cao, H. Ha, P. Raina, C. Kozyrakis, and M. Horowitz, “Interstellar: Using halide’s scheduling language to analyze dnn accelerators,” in Proceedings of the Twenty-Fifth International Conference on Architectural

Support for Programming Languages and Operating Systems (ASPLOS), 2020.

[9] Sharifian, Amirali & Hojabr, Reza & Rahimi, Navid & Liu, Sihao & Guha, Apala & Nowatzki, Tony & Shriraman, Arrvindh. (2019). μIR -An intermediate representation for transforming and optimizing the microarchitecture of application accelerators. 940-953. 10.1145/3352460.3358292.

[10] David Koeplinger, MatthewFeldman, Raghu Prabhakar, Yaqi Zhang, Stefan Hadjis, Ruben Fiszel, Tian Zhao, Luigi Nardi, Ardavan Pedram, Christos Kozyrakis, and Kunle Olukotun. 2018. Spatial: A Language and Compiler for Application Accelerators. In Proceedings of the 39th ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI 2018).

[11] Geonhwa Jeong, Gokcen Kestor, Prasanth Chatarasi, Angshuman Parashar, Po-An Tsai, Sivasankaran Rajamanickam, Roberto Gioiosa, Tushar Krishna: Union: A Unified HW-SW Co-Design Ecosystem in MLIR for Evaluating Tensor Operations on Spatial Accelerators. CoRR abs/2109.07419 (2021)

[1] Chris Lattner, Mehdi Amini,Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle,Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: A compiler infrastructure for the end of moore\'s law, 2020

[2] MLIR：https://mlir.llvm.org/

[3] Tze Meng Low, etc. Analytical Modeling Is Enough for High-Performance BLIS. 2016.

[4] UdayBondhugula, High Performance Code Generation in MLIR: An Early Case Study With GEMM. 2020.

 

人工智能芯片与自动驾驶

以上是关于MLIR编译器手册，Dialect及Operation详解的主要内容，如果未能解决你的问题，请参考以下文章

新建MLIR一个Dialect，lowering，opt

从零开始学深度学习编译器十六，MLIR ODS要点总结上篇

从零开始学深度学习编译器十六，MLIR ODS要点总结上篇

从零开始学深度学习编译器十六，MLIR ODS要点总结上篇

MLIR设计与Dialect体系分析

MLIR编译器调度与优化点滴