Mindspore实习-AKG SIG算子addlayernorm编辑和合并
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 之中。
【需求描述】
- 在MindSporeOps.td中添加对应算子的实现
- 打通算子的Lower流程(直接lower到Linalg或者lower到TOSA),相关Dialect的代码(如MindSporeToTosa.cpp或者MindSporeToLinalg.cpp)
- 提供对应的测试用例(算子的info文件和对应的op_dsl)
【参考资料】
- 样例代码仓:https://gitee.com/mindspore/akg/pulls/989
- 代码添加教程wiki:https://gitee.com/monkeykingd/akg/wikis/AKG-MLIR算子添加用例
- 测试用例教程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代码但是详细细节写了个寂寞,然后全是独立构建方法
从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 | # 以上依赖的安装过程 |
1 | # 编译安装mindspore 2.3.1 |
最终的解决方案mindspore的whl安装+AKG源码安装
ok fine ,最后终于成功了,整理一下思路,就是说从官网只下载AKG的源码,当然mindspore的仓库里有AKG这个子仓库,然后我用官方编译好的whl进行安装,安装版本为2.3.1,这个安装没有出现问题
1 | conda create -n mindspore39NLP python=3.9 |
AddLayerNorm算子的定义
首先addlayernorm的公式是怎样的,看一下pytorch没找到,但是找到了升腾的一个方案。文献没有找到,但是在 Transformer 架构中,残差连接(Add)与层归一化(LayerNorm)常被结合使用,是将输入的X1和X2进行逐元素相加,然后对结果进行layernorm。LayerNorm 最初由 Jimmy Lei Ba 等人在 2016 年的论文 《Layer Normalization》
任务1在MindSporeOps.td中添加对应算子的实现
开头已经了解了mindspore AKG是什么了,里面提供了基本算子的定义和Lower流程,但现在需要添加新的复合算子AddLayerNorm到MindSpore Dialect中。
这里需要会LLVM编译器的相关操作,所谓td文件就是其中的编辑文件
- TableGen是LLVM生态中用于生成代码的声明式语言,通过
.td
文件定义数据模型(如操作、指令、寄存器)。 - 在MLIR中:td文件用于定义Dialect中的算子(Op)、类型(Type)、属性(Attribute)等,生成C++代码框架。
如果要写一个addLaynorm,可能如下
1 | //有三个输入,一个输出,所以在这里定义了四个MindSpore_Tensor类型的tensor,对应五个输入X1、X2、gamma、beta和一个输出y。但是升腾的方案写的很复杂,还包括了是否要输出中间的相加结果 |
ins后接输入参数,out后接输出参数。在编译后会自动生成get方法来获取输入输出,并将蛇形名称转换为驼峰式命名法。在akg-mlir/build/include/akg/Dialect/MindSpore/IR/MindSporeOps.cpp.inc
可以找到编译后的get方法
比如:
1 | ::mlir::TypedValue<::mlir::TensorType> AccMulOp::getInputA() { |
任务2打通算子的Lower流程
Lower流程讲解
直接lower到Linalg或者lower到TOSA,相关Dialect的代码(如MindSporeToTosa.cpp或者MindSporeToLinalg.cpp)
在MLIR(Multi-Level Intermediate Representation)生态中,Dialect、Linalg 和 TOSA 是不同层次的中间表示(IR),用于描述和优化计算任务。
-
Dialect:是 MLIR 中的用于定义特定领域或抽象层次的中间表示。每个 Dialect 包含一组操作(Op)、类型(Type)和属性(Attribute),用于描述特定领域的计算任务。
例如:
- MindSpore Dialect:描述 MindSpore 框架中的算子(如
Add
、LayerNorm
)。 - TOSA Dialect:描述面向硬件后端的张量运算操作。
- Linalg Dialect:描述线性代数操作(如矩阵乘法、卷积)。
- MindSpore Dialect:描述 MindSpore 框架中的算子(如
-
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.add
、linalg.mul
等操作。 - MindSpore Dialect 到 TOSA:将 MindSpore 的算子转换为 TOSA 的标准化操作。将
AddLayerNorm
分解为tosa.add
、tosa.mul
等操作。
MindSpore的Lower过程这里有两条路,可以从MindSpore转到Tosa再转到Linalg,也可以直接从MindSpore转到Linalg
- MindSpore Dialect 到 Linalg:将 MindSpore 的算子(如
编写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 | this->mindOpFactory["addlayernorm"] = &MindBuilder::convertTernaryOp<mindspore::AddLayerNormOp>; |
MindSporeToTosa
首先这个是一个AccMulOp的范例
1 | //在MindSporeToTosa.cpp中 |
再在struct ConvertMindSporeToTosaPass : public ConvertMindSporeToTosaBase<ConvertMindSporeToTosaPass>
中添加patterns-AccMulOpConverter
1 | struct ConvertMindSporeToTosaPass : public ConvertMindSporeToTosaBase<ConvertMindSporeToTosaPass> |
同时,在/akg-mlir/compiler/lib/Conversion/MindSporeToLinalg/MindSporeToLinalg.cpp
和/akg-mlir/compiler/lib/Conversion/MindSporeFinalizingLower/MindSporeFinalizingLower.cpp
中对此算子添加了legal,通过手动添加legal,可以保证AccMul算子不在这里被lower,在此类转换中被保留
1 | void runOnOperation() override |
然后是一个deepseek写的addlayernorm的方案
1 | // MindSporeToTosa.cpp |
然后要注册这个写好的pattern
1 | void populateMindSporeToTosaPatterns(MLIRContext *context, RewritePatternSet &patterns) { |
更新Lower流程入口确保在Lower流程中调用上述Pattern:
1 | void LowerMindSporeToTosa(ModuleOp module) { |
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旨在容纳基本的整数和浮点数学运算
同上的流程
- 编写MindsporetoLinalg
- 然后编写一个pattern
- 还要在
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 | def accmul_str(inputs, output, attr): |
这里指用np的multiply函数做计算。如果算子lower成功,期待的输出如下:
1 | Start running xxx |
测试用例怎么写
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
对应目录下