Mindspore实习-AKG SIG算子addlayernorm编辑和合并

什么是算子

计算图和算子在计算本质上是一致的。算子是打包后的计算图,计算图是拆包后的算子

比如sigmod复合算子可以看作一个计算图,由基础算子Exp,Add,Reciprocal等基础算子组成,用小规模的“基本算子”集合就可以表达任意现有计算图。

计算图可以完全由基础算子组成,但是还是得定义复合算子,比如sigmod这种,因为对于基本算子计算图来说,相邻算子之间只能通过全局内存(或显存)进行数据传递。而对于复合算子来说,相邻的基本计算之间则可以通过局部内存或者寄存器进行数据传递。除了性能之外,在一些场景下,通过算子融合也能有效减少对全局内存的的实际占用。

mindspore的算子融合方案是什么

在TVM、XLA等自动算子编译技术出现之前,AI框架主流采用手工融合的方式解决如上问题。主要思路是:

  • 手工融合: 1)识别常见的热点算子组合子图,比如: Add(Mul(x, y))。然后针对该算子子图手工实现对应融合算子; 2)将融合算子注册到AI框架,并在AI框架中增加对应的优化pass,将匹配的算子子图替换为融合算子节点。 这种方式的缺点是显而易见的。因为它只能针对若干热点场景进行融合,所以是无法做到通用和泛化的。

  • XLA:XLA最早基于TensorFlow开发。不过它采用了与TensorFlow不同的计算图IR表示。所以在TensorFlow中,需要将TensorFlow的计算图IR首先转换为XLA的计算图IR(HLO)。然后XLA基于HLO进行融合优化等,最后通过LLVM等后端编译生成相应融合算子。由于采用独立的IR,XLA具有较好的可移植性,目前在PyTorch、JAX等非TensorFlow框架中也有不错的表现。

  • TVM:TVM主要用于推理场景。在架构上,主要包括relay和tir两层。其通过relay导入推理模型,然后进行融合优化,最后通过tir生成融合算子。TVM在算子编译方面采用compute和schedule分离的技术,并且不同算子compute所需要的schedule通常是不同的。为了更好支持不同融合算子场景,TVM支持对算子进行自动tuning,来生成较优的切分参数甚至schedule。由于tuning空间较大,目前tuning时间相对还是比较长的。

基于以上问题背景,并结合MindSpore自身需求,我们提出图算融合解决方案。其主要思路是:

  • 图算融合解决方案: 1)3.9以通用的pattern识别计算图中的融合场景,并生成融合算子子图;2)将生成的融合算子子图通过自动算子编译技术(AKG)生成对应的融合算子。

什么是图算融合编译技术

我们如果需要把模型部署到CPU甚至手机上去,此时需要其他硬件及其架构的支持。开发者们往往会根据实际情况选择各种各样的深度学习顶层框架训练模型,例如昇思MindSpore等,再把训练好的模型部署到各种各样的设备后端包括GPU、CPU、FPGA、昇腾AI处理器及其它新型的AI加速器上。

考虑到不同硬件设备的特性千差万别、现有算子库中算子包含范围不同、新型加速器算子库支持不足、非常规的神经网络中存在不常见的layer等等情况,开发者要完成手写算子并保证性能,学习成本和时间成本都变得很高,所以图算融合编译技术的出现变得非常有必要。

这里有一个介绍

什么是自动算子编译技术AKG

AKG是Auto Kernel Generator的简称,"SIG"是一个常见的缩写,全称为 “Special Interest Group”,中文直译为“特殊兴趣小组”或者”特殊利益集团“。

AKG(Auto Kernel Generator)对深度神经网络中的算子进行优化,并提供特定模式下的算子自动融合功能。AKG与MindSpore的图算融合功能协同工作,可提升在不同硬件后端上运行网络的性能。AKG基于polyhedral技术,多面体。

任务-addlayernorm

【任务背景】
AKG-MLIR 已经在 MindSpore Dialect 内提供了基础算子的定义和对应的Lower流程。但是,当前MindSpore Dialect算子集是基于传统的Bert/Transformer类网络的需求梳理和统计的。随着网络的变化,我们有更多算子的需求希望添加在 MindSpore Dialect 之中。

