理论 + 实践,揭秘昇腾 CANN 算子开发
摘要: 本文介绍了 CANN 自定义算子开发的几种开发方式和算子的编译运行流程。然后以开发一个 DSL Add 算子为例,讲解算子开发的基本流程。
本文分享自华为云社区《昇腾CANN算子开发揭秘》,作者:昇腾 CANN 。
开发者在利用昇腾硬件进行神经网络模型训练或者推理的过程中,可能会遇到以下场景:
训练场景下,将第三方框架(例如 TensorFlow、PyTorch 等)的网络训练脚本迁移到昇腾 AI 处理器时遇到了不支持的算子。
推理场景下,将第三方框架模型(例如 TensorFlow、Caffe、ONNX 等)使用 ATC 工具转换为适配昇腾 AI 处理器的离线模型时遇到了不支持的算子。
网络调优时,发现某算子性能较低,影响网络性能,需要重新开发一个高性能算子替换性能较低的算子。
推理场景下,应用程序中的某些逻辑涉及到数学运算(例如查找最大值,进行数据类型转换等),希望通过自定义算子的方式实现这些逻辑,从而获得性能提升。
此时我们就需要考虑进行自定义算子的开发,本期我们主要带您了解 CANN 自定义算子的几种开发方式和基本开发流程,让您对 CANN 算子有宏观的了解。
一、算子基本概念
相信大家对算子的概念并不陌生,这里我们来做简单回顾。深度学习算法由一个个计算单元组成,我们称这些计算单元为算子(Operator,简称 OP)。
在网络模型中,算子对应层中的计算逻辑,例如:卷积层(Convolution Layer)是一个算子;全连接层(Fully-connected Layer, FC layer)中的权值求和过程,是一个算子。
再例如:tanh、ReLU 等,为在网络模型中被用做激活函数的算子。
二、CANN 自定义算子开发方式
学习 CANN 自定义算子开发方式之前,我们先来了解一下 CANN 算子的运行位置:包括 AI Core 和 AI CPU。
AI Core 是昇腾 AI 处理器的计算核心,负责执行矩阵、向量、标量计算密集的算子任务。
AI CPU 负责执行不适合跑在 AI Core 上的算子,是 AI Core 算子的补充,主要承担非矩阵类、逻辑比较复杂的分支密集型计算。
CANN 支持用户使用多种方式来开发自定义算子,包括 TBE DSL、TBE TIK、AICPU 三种开发方式。其中 TBE DSL、TBE TIK 算子运行在 AI Core 上,AI CPU 算子运行在 AI CPU 上。
1. 基于 TBE 开发框架的算子开发
TBE(Tensor Boost Engine:张量加速引擎)是 CANN 提供的算子开发框架,开发者可以基于此框架使用 Python 语言开发自定义算子,通过 TBE 进行算子开发有 TBE DSL、TBE TIK 两种方式。
TBE DSL(Domain-Specific Language ,基于特性域语言)开发方式
为了方便开发者进行自定义算子开发,CANN 预先提供一些常用运算的调度,封装成一个个运算接口,称为基于 TBE DSL 开发。DSL 接口已高度封装,用户仅需要使用 DSL 接口完成计算过程的表达,后续的算子调度、算子优化及编译都可通过已有的接口一键式完成,适合初级开发用户。
TBE TIK(Tensor Iterator Kernel)开发方式
TIK(Tensor Iterator Kernel)是一种基于 Python 语言的动态编程框架,呈现为一个 Python 模块,运行于
Host CPU 上。开发者可以通过调用 TIK 提供的 API 基于 Python 语言编写自定义算子,TIK 编译器会将其编译为昇腾 AI 处理器应用程序的二进制文件。
TIK 需要用户手工控制数据搬运和计算流程,入门较高,但开发方式比较灵活,能够充分挖掘硬件能力,在性能上有一定的优势。
2. AI CPU 算子开发方式
AI CPU 算子的开发接口即为原生 C++接口,具备 C++程序开发能力的开发者能够较容易的开发出 AI CPU 算子。AI CPU 算子在 AI CPU 上运行。
下面的开发方式一览表,对上述几种开发方式作对比说明,您可以根据各种开发方式的适用场景选择适合的开发方式。
三、CANN 算子编译运行
算子构成
一个完整的 CANN 算子包含四部分:算子原型定义、对应开源框架的算子适配插件、算子信息库和算子实现。这四个组成部分会在算子编译运行的过程中使用。
算子编译
推理场景下,进行模型推理前,我们需要使用 ATC 模型转换工具将原始网络模型转换为适配昇腾 AI 处理器的离线模型,该过程中会对网络中的算子进行编译。
训练场景下,当我们跑训练脚本时,CANN 内部实现逻辑会先将开源框架网络模型下发给 Graph Engine 进行图编译,该过程中会对网络中的算子进行编译。
CANN 算子的编译逻辑架构如下:
具体的 CANN 算子编译流程如下,在编译流程中会用到上文提到的算子的四个组成部分。
Graph Engine 调用算子插件,将原始网络模型中的算子映射为适配昇腾 AI 处理器的算子,从而将原始开源框架图解析为适配昇腾 AI 处理器的图。
调用算子原型库校验接口进行基本参数的校验,校验通过后,会根据原型库中的推导函数推导每个节点的输出 shape 与 dtype,进行输出 tensor 的静态内存的分配。
Graph Engine 根据图中数据将图拆分为子图并下发给 FE。FE 在处理过程中根据算子信息库中算子信息找到算子实现,将其编译成算子 kernel,最后将优化后子图返回给 Graph Engine。
Graph Engine 进行图编译,包含内存分配、流资源分配等,并向 FE 发送 tasking 请求,FE 返回算子的 taskinfo 信息给 Graph Engine,图编译完成后生成适配昇腾 AI 处理器的模型。
算子运行
推理场景下,使用 ATC 模型转换工具将原始网络模型转换为适配昇腾 AI 处理器的离线模型后,开发 AscendCL 应用程序,加载转换好的离线模型文件进行模型推理,该过程中会进行算子的调用执行。
训练场景下,当我们跑训练脚本时,内部实现逻辑将开源框架网络模型下发给 Graph Engine 进行图编译后,后续的训练流程会进行算子的调用执行。
CANN 算子的运行逻辑架构如下:
具体流程如下,首先 Graph Engine 下发算子执行请求给 Runtime,然后 Runtime 会判断算子的 Task 类型,若是 TBE 算子,则将算子执行请求下发到 AI Core 上执行;若是 AI CPU 算子,则将算子执行请求下发到 AI CPU 上执行。
四、算子开发流程
本章节以通过 DSL 开发方式开发一个 Add 算子为例,带您快速体验 CANN 算子开发的流程。流程图如下:
算子开发准备
环境准备:准备算子开发及运行验证所依赖的开发环境与运行环境。工程创建:创建算子开发工程,有以下几种实现方式:
基于 MindStudio 工具进行算子开发,直接使用 MindStudio 工具创建算子工程,会自动生成算子工程及代码模板。
基于 msopgen 工具进行开发,会自动生成算子工程及代码模板。
基于自定义算子样例工程进行开发,开发者需要自己创建算子相关实现文件,或者基于已有样例进行修改。
下面以 msopgen 工具创建算子开发工程为例进行介绍:
定义 AddDSL 算子的原型定义 json 文件,用于生成 AddDSL 的算子开发工程。例如,定义的 json 文件的名字为 add_dsl.json,存储路径为:$HOME/sample,文件内容如下:
使用 msopgen 工具生成 AddDSL 算子的开发工程。
“$HOME/Ascend”为 CANN 软件安装目录;
“-f tf”参数代表选择的原始框架为 TensorFlow;
“ai_core-<soc_version>”代表算子在 AI Core 上运行,<soc_version>为昇腾 AI 处理器的型号。
此命令执行完后,会在 $HOME/sample/AddDsl 目录下生成算子工程,工程中包含各交付件的模板文件,编译脚本等,如下所示:
算子开发过程
实现 AddDSL 算子的原型定义。
算子原型定义文件包含算子注册代码的头文件(*.h)以及实现基本校验、Shape 推导的实现文件(*.cc)。
msopgen 工具根据 add_dsl.json 文件在“op_proto/add_dsl.h”中生成了算子注册代码,开发者需要检查自动生成的代码逻辑是否正确,一般无需修改。
改“op_proto/add_dsl.cc”文件,实现算子的输出描述推导函数及校验函数。
在 IMPLEMT_COMMON_INFERFUNC(AddDSLInferShape)函数中,填充推导输出描述的代码,针对 AddDSL 算子,输出 Tensor 的描述信息与输入 Tensor 的描述信息相同,所以直接将任意一个输入 Tensor 的描述赋給输出 Tensor 即可。
在 IMPLEMT_VERIFIER(AddDSL, AddDSLVerify)函数中,填充算子参数校验代码。
实现 AddDSL 算子的计算逻辑。
“tbe/impl/add_dsl.py”文件中已经自动生成了算子代码的框架,开发者需要在此文件中修改 add_dsl_compute 函数,实现此算子的计算逻辑。
add_dsl_compute 函数的实现代码如下:
@register_op_compute("add_dsl")def add_dsl_compute(x1, x2, y, kernel_name="add_dsl"):
调用 dsl 的 vadd 计算接口
return res 配置算子信息库。
算子信息库的路径为“tbe/op_info_cfg/ai_core/<soc_version>/add_dsl.ini”,包含了算子的类型,输入输出的名称、数据类型、数据排布格式等信息,msopgen 工具已经根据 add_dsl.json 文件将上述内容自动填充,开发者无需修改。
AddDSL 算子的信息库如下:
实现算子适配插件。
算子适配插件实现文件的路径为“framework/tf_plugin/tensorflow_add_dsl_plugin.cc”,针对原始框架为 TensorFlow 的算子,CANN 提供了自动解析映射接口“AutoMappingByOpFn”,如下所示:
以上为工程自动生成的代码,开发者仅需要修改.OriginOpType("AddDSL")中的算子类型即可。此处我们仅展示算子开发流程,不涉及原始模型,我们不做任何修改。
至此,AddDSL 算子的所有交付件都已开发完毕。
算子工程编译及算子包部署
算子开发过程完成后,需要编译自定义算子工程,生成自定义算子安装包并进行自定义算子包的安装,将自定义算子部署到算子库。
算子工程编译
1. 修改 build.sh 脚本,配置算子编译所需环境变量。
将 build.sh 中环境变量 ASCEND_TENSOR_COMPILER_INCLUDE 配置为 CANN 软件头文件所在路径。
修改前,环境变量配置的原有代码行如下:
修改后,新的代码行如下:
${INSTALL_DIR}请替换为 CANN 软件安装后文件存储路径。例如,若安装的 Ascend-cann-toolkit 软件包,则安装后文件存储路径为:$HOME/Ascend/ascend-toolkit/latest。
2. 在算子工程目录下执行如下命令,进行算子工程编译。
编译成功后,会在当前目录下创建 build_out 目录,并在 build_out 目录下生成自定义算子安装包 custom_opp_<target os>_<target architecture>.run。
自定义算子安装包部署
以运行用户执行如下命令,安装自定义算子包。
命令执行成功后,自定义算子包中的相关文件部署到 CANN 算子库中。
算子运行验证
算子包部署完成后,可以进行 ST 测试(System Test)和网络测试,对算子进行运行验证。
ST 测试
ST 测试的主要功能是:基于算子测试用例定义文件*.json 生成单算子的 om 文件;使用 AscendCL 接口加载并执行单算子 om 文件,验证算子执行结果的正确性。ST 测试会覆盖算子实现文件,算子原型定义与算子信息库,不会对算子适配插件进行测试。
网络测试
你可以将算子加载到网络模型中进行整网的推理验证,验证自定义算子在网络中运行结果是否正确。网络测试会覆盖算子开发的所有交付件,包含实现文件,算子原型定义、算子信息库以及算子适配插件。
具体的验证过程请参考“昇腾文档中心”。
以上就是 CANN 自定义算子开发的相关知识点,您也可以在“昇腾社区在线课程”板块学习视频课程,学习过程中的任何疑问,都可以在“昇腾论坛”互动交流!
相关参考:
[1]昇腾文档中心
[2]昇腾社区在线课程
[3]昇腾论坛
版权声明: 本文为 InfoQ 作者【华为云开发者联盟】的原创文章。
原文链接:【http://xie.infoq.cn/article/ef034a7c87edf058e3a05c9fd】。文章转载请联系作者。
评论