ARM指令SVE、SME在MLIR中的实现方式

ARM指令SVE、SME在MLIR中的实现方式

1. MLIR基本概念

MLIR（Multi-Level Intermediate Representaion，多级中间表示）是一种用来构建可重用和可扩展编译的新方法。MLIR的设计初衷是为了解决软件碎片化问题，改进异构硬件的编译，显著减少构建特定领域编译器的成本以及帮助连接现有的编译器。当前很多语言都拥有自己的IR：

 

图1

MLIR的实现原理是通过一种通用的架构来实现IR高层转换：

 

图2

MLIR的核心组成部分包括dialect、operation、region等。

 

图3

dialect可以粗略理解为一个class，operation则是class中封装的函数。例如，对于arith.constant，arith是dialect名称，constant是operation名称，涵义为：调用arith方言中的constant operation（写法上类似于调用arith class中的constant成员函数）。通常，MLIR会将高层/高抽象化的方言，转换成底层/低抽象化的方言，最终生成LLVM IR。例如：方言A ->方言B -> ...... -> LLVM方言-> LLVM IR。以下是一个arith方言-> LLVM方言-> LLVM IR的转换示例：

arith.mlir：

func.func @vector_ops(%arg0: vector<[4]xi32>) -> vector<[4]xi32> {

  %0 = arith.constant dense<2> : vector<[4]xi32>

  %1 = arith.addi %arg0, %0 : vector<[4]xi32>

  return %1 : vector<[4]xi32>

}

转换成LLVM方言，结果如下：

arith_out.mlir：

module attributes {llvm.data_layout = ""} {

  llvm.func @vector_ops(%arg0: vector<[4]xi32>) -> vector<[4]xi32> {

    %0 = llvm.mlir.constant(dense<2> : vector<[4]xi32>) : vector<[4]xi32>

    %1 = llvm.add %arg0, %0  : vector<[4]xi32>

    llvm.return %1 : vector<[4]xi32>

    }

}

LLVM方言转换成LLVM IR，结果如下：

arith.ll：

; ModuleID = 'LLVMDialectModule'

source_filename = "LLVMDialectModule"

 

declare ptr @malloc(i64)

 

declare void @free(ptr)

 

define <vscale x 4 x i32> @vector_ops(<vscale x 4 x i32> %0) {

  %2 = add <vscale x 4 x i32> %0, shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 2, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer)

  ret <vscale x 4 x 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 模块用于定义operation，DRR 模块用于实现两个方言之间的转换。

1.2 操作

operation 是方言的重要组成部分，是抽象和计算的核心单元，可以看成是方言语义的基本元素。

%3 = "tosa.sub"(%2, %1) : (tensor<1xf16>, tensor<32x1x1x128xf16>) -> tensor<32x1x1x128xf16>

上面的例子中，tosa.sub 中 sub 就是operation的名称，而它被定义在 tosa方言中，用于做tensor的减法运算。

1.3 region

region是嵌套在function当中的blocks的集合。在LLVM IR当中，function的主体是控制流图，控制流图由block组成；在MLIR当中，function以operation的形式出现，operation中含有一个或者多个region，region由block组成，而block中又含有operation。

 

图4

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方言/内联 operation的转换；
  3. 定义其它上层方言到该方言的转换，比如tosa方言-> arm_sve方言。

在完成上述步骤之后，便初步打通了高层方言-> arm_sve方言-> llvm方言/内联 op -> LLVM IR的编译路径。

2.1 SVE的特性

SVE（Scalable Vector Extension，可缩放向量扩展）是Arm Aarch64架构下的下一代SIMD指令集，旨在加速高性能计算，SVE引入了很多新的架构特点，比如：

  • 可变矢量长度
  • Predicate标注SVE向量寄存器中参与计算的部分
  • 聚集加载和分散存储
  • 横向操作

SVE中没有对于vector长度的定义，所以同样的二进制可以在不同向量长度的架构上面运行。SVE vector的长度为128bit到2048bit（128bit的倍数），具体长度根据runtime中寄存器的avaliability决定。在特定情况下，向量化性能会超过传统向量化方式，因为SVE可能会增加向量长度。虽然没有明确定义vector的长度，但是SVE可以在指令中动态获取硬件的向量长度，并在vector的loop中以此为增量。下图分别列举了标量、NEON和SVE在计算一个长度为40字节的向量时的处理方法：

 

图5

Predicate是一个vector，由0和1组成。需要的部分标注为1，不需要的部分标注为0。它的功能与bit mask类似，不过bit mask会计算整个vector，然后将不需要的部分丢弃；而Predicate则会关闭寄存器不需要的部分，只计算需要的部分。Predicate应用如下图所示：

 

图6

2.2 SVE的实现

MLIR可以将高层方言下译至LLVM IR中的SVE内联。由于SVE采用了单独的指令集和独特的实现方式，社区增加了arm_sve方言，其中含有所有的SVE operation。例如：

%1 = arm_sve.masked.addi %0, %arg0, %cst : vector<[4]xi1>, vector<[4]xi32>

这条operation代表arm_sve方言中的masked addition operation。它的定义被写在 .td 文件中由TableGen进行生成。于是可以将所有SVE能够实现的计算定义为arm_sve方言的各种operation。

之后需要定义每条operation到SVE 内联的下译方式（改写方式）。例如，上面的masked addition operation就会在转换成LLVM方言时被下译成如下的SVE 内联：

%9 = "arm_sve.intr.add"(%8, %arg0, %0) : (vector<[4]xi1>, vector<[4]xi32>, vector<[4]xi32>) -> vector<[4]xi32>

最终在LLVM IR中，这条内联会被改写为如下形式：