【需求描述】

  1. 在MindSporeOps.td中添加对应算子的实现
  2. 打通算子的Lower流程(直接lower到Linalg或者lower到TOSA),相关Dialect的代码(如MindSporeToTosa.cpp或者MindSporeToLinalg.cpp)
  3. 提供对应的测试用例(算子的info文件和对应的op_dsl)

【参考资料】

  1. 样例代码仓:https://gitee.com/mindspore/akg/pulls/989
  2. 代码添加教程wiki:https://gitee.com/monkeykingd/akg/wikis/AKG-MLIR算子添加用例
  3. 测试用例教程wiki:https://gitee.com/monkeykingd/akg/wikis/AKG-MLIR单算子测试用例构建

【验收标准】
根据参考资料中的参考测试文件(demo.info)构建测试用例,使用python ${path_to_py_benchmark}/py_benchmark.py -e cpu -f mul.info --dump_ir 1运行命令,和numpy模拟的算子结果一致

【任务技术要求】
Python, MindSpore > 2.3

我要做什么?
把mindspore的AKG安装起来然后根据参考资料里的,利用基础算子写addlayernorm这个算子的实现然后添加lower流程然后再写一个测试示例,把整个代码push上去,让导师通过。

AKG安装

官方文档里写了从MindSpore侧构建运行AKG代码但是详细细节写了个寂寞,然后全是独立构建方法

AKG安装过程

从mindspore的whl安装,失败,但是后面还是这么做了

独立构建方法网上有,我偏不用,看看这个从MindSpore侧怎么安装,理论上AKG是mindspore的子仓库应该是自动安装了。

  • 首先我创建一个一个python3.9 conda环境然后安装Mindspore2.3安装方法是下载whl包,这个没有出现问题日
1
pip install mindspore-2.3.1-cp39-cp39-linux_x86_64.whl

不过如果是实习的话,需要推自己的改动,那我岂不是最好git下载然后source安装mindspore,算了。后来发现确实是这样,因为这样找不到tests文件夹,就无法运行官方文档里面的设置环境变量和运行测试用例,想了一下还是进行源码安装,因为又可以看源码又可以安装环境但是后来

从mindspore的源码安装

我选择安装的是CPU的ubuntu x86架构的源码,参考mindspore的仓库主页CPU安装源码

首先安装以下依赖软件,我的麻

软件名称 版本 作用
Ubuntu 18.04 编译和运行MindSpore的操作系统
Python 3.9-3.11 MindSpore的使用依赖Python环境
wheel 0.32.0及以上 MindSpore使用的Python打包工具
setuptools 44.0及以上 MindSpore使用的Python包管理工具
PyYAML 6.0-6.0.2 MindSpore里的算子编译功能依赖PyYAML模块
Numpy 1.19.3-1.26.4 MindSpore里的Numpy相关功能依赖Numpy模块
GCC 7.3.0到9.4.0之间 用于编译MindSpore的C++编译器
git - MindSpore使用的源代码管理工具
CMake 3.22.2及以上 编译构建MindSpore的工具
tclsh - MindSpore sqlite编译依赖
patch 2.5及以上 MindSpore使用的源代码补丁工具
NUMA 2.0.11及以上 MindSpore使用的非一致性内存访问库
LLVM 12.0.1 MindSpore使用的编译器框架(可选,图算融合以及稀疏计算需要)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 以上依赖的安装过程
conda create -n mindspore39AKGSource python=3.9
pip install wheel
pip install -U setuptools
pip install pyyaml
pip install "numpy>=1.19.3,<=1.26.4"
sudo apt-get install gcc-7 git tcl patch libnuma-dev -y # 发现都安装了
# cmake已经安装了,但是发现了一个apt安装cmake的方法如下
wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | sudo apt-key add -
sudo apt-add-repository "deb https://apt.kitware.com/ubuntu/ $(lsb_release -cs) main"
sudo apt-get install cmake -y
# # 以上依赖中的LLVM的安装过程,安装LLVM是可选的,但是我不知道哪里可选
wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
sudo add-apt-repository "deb http://apt.llvm.org/bionic/ llvm-toolchain-bionic-12 main"
sudo apt-get update
sudo apt-get install llvm-12-dev -y
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 编译安装mindspore 2.3.1 
git clone https://gitee.com/mindspore/mindspore.git
git checkout v2.3.1 #要更换版本,但是后来想了一下好像也没必要
bash build.sh -e cpu -j4 -S on # 默认从github下载依赖源码,当-S选项设置为on时,从对应的gitee镜像下载。这一步只是编译成whl,要安装还得install

