MLIR中ARM指令SVE、SME的实现方式
1. MLIR基本概念
MLIR（Multi-Level Intermediate Representaion，多级中间表示）是一种用来构建可重用与可扩展编译的新方法。MLIR的设计初衷是为了解决软件碎片化问题，改进异构硬件的编译，显著减少构建特定领域编译器的成本以及帮助连接现有的编译器。当前很多语言都拥有自己的IR，如图1-33所示。
图1-33 很多语言都也有相应的IR
MLIR的实现原理是通过一种通用的架构来实现IR高层转换，如1-34所示。
图1-34 MLIR通过一种通用的架构来实现IR高层转换
核心组成部分包括dialect、operation、region等，如图1-35所示。
图1-35 方言核心组成包括方言、操作、区域等
dialect可以粗略理解为一个类，operation则是类中封装的函数。例如，对于arith.constant，arith是dialect名称，constant是operation名称，涵义为：调用arith方言中的constant操作（写法上类似于调用arith class中的constant成员函数）。通常，MLIR会将高层/高抽象化的方言，转换成底层/低抽象化的方言，最终生成LLVM IR。例如：方言A ->方言B ->...... -> LLVM方言-> LLVM IR。以下是一个arith方言-> LLVM方言-> LLVM IR的转换示例，代码如下：
//第1章/arith.mlir
func.func @vector_ops(%arg0: vector<[4]×i32>) -> vector<[4]×i32> {
%0 = arith.constant dense<2>: vector<[4]×i32>
%1 = arith.addi %arg0, %0: vector<[4]×i32>
return %1: vector<[4]×i32>
}
转换成LLVM方言，代码如下：
//第1章/arith_out.mlir
module attributes {llvm.data_layout = ""} {
llvm.func @vector_ops(%arg0: vector<[4]×i32>) -> vector<[4]×i32> {
%0 = llvm.MLIR.constant(dense<2>: vector<[4]×i32>): vector<[4]×i32>
%1 = llvm.add %arg0, %0: vector<[4]×i32>
llvm.return %1: vector<[4]×i32>
}
}
LLVM方言转换成LLVM IR，代码如下：
//第1章/arith.ll
; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"
declare ptr @malloc(i64)
declare void @free(ptr)
define <vscale×4×i32> @vector_ops(<vscale×4×i32> %0) {
%2 = add <vscale×4×i32> %0, shufflevector (<vscale×4×i32> insertelement (<vscale×4×i32> poison, i32 2, i64 0), <vscale×4×i32> poison, <vscale×4×i32> zeroinitializer)
ret <vscale×4×i32> %2
}
!llvm.module.flags = !{!0}
!0 = !{i32 2, !"Debug Info Version", i32 3}
介绍一些MLIR的基本概念。
1.1方言
方言主要由type、operations、interface、passes等构成。同时存在 ODS 与 DRR 两个重要的模块，这两个模块都是基于TableGen模块，ODS 模块用于定义操作，DRR 模块用于实现两个方言之间的转换。
1.2 操作
operation 是方言的重要组成部分，是抽象与计算的核心单元，可以看成是方言语义的基本元素，代码如下：
//第1章/tosa.sub.ll
%3 = "tosa.sub"(%2, %1): (tensor<1×f16>, tensor<32×1×1×128×f16>) -> tensor<32×1×1×128×f16>
上面的例子中，tosa.sub 中sub就是operation的名称，而它被定义在tosa方言中，用于做张量的减法运算。
1.3 区域
区域是嵌套在函数中的块集合，如图1-36所示。在LLVM IR当中，函数的主体是控制流图，控制流图由block组成；在MLIR当中，函数以operation的形式出现，operation中含有一个或者多个region，区域由块组成，而块中又含有操作。
图1-36 region是嵌套在函数中快的集合
MLIR的方言以及operation被定义在.td文件中，通过 tablegen 自动生成.inc 文件。例如 MLIR/include/MLIR/Dialect/ArmSME/ArmSME.td 中定义的arm_sme方言，通过tablegen生成 build/tools/MLIR/include/MLIR/Dialect/ArmSME/ArmSME.h.inc。
接下来会以SVE、SME为例，简要介绍它们在MLIR中的实现方式，以及本身的概念。SVE与SME都是Arm架构下的新指令集，分别作用于加速向量计算与矩阵计算。在进行向量/矩阵的计算时，SVE与SME有自己独特的指令以及寄存器，去完成这些工作并达到优化效果。新的指令意味着在LLVM IR中，SVE与SME有对应的内联。MLIR可以通过定义方言以及方言转换的方式，从高层方言的操作逐步下译，最终生成SVE/SME对应的内联，从而生成对应的LLVM IR。
2. SVE在MLIR中的实现
SME是基于SVE做的扩展，所以在了解SME之前，先简单介绍一下SVE。SVE与SME在MLIR当中都有自己的实现。SVE主要作用于向量计算，在后端拥有对应的内联操作之后。为了在MLIR中实现相关功能，需要：
1） 定义方言，如arm_sve方言、arm_sme方言，之后为方言添加相应的operation；
2）定义该方言中的operation向下下译成llvm方言/内联操作的转换；
3）定义其它上层方言到该方言的转换，比如tosa方言-> arm_sve方言。
在完成上述步骤之后，便初步打通了高层方言-> arm_sve方言-> llvm方言/内联操作-> LLVM IR的编译路径。
2.1 SVE的特性
SVE（Scalable Vector Extension，可缩放向量扩展）是Arm Aarch64架构下的下一代SIMD指令集，旨在加速高性能计算，SVE引入了很多新的架构特点，比如：
1） 可变向量长度；
2） 谓词标注SVE向量寄存器中参与计算的部分；
3）聚集加载与分散存储；
4）横向操作。
SVE中没有对于向量长度的定义，所以同样的二进制可以在不同向量长度的架构上运行。SVE向量的长度为128bit到2048bit（128bit的倍数），具体长度根据运行时中寄存器的可验证性决定。在特定情况下，向量化性能会超过传统向量化方式，因为SVE可能会增加向量长度。虽然没有明确定义向量的长度，但是SVE可以在指令中动态获取硬件的向量长度，并在向量的loop中以此为增量。下图分别列举了标量、NEON与SVE在计算一个长度为40字节的向量时的处理方法，如图1-37所示。
图1-37 Aarch64、NEON向量、SVE向量引擎
谓词是一个向量，由0与1组成。需要的部分标注为1，不需要的部分标注为0。它的功能与bit mask类似，不过bit mask会计算整个向量，然后将不需要的部分丢弃；而谓词则会关闭寄存器不需要的部分，只计算需要的部分。谓词应用示例，如图1-38所示。
图1-38 谓词应用示例
2.2 SVE的实现
MLIR可以将高层方言下译至LLVM IR中的SVE内联。由于SVE采用了单独的指令集与独特的实现方式，增加了arm_sve方言，其中含有所有的SVE操作。例如：
%1 = arm_sve.masked.addi %0, %arg0, %cst: vector<[4]×i1>, vector<[4]×i32>
这条操作代表arm_sve方言中的掩蔽加法运算。它的定义被写在.td 文件中，由TableGen进行生成。可将所有SVE能够实现的计算定义为arm_sve方言的各种操作。
需要定义每条操作到SVE 内联的下译方式（改写方式）。例如，上面的掩蔽加法运算就会在转换成LLVM方言时，被下译成如下的SVE内联：
%9 = "arm_sve.intr.add"(%8, %arg0, %0): (vector<[4]×i1>, vector<[4]×i32>, vector<[4]×i32>) -> vector<[4]×i32>
最终在LLVM IR中，这条内联会被改写为如下形式：
%4 = call <vscale×4×i32> @llvm.aarch64.sve.add.nxv4i32(<vscale×4×i1> %3, <vscale×4×i32> %0, <vscale×4×i32> shufflevector (<vscale×4×i32> insertelement (<vscale×4×i32> poison, i32 2, i64 0), <vscale×4×i32> poison, <vscale×4×i32> zeroinitializer))
这样，从arm_sve方言到LLVM IR的通道就已经打通了。接着，需要实现从高层方言到arm_sve方言的下译。对于SVE来说，需要定义arith方言到arm_sve方言之间的转换。例如，arm_sve.masked.addi是由arith.addi转换形成的。由于arm_sve.masked.addi是带有mask的向量加法操作，需要在下译时添加对于mask的定义。
%1 = arith.addi %arg0, %0: vector<[4]×i32>
下译成arm_sve方言如下所示：
//第1章/arm_sve.mlir
%c4 = arith.constant 4: index
%0 = vector.create_mask %c4: vector<[4]×i1>
%1 = arm_sve.masked.addi %0, %arg0, %cst: vector<[4]×i1>, vector<[4]×i32>
这样，从高层方言-> arm_sve方言-> LLVM IR的流程就已经完成了。
3. SME在MLIR中的实现
SME（Scalable Matrix Extension，可缩放矩阵扩展）在SVE基础之上，增加了对于矩阵的一些高效处理方法。它的原理是通过外积计算矩阵乘法，从而减少load次数，达到优化的效果。在MLIR当中，可以通过与SVE相同的方式进行定义，即定义SME方言，SME方言中的操作到LLVM方言/LLVM 内联的下译，以及上层方言（如向量方言）到SME方言的下译。
SME建立在可扩展向量扩展（SVE与SVE2）的基础上，增加了有效处理矩阵的新功能。主要功能包括：
1）矩阵网格存储；
2）加载、存储、插入与提取网格向量，包括动态换位；
3）SVE向量的外积；
4）流式SVE模式。
Arm上的矩阵乘法
矩阵乘法是许多关键工作负载的重要组成部分，如科学模拟、计算机视觉、机器学习（ML）的某些方面与增强现实（AR）。Arm架构随着时间的推移而不断发展，获得了提高这些操作的性能与效率的功能，如图1-39所示。
图1-39 ARM上的普通矩阵乘法
1）Armv8.4-A：支持8位整数DOT产品指令
2）Armv8.6-A：支持向量内整数与浮点矩阵乘法指令，以及BFloat16数据类型。
3）Armv9-A：支持SVE2中更宽的向量。
SME是这一流程的下一步，能够显著提高CPU矩阵处理吞吐量与效率。
3.1向量的外积
普通矩阵乘法，如图1-40所示。
图1-40 普通矩阵乘法
令A，B，C为4×4 f32矩阵，计算A * B = C的过程为：取A的一行，B的一列，点乘得到C中一个元素，代码如下：
//第1章/matrix_mul.mlir
C(0, 0) = A(0,:) * B(:, 0)
C(0, 1) = A(0,:) * B(:, 1)
...
C(3, 3) = A(3,:) * B(:, 3)
假如只有两个128bit向量寄存器，需要16次load操作。
外积
矩阵外积计算，如图1-41所示。
图1-41 矩阵外积
令A、B、C为4×4 f32矩阵，计算A * B = C的过程为：取A的一列，取B的一行，通过外积计算出一个4×4的矩阵，然后累加到C上，代码如下：
//第1章/matrix_mul_accumulation.mlir
C = A(:, 0) * B(0,:)
C += A(:, 1) * B(1,:)
C += A(:, 2) * B(2,:)
C += A(:, 3) * B(3,:)
假如只有两个128bit向量寄存器，需要8次load操作。
3.2 SVE模式流
在模式流之下，SVE/SME允许改变原有的向量长度。新的向量长度叫做SVL（流向量长度，流向量长度）。它的长度不定，但必须为128bit ~ 2048bit中128bit的倍数。假如流向量长度为256bit，代码如下：
//第1章/SVL_low.arith
SVL-B：SVL中字节数量 = 256 / 8 = 32
SVL-H：SVL中16bit数量 = 256 / 16 = 16
SVL-S：SVL中32bit数量 = 256 / 32 = 8
SVL-D：SVL中64bit数量 = 256 / 64 = 4
SVL-Q：SVL中128bit数量 = 256 / 128 = 2
...
3.3 SME ZA存储
ZA存储是SME独特的数据储存方式，这是一个大小SVL-B * SVL-B字节的寄存器。例如：流向量长度为256bit，则SVL-B为256 / 8 = 32。ZA存储大小为32 * 32 * 8 bit。下面是一些相关概念的定义：
1）ZA存储为一个寄存器，含有一个大小为32×32字节的2d ZA数组；
2）ZA数组的大小为32×32字节；
3）ZA网格是ZA数组的一部分，一个或多个ZA网格组成一个ZA数组；
4）ZA切片是ZA网格中的一行或一列。
一个ZA数组，每个方格代表8 bit，大小为32 * 32 个8 bit，如图1-42所示。
图1-42 ZA数组，每个方格代表8 bit，大小为32 * 32 个8 bit
对于不同带宽的数据，ZA 网格的数量、大小、表示方法也都不同。例如：
1）对于8 bit数据：大小为SVL-B * SVL-B * 8 bit，即32 * 32 * 8 bit。一个ZA数组中只有一个8 bit类的ZA网格，大小为32 * 32。其中：
1） ZA0B：整个ZA数组。
3）对于16 bit数据：大小：SVL-H * SVL-H * 16 bit，即16 * 16 * 16 bit。一个ZA数组中有2个16 bit类的ZA网格，大小为16 * 32。其中：
2） ZA0H：ZA数组的0，2，4，6...行；
3） ZA1H：ZA数组的1，3，5，7...行。
6）对于32 bit数据：大小：SVL-S * SVL-S * 32 bit，即8 * 8 * 32 bit。一个ZA数组中有4个32 bit类的ZA网格，大小为8 * 32。其中：
7）ZA0S：ZA数组的0，4，8，12...行；
8） ZA1S：ZA数组的1，5，9，13...行；
9）ZA2S：ZA数组的2，6，10，14...行；
10）ZA3S：ZA数组的3，7，11，15...行。
以此类推。
ZA切片为ZA网格中的一行或一列，拥有特定表示方法，例如：ZA0H.S[2]表示 ZA0S这个网格中的第2行，H表示行（若是V则表示列），2表示行号（列号）。同理，ZA4V.D[2]表示 ZA4D网格中的第二列，ZA2V.S[1]表示ZA2S网格中的第一列，如图1-43所示。
图1-43 ZA数组与网格分布
3.4 SME的实现
与SVE十分类似，同样去为SME定义自己的方言。之后会在arm_sme方言中添加相应的操作，并定义这些操作到arm_sme 内联的下译流程，以及从高层的向量方言到arm_sme方言的下译流程。例如SME的mopa操作如下所示：
arm_sme.mopa za0s, %pred.32, %pred.32, %1, %1: vector<[4]×i1>, vector<[4]×i1>, vector<[4]×f32>, vector<[4]×f32>
这条操作会对两个向量进行外积计算，并将结果累加到za0d所对应的ZA网格中。它会被下译成如下的SME 内联：
"arm_sme.intr.mopa"(%28, %27, %27, %arg0, %arg0): (i32, vector<[4]×i1>, vector<[4]×i1>, vector<[4]×f32>, vector<[4]×f32>) -> ()
转换成LLVM IR之后形式如下：
call void @llvm.aarch64.sme.mopa.nxv4f32(i32 0, <vscale×4×i1> %11, <vscale×4×i1> %11, <vscale×4×float> %0, <vscale×4×float> %0)
参考文献
- https://zhuanlan.zhihu.com/p/420729459
- https://iq.opengenus.org/MLIR-compiler-infrastructure/
- https://MLIR.llvm.org
- https://www.youtube.com/watch?v=Y4SvqTtOIDk
- https://arxiv.org/abs/2002.11054v2
- https://zhuanlan.zhihu.com/p/189589184
- https://www.youtube.com/watch?v=eGCcPo4UAHs
- https://community.arm.com/arm-community-blogs/b/architectures-and-processors-blog/posts/scalable-matrix-extension-armv9-a-architecture
https://developer.arm.com/documentation/ddi0616/latest