%4 = call <vscale x 4 x i32> @llvm.aarch64.sve.add.nxv4i32(<vscale x 4 x i1> %3, <vscale x 4 x i32> %0, <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 2, i64 0), <vscale x 4 x i32> poison, <vscale x 4 x 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]xi32>

下译成arm_sve方言如下所示：

%c4 = arith.constant 4 : index

%0 = vector.create_mask %c4 : vector<[4]xi1>

%1 = arm_sve.masked.addi %0, %arg0, %cst : vector<[4]xi1>, vector<[4]xi32>

这样，从高层方言-> arm_sve方言-> LLVM IR的流程就已经完成了。

3. SME在MLIR中的实现

SME（Scalable Matrix Extension，可缩放矩阵扩展）在SVE基础之上，增加了对于矩阵的一些高效处理方法。它的原理是通过外积计算矩阵乘法，从而减少load次数，达到优化的效果。在MLIR当中，可以通过与SVE相同的方式进行定义，即定义SME方言，SME方言中的op到LLVM方言/LLVM 内联的下译，以及上层方言（如向量方言）到SME方言的下译。

3.1向量的外积

普通矩阵乘法

 

图7

令A，B，C为4x4 f32矩阵，计算A * B = C的过程为：取A的一行，B的一列，点乘得到C中一个元素。

  C(0, 0) = A(0, :) * B(:, 0)

  C(0, 1) = A(0, :) * B(:, 1)

  ...

  C(3, 3) = A(3, :) * B(:, 3)

假如只有两个128bit向量寄存器，需要16次load操作。

外积

 

图8

令A、B、C为4x4 f32矩阵，计算A * B = C的过程为：取A的一列，取B的一行，通过外积计算出一个4x4的矩阵，然后累加到C上。

  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：

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。下面是一些相关概念的定义：

  • ZA存储为一个寄存器，含有一个大小为32 x 32字节的2d ZA数组。
  • ZA数组的大小为32 x 32字节。
  • ZA瓦片是ZA数组的一部分，一个或多个ZA瓦片组成一个ZA数组。
  • ZA切片是ZA瓦片中的一行或一列。

下图为一个ZA数组，每个方格代表8 bit，大小为32 * 32 个8 bit。

 

图9

对于不同带宽的数据，ZA 瓦片的数量、大小、表示方法也都不同。例如：

  • 对于8 bit数据：大小为SVL-B * SVL-B * 8 bit，即32 * 32 * 8 bit。一个ZA数组中只有一个8 bit类的ZA瓦片，大小为32 * 32。其中：
  • ZA0B：整个ZA数组。
  • 对于16 bit数据：大小：SVL-H * SVL-H * 16 bit，即16 * 16 * 16 bit。一个ZA数组中有2个16 bit类的ZA瓦片，大小为16 * 32。其中：
  • ZA0H：ZA数组的0，2，4，6 ...行；
  • ZA1H：ZA数组的1，3，5，7 ...行。
  • 对于32 bit数据：大小：SVL-S * SVL-S * 32 bit，即8 * 8 * 32 bit。一个ZA数组中有4个32 bit类的ZA瓦片，大小为8 * 32。其中：
  • ZA0S：ZA数组的0，4，8，12 ...行；
  • ZA1S：ZA数组的1，5，9，13 ...行；
  • ZA2S：ZA数组的2，6，10，14 ...行；
  • ZA3S：ZA数组的3，7，11，15 ...行。

以此类推。

ZA切片为ZA瓦片中的一行或一列，拥有特定表示方法，例如：ZA0H.S[2]表示 ZA0S这个瓦片中的第2行，H表示行（若是V则表示列），2表示行号（列号）。同理，ZA4V.D[2]表示 ZA4D瓦片中的第二列，ZA2V.S[1] 表示ZA2S瓦片中的第一列，如下图所示：

 

图10

3.4 SME的实现

与SVE十分类似，同样去为SME定义自己的方言。之后会在arm_sme方言中添加相应的operation，并定义这些operation到arm_sme 内联的下译流程，以及从高层的向量方言到arm_sme方言的下译流程。例如SME的mopa operation如下所示：

arm_sme.mopa za0s, %pred.32, %pred.32, %1, %1 : vector<[4]xi1>, vector<[4]xi1>, vector<[4]xf32>, vector<[4]xf32>

这条operation会对两个向量进行外积计算，并将结果累加到za0d所对应的ZA瓦片中。它会被下译成如下的SME 内联：

"arm_sme.intr.mopa"(%28, %27, %27, %arg0, %arg0) : (i32, vector<[4]xi1>, vector<[4]xi1>, vector<[4]xf32>, vector<[4]xf32>) -> ()

转换成LLVM IR之后形式如下：

call void @llvm.aarch64.sme.mopa.nxv4f32(i32 0, <vscale x 4 x i1> %11, <vscale x 4 x i1> %11, <vscale x 4 x float> %0, <vscale x 4 x float> %0)

 

参考文献

  1. https://zhuanlan.zhihu.com/p/420729459
  2. https://iq.opengenus.org/mlir-compiler-infrastructure/
  3. https://mlir.llvm.org
  4. https://www.youtube.com/watch?v=Y4SvqTtOIDk
  5. https://arxiv.org/abs/2002.11054v2
  6. https://zhuanlan.zhihu.com/p/189589184
  7. https://www.youtube.com/watch?v=eGCcPo4UAHs
  8. 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

https://mp.weixin.qq.com/s/W4Dc-3ESWB9k6Z7nLsf8ag

 

posted @ 2024-04-01 04:41  吴建明wujianming  阅读(13)  评论(0)  编辑  收藏  举报