#报错了configure: WARNING: *** Could not find Flex on your system
sudo apt-get install flex
sudo apt-get install bison
#又报错了PYTHON_INCLUDE_DIRS = /usr/include/python3.10 PYTHON_LIBRARIES = /usr/lib/x86_64-linux-gnu/libpython3.10.so MS LIBS CACHE PATH: /home/outbreak/mindspore/AKG/mindspore/build/mindspore/.mslib
#CMake Error at cmake/gencode.cmake:13 (message):Generate operator python/c++ definitions FAILED.
# 然后这个报错我解决不了,艘不到也看不懂

# 但是我突然想到,这一步不就是生成whl的吗,所以我干脆只下载源码但是用网上编译好的whl算了,然后结合同一个版本的mindspore仓库进行编辑测试。后经过了测试,发现这样并不行,因为运行akg/test文件的时候并没有找到test_add.py,报错ERROR test_abs.py - RuntimeError: Cannot find the files.
#我觉得是我的whl安装过程并没有链接这个源码中的AKG的文件夹,我的猜测是这样的。所以最终怎么安装请看下面的最终解决方案

# 如果bash build.sh -e cpu -j4 -S on成功的话请继续
pip install output/mindspore-*.whl -i https://pypi.tuna.tsinghua.edu.cn/simple
# 这一步在联网状态下,安装whl包时会自动下载mindspore安装包的依赖项(依赖项详情参见setup.py中的required_package),其余情况需自行安装。运行模型时,需要根据ModelZoo中不同模型指定的requirements.txt安装额外依赖,常见依赖可以参考requirements.txt。
# 验证安装是否成功
python -c "import mindspore;mindspore.set_device(device_target='CPU');mindspore.run_check()"
# 这一步应该会输出MindSpore version: 版本号
The result of multiplication calculation is correct, MindSpore has been installed on platform [CPU] successfully!

最终的解决方案mindspore的whl安装+AKG源码安装

ok fine ,最后终于成功了,整理一下思路,就是说从官网只下载AKG的源码,当然mindspore的仓库里有AKG这个子仓库,然后我用官方编译好的whl进行安装,安装版本为2.3.1,这个安装没有出现问题

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
conda create -n mindspore39NLP python=3.9
pip install mindspore-2.3.1-cp39-cp39-linux_x86_64.whl #这就安装完成了
# 以下为AKG的独立构建方式,针对CPU
git clone https://gitee.com/mindspore/akg.git
cd akg
bash build.sh -e cpu -j8

# 中途报错ModuleNotFoundError: No module named 'decorator'
pip install decorator

# 中途报错OSError: /home/outbreak/anaconda3/envs/mindspore39NLP/bin/../lib/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /home/outbreak/mindspore/AKG/akg/build/libakg.so)
strings /home/outbreak/anaconda3/envs/mindspore39NLP/lib/libstdc++.so.6 | grep GLIBC #这个是查询当前的conda环境里有哪些stdc++,其中发现只有GLIBCXX_3.4.29,并没有GLIBCXX_3.4.30
strings /usr/lib/x86_64-linux-gnu/libstdc++.so.6 | grep GLIBCXX #这个是查询本地环境中有哪些stdc++,conda环境里某些包需要某个版本的c++,一般会从本地环境中复制进conda环境中,如果没有则报错,解决方法是链接进conda环境里
ln -sf /usr/lib/x86_64-linux-gnu/libstdc++.so.6 /home/outbreak/anaconda3/envs/mindspore39NLP/bin/../lib/libstdc++.so.6


cd tests
# 设置环境变量,这里有很大的作用,source的位置一定要对,要不然找不到测试文件,卡bug
source ./test_env.sh cpu
# AKG测试
cd tests/st
python run.py -e cpu -o add -l level0 # 执行CPU Add算子的level0用例,包含了下面的test_abs
cd tests/st/ops/
pytest -s test_abs.py -m "level0 and platform_x86_cpu" # 运行CPU level0测试用例

AddLayerNorm算子的定义

首先addlayernorm的公式是怎样的,看一下pytorch没找到,但是找到了升腾的一个方案。文献没有找到,但是在 Transformer 架构中,残差连接(Add)与层归一化(LayerNorm)常被结合使用,是将输入的X1和X2进行逐元素相加,然后对结果进行layernorm。LayerNorm 最初由 Jimmy Lei Ba 等人在 2016 年的论文 《Layer Normalization》

image-20250210193652012

任务1在MindSporeOps.td中添加对应算子的实现

开头已经了解了mindspore AKG是什么了,里面提供了基本算子的定义和Lower流程,但现在需要添加新的复合算子AddLayerNorm到MindSpore Dialect中。

这里需要会LLVM编译器的相关操作,所谓td文件就是其中的编辑文件

  • TableGen是LLVM生态中用于生成代码的声明式语言,通过.td文件定义数据模型(如操作、指令、寄存器)。
  • 在MLIR中:td文件用于定义Dialect中的算子(Op)、类型(Type)、属性(Attribute)等,生成C++代码框架。

如果要写一个addLaynorm,可能如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
//有三个输入,一个输出,所以在这里定义了四个MindSpore_Tensor类型的tensor,对应五个输入X1、X2、gamma、beta和一个输出y。但是升腾的方案写的很复杂,还包括了是否要输出中间的相加结果
// arguments:输入张量和属性(如epsilon)。
// results:输出张量。
// assemblyFormat:定义文本格式的语法,用于调试和序列化。
// MindSporeOps.td中
//===----------------------------------------------------------------------===//
// MindSpore Operator: addlayernorm
//===----------------------------------------------------------------------===//
def MindSpore_AddLayerNormOp : MindSpore_Op<"addlayernorm", [Pure]> {
let summary = "Add followed by Layer Normalization";
let arguments = (ins
MindSpore_Tensor:$input_a,
MindSpore_Tensor:$input_b,
MindSpore_Tensor:$input_gamma,
MindSpore_Tensor:$input_beta
);
let results = (outs
MindSpore_Tensor:$output
);
}

ins后接输入参数,out后接输出参数。在编译后会自动生成get方法来获取输入输出,并将蛇形名称转换为驼峰式命名法。在akg-mlir/build/include/akg/Dialect/MindSpore/IR/MindSporeOps.cpp.inc可以找到编译后的get方法

比如:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
::mlir::TypedValue<::mlir::TensorType> AccMulOp::getInputA() {
return *getODSOperands(0).begin();
}

::mlir::TypedValue<::mlir::TensorType> AccMulOp::getInputB() {
return *getODSOperands(1).begin();
}

::mlir::TypedValue<::mlir::TensorType> AccMulOp::getInputGamma() {
return *getODSOperands(2).begin();
}
::mlir::TypedValue<::mlir::TensorType> AccMulOp::getInputBeta() {
return *getODSOperands(2).begin();
}

任务2打通算子的Lower流程

Lower流程讲解

直接lower到Linalg或者lower到TOSA,相关Dialect的代码(如MindSporeToTosa.cpp或者MindSporeToLinalg.cpp)

在MLIR(Multi-Level Intermediate Representation)生态中,DialectLinalgTOSA 是不同层次的中间表示(IR),用于描述和优化计算任务。

  • Dialect:是 MLIR 中的用于定义特定领域或抽象层次的中间表示。每个 Dialect 包含一组操作(Op)、类型(Type)和属性(Attribute),用于描述特定领域的计算任务。

    例如:

    • MindSpore Dialect:描述 MindSpore 框架中的算子(如 AddLayerNorm)。
    • TOSA Dialect:描述面向硬件后端的张量运算操作。
    • Linalg Dialect:描述线性代数操作(如矩阵乘法、卷积)。
  • Linalg:是 MLIR 中的一个 Dialect,专注于线性代数操作(如矩阵乘法、卷积、点积等)。适合循环优化和高层次中间表示。

  • TOSA:全称Tensor Operator Set Architecture。里面都是张量运算,也就是比Linalg高了一个等级,它定义了一组硬件友好的操作,适合在 AI 加速器(如 NPU、GPU)上执行。适合面向硬件后端的标准化操作。在Tosa中已经定义有矩阵逐元素乘的实现,如下

  • Lower流程(Lowering):是将高层抽象的 Dialect 逐步转换为低层抽象的 Dialect 的过程。(如从MindSpore Dialect到Linalg/TOSA)。

    • MindSpore Dialect 到 Linalg:将 MindSpore 的算子(如 AddLayerNorm)转换为 Linalg 的线性代数操作。将 AddLayerNorm 分解为 linalg.addlinalg.mul 等操作。
    • MindSpore Dialect 到 TOSA:将 MindSpore 的算子转换为 TOSA 的标准化操作。将 AddLayerNorm 分解为 tosa.addtosa.mul 等操作。

    MindSpore的Lower过程这里有两条路,可以从MindSpore转到Tosa再转到Linalg,也可以直接从MindSpore转到Linalg

编写AddLayerNorm的Lower过程

参考的算子编写示例中讲解,测试的时候给算子输入的是JSON格式的文件,然后呢需要首先将JSON的格式转化为MindSpore Dialect 。

然后有两个路,一个是MindSpore Dialect 转化为Tosa Dialect 再转化为Linalg Dialect ,另一个是MindSpore Dialect 直接转化为Linalg Dialect。

JSONtoMindSpore

akg-mlir/compiler/lib/Target/MindsporeDialect/TranslateToMindsporeDialect.cpp下的void MindBuilder::initMindOpFactory()中,case by case的注册刚写好的AddLayerNormOp

1
2
this->mindOpFactory["addlayernorm"] = &MindBuilder::convertTernaryOp<mindspore::AddLayerNormOp>;

MindSporeToTosa

首先这个是一个AccMulOp的范例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
//在MindSporeToTosa.cpp中
//这里写一个函数用来计算中间过程的输出尺度
//broadcast type
SmallVector<int64_t> broadcast(OpBuilder &rewriter, Value M1, Value M2, Location loc)
{
SmallVector<int64_t> R1dim1;
//shape
auto M1Type = M1.getType().dyn_cast<ShapedType>();
auto M2Type = M2.getType().dyn_cast<ShapedType>();
auto shape1 = M1Type.getShape();
auto shape2 = M2Type.getShape();
int64_t R1dimi;
//broadcast
if (shape1.size() == shape2.size())
{
for (unsigned int i = 0; i < shape1.size(); i++)
{

if (shape1[i] == shape2[i])
{
R1dimi = shape1[i];
}
else if (shape1[i] != 1 && shape2[i] != 1)
{
assert("error");
}
else if (shape1[i] == 1)
{
R1dimi = shape2[i];
}
else if (shape2[i] == 1)
{
R1dimi = shape1[i];
}
R1dim1.push_back(R1dimi);
}
}
return R1dim1;
};
//AccMulOp——Result = Mul(Mul(Q, K), V)
class AccMulOpConverter : public OpConversionPattern<mindspore::AccMulOp>
{

public:
using OpConversionPattern<mindspore::AccMulOp>::OpConversionPattern;
LogicalResult matchAndRewrite(mindspore::AccMulOp mindsporeOp, AccMulOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override
{
//Location
Location loc = mindsporeOp.getLoc();
//resultTypes
Type resultTypes = mindsporeOp.getResult().getType();

//Q = args[0], K = args[1], V = args[2]
ValueRange args = adaptor.getOperands();

Type elemType = mindsporeOp.getResult().getType().cast<ShapedType>().getElementType();

SmallVector<int64_t> result1 = broadcast(rewriter, args[0], args[1], loc);
ArrayRef<int64_t> newArrayRef(result1);
auto Type1 = RankedTensorType::get(newArrayRef, elemType);

//Q = args[0], K = args[1], V = args[2]
//mulOp11 = Mul(Q,K)
Value mulOp11 = rewriter.create<mlir::tosa::MulOp>(loc, Type1, args[0], args[1], 0)->getResult(0);
//Mul(mulOp11,V)
auto mulOp2 = rewriter.create<mlir::tosa::MulOp>(loc, resultTypes, mulOp11, args[2], 0);

rewriter.replaceOp(mindsporeOp, mulOp2.getResult());
return success();
}
};

再在struct ConvertMindSporeToTosaPass : public ConvertMindSporeToTosaBase<ConvertMindSporeToTosaPass>中添加patterns-AccMulOpConverter

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
struct ConvertMindSporeToTosaPass : public ConvertMindSporeToTosaBase<ConvertMindSporeToTosaPass>
{
public:
ConvertMindSporeToTosaPass() = default;
...
void runOnOperation() override
{

// clang-format off
(void)patterns.add<
...
//AccMulOp
AccMulOpConverter,
ConvertMindSporePadOp<mindspore::PadOp>
>(patterns.getContext());
mlir::populateMindSporeLowerPattern(patterns);

...
}
};

同时,在/akg-mlir/compiler/lib/Conversion/MindSporeToLinalg/MindSporeToLinalg.cpp/akg-mlir/compiler/lib/Conversion/MindSporeFinalizingLower/MindSporeFinalizingLower.cpp中对此算子添加了legal,通过手动添加legal,可以保证AccMul算子不在这里被lower,在此类转换中被保留

1
2
3
4
5
6
7
8
void runOnOperation() override
{
...

//AccmulOp
target.addLegalOp<mindspore::AccMulOp>();
...
}

然后是一个deepseek写的addlayernorm的方案

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
// MindSporeToTosa.cpp
struct ConvertAddLayerNormToTosa : public OpConversionPattern<MindSpore::AddLayerNormOp> {
using OpConversionPattern::OpConversionPattern;

LogicalResult matchAndRewrite(
MindSpore::AddLayerNormOp op,
OpAdaptor adaptor,
ConversionPatternRewriter &rewriter
) const override {
// 获取输入和属性
Value input = adaptor.getInput();
Value gamma = adaptor.getGamma();
Value beta = adaptor.getBeta();
float epsilon = op.getEpsilon().convertToFloat();

// 将AddLayerNorm分解为TOSA操作:
// 1. 计算输入均值(mean)
// 2. 计算方差(variance)
// 3. 执行归一化:(input - mean) / sqrt(variance + epsilon)
// 4. 应用gamma和beta
// 具体实现需调用TOSA的现有操作(如tosa.add, tosa.mul等)

// 伪代码示例:
auto mean = rewriter.create<tosa::ReduceMeanOp>(...);
auto variance = rewriter.create<tosa::SubOp>(...);
auto normalized = rewriter.create<tosa::DivOp>(...);
auto scaled = rewriter.create<tosa::MulOp>(normalized, gamma);
auto output = rewriter.create<tosa::AddOp>(scaled, beta);

rewriter.replaceOp(op, output);
return success();
}
};

然后要注册这个写好的pattern

1
2
3
void populateMindSporeToTosaPatterns(MLIRContext *context, RewritePatternSet &patterns) {
patterns.add<ConvertAddLayerNormToTosa>(context);
}

更新Lower流程入口确保在Lower流程中调用上述Pattern:

1
2
3
4
5
6
7
8
9
10
11
void LowerMindSporeToTosa(ModuleOp module) {
ConversionTarget target(*module.getContext());
target.addLegalDialect<Tosa::TosaDialect>();
target.addIllegalDialect<MindSpore::MindSporeDialect>();

RewritePatternSet patterns(&getContext());
populateMindSporeToTosaPatterns(patterns);
if (failed(applyPartialConversion(module, target, std::move(patterns)))) {
signalPassFailure();
}
}

MindsporetoLinalg

利用Linalg 在这一层并没有直接的矩阵乘加指令。但在
akg/akg-mlir/compiler/lib/Conversion/MindSporeToLinalg/MindSporeToLinalg.cpp
中有已经提供了写elementwise计算的方法,所以直接在
static Value createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args, ArrayRef<Type> resultTypes, PatternRewriter &rewriter)中添加对应逻辑,其中arith旨在容纳基本的整数和浮点数学运算

同上的流程

  1. 编写MindsporetoLinalg
  2. 然后编写一个pattern
  3. 还要在akg-mlir/compiler/lib/Conversion/MindSporeToTosa/MindSporeToTosa.cpp对此算子手动添加egal,保证AccMul算子不在这里被lower

但是我具体没有搞明白

任务3编写对应的测试用例(算子的info文件和对应的op_dsl)

编写测试用例的参考资料

这里老师不仅讲了怎么写测试用例,还讲了怎么测试lower过程

测试的时候首先有一个正确结果benchmark,还有一个测试数据集,然后运行测试程序就可以把算子计算结果和正确结果进行对比

Benchmark构建

用numpy验证计算的正确性

关benchmark记录在akg-mlir/python/akg_v2/utils/op_dsl.py中的get_op_dsl()函数中,每一个算子对应一个lambda函数作为dsl。以AccMul为例

1
"AccMul": lambda inputs, output, attr: accmul_str(inputs, output, attr),

对应的dsl为

1
2
3
def accmul_str(inputs, output, attr):
s="%s=np.multiply(np.multiply(%s,%s),%s)"% ( output[0]['tensor_name'], get_input(inputs[0][0]),get_input(inputs[1][0]),get_input(inputs[2][0]))
return s

这里指用np的multiply函数做计算。如果算子lower成功,期待的输出如下:

1
2
Start running xxx
xxxprecision correct

测试用例怎么写

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
{
# "input_desc"为一个list,里面包含了每个input的信息,包括data_type, format, shape和tensor_name,这个决定了融合算子的输入
"input_desc": [
[
{
"data_type": "float32",
"format": "DefaultFormat",
"shape": [
4096,
7680
],
"tensor_name": "Query"
}
],
[
{
"data_type": "float32",
"format": "DefaultFormat",
"shape": [
4096,
7680
],
"tensor_name": "Key"
}
],
[
{
"data_type": "float32",
"format": "DefaultFormat",
"shape": [
4096,
7680
],
"tensor_name": "Value"
}
]
],

# "op"为整个融合算子的名字,对应到mlir文件的func的名字。
"op": "AccMulOp_fusion",
# "output_desc"为一个list,表示融合算子中每个单算子的名字,每一项表示每个算子的特点:
# attr:没有的时候要写成null
# input_desc:算子的输入,要求同上;
# name:算子的名字,必须和mindspore td内注册的算子对应;
# output_desc:算子的输入,要求同输入;值得注意的是,如果融合算子中一个算子的输入为前序算子的输出,tensor_name必须一致以获得对应关系;
"op_desc": [
{
"attr": null,
"input_desc": [
[
{
"data_type": "float32",
"format": "DefaultFormat",
"name": "Query",
"shape": [
4096,
7680
],
"tensor_name": "Query"
}
],
[
{
"data_type": "float32",
"format": "DefaultFormat",
"name": "Key",
"shape": [
4096,
7680
],
"tensor_name": "Key"
}
],
[
{
"data_type": "float32",
"format": "DefaultFormat",
"name": "Value",
"shape": [
4096,
7680
],
"tensor_name": "Value"
}
]
],
"name": "AccMul",
"output_desc": [
{
"data_type": "float32",
"format": "DefaultFormat",
"name": "output",
"shape": [
4096,
7680
],
"tensor_name": "output"
}
]
}
],
# output_desc:整个融合算子的输出。
"output_desc": [
{
"data_type": "float32",
"format": "DefaultFormat",
"shape": [
4096,
7680
],
"tensor_name": "output"
}
],
# 后面的不管
"platform": "AKG",
"process": "cpu",
"target_info": {
"arch": "aarch64",
"feature": "neon",
"system": "linux"
},
"version": 1
}

测试文件构建完成之后,我们使用python/akg_v2/exec_tools/目录下py_benchmark.py进行测试,基本指令为

1
python ${path_to_py_benchmark}/py_benchmark.py -e cpu -f mul.info --dump_ir 1

其中

  • -e决定了后端,这里使用cpu
  • -f决定了测试文件,这里使用是mul.info
  • --dump_ir 1表示dump中间结果,相关结果保存在py_benchmark.py对应目录下