PyTorch框架使用DSL进行TBE算子开发全流程

PyTorch框架使用DSL进行TBE算子开发全流程

bilibili视频链接

PyTorch框架使用DSL进行TBE算子开发全流程

1. DSL算子基本概念介绍

1.1 什么是算子

深度学习算法由一个个计算单元组成，我们称这些计算单元为算子（Operator，简称 OP）。在网络模型中，算子对应层中的计算逻辑，例如：卷积层（Convolution Layer）是一个算子；全连接层（Fully-connected Layer， FC layer）中的权值求和过 程，是一个算子。
对每一个独立的算子，用户需要编写算子描述文件，描述算子的整体逻辑、计算步骤以及相关硬件平台信息等。然后用深度学习编译器对算子描述文件进行编译，生成可在特定硬件平台上运行的二进制文件后，将待处理数据作为输入，运行算子即可得到期望输出。将神经网络所有被拆分后的算子都按照上述过程处理后， 再按照输入输出关系串联起来即可得到整网运行结果。

图1：算子全流程执行

1.2 什么是TBE算子

TBE（Tensor Boost Engine）提供了基于开源深度学习编译栈TVM框架的自定义算子开发能力，通过TBE提供的API可以完成相应神经网络算子的开发。

一个完整的TBE算子包含四部分：算子原型定义、对应开源框架的算子适配插件、算子信息库定义和算子实现。

图2：TBE算子全流程执行

昇腾AI软件栈提供了TBE（Tensor Boost Engine：张量加速引擎）算子开发框架，开发者可以基于此框架使用Python语言开发自定义算子，通过TBE进行算子开发有以下几种方式：

• DSL（Domain-Specific Language）开发
• TIK（Tensor Iterator Kernel）开发

1.3 什么是DSL算子

为了方便开发者进行自定义算子开发，TBE（Tensor Boost Engine）提供了一套计算接口供开发者用于组装算子的计算逻辑，使得70%以上的算子可以基于这些接口进行开发，极大的降低自定义算子的开发难度。TBE提供的这套计算接口，称之为DSL（Domain-Specific Language）。基于DSL开发的算子，可以直接使用TBE提供的Auto Schedule机制，自动完成调度过程，省去最复杂的调度编写过程。

图3：DSL算子介绍

1.3.1 DSL功能框架

1. 开发者调用TBE提供的DSL接口进行计算逻辑的描述，指明算子的计算方法和步骤。
2. 计算逻辑开发完成后，开发者可调用Auto Schedule接口，启动自动调度，自动调度时TBE会根据计算类型自动选择合适的调度模板，完成数据切块和数据流向的划分，确保在硬件执行上达到最优。
3. Auto Schedule调度完成后，会生成类似于TVM的I R（Intermediate Representation）的中间表示。
4. 编译优化（Pass）会对算子生成的IR进行编译优化，优化的方式有双缓冲（Double Buffer）、流水线（Pipeline）同步、内存分配管理、指令映射、分块适配矩阵计算单元等。
5. 算子经Pass处理后，由CodeGen生成类C代码的临时文件，这个临时代码文件可通过编译器生成算子的实现文件，可被网络模型直接加载调用。

1.3.2 DSL计算接口

TBE DSL提供的计算接口主要涵盖向量运算，包括Math、NN、Reduce、卷积、矩阵计算等接口。

2. 算子开发流程介绍

算子开发通常包括以下步骤：

• 环境准备：配置MindStudio算子开发环境
• 算子分析：明确算子的功能及数学表达式，选择算子开发方式
• 工程创建：利用MindStudio创建算子工程
• 算子原型定义：规定在昇腾AI处理器上可运行算子的约束
• 算子代码实现：使用DSL或TIK实现算子
• 算子信息库定义：规定算子在昇腾AI处理器上的具体实现规格
• 算子UT测试：测试算子代码的正确性
• 算子工程编译：编译出可直接安装的自定义算子run包
• 算子工程部署：将自定义算子部署到opp算子库
• 算子ST测试：在真实的硬件环境中，使用离线模型验证算子功能的正确性
  

图4：算子开发流程介绍

3.算子分析

进行算子开发前，开发者应首先进行算子分析，算子分析包含：明确算子的功能及数学表达式，选择算子开发方式（DSL方式或者TIK方式），最后细化并明确算子规格，分析流程如下所示：

图5：TBE开发流程介绍

3.1分析算子算法原理，提取算子的数学表达式

当$x={\left[x_1,x_2,\cdots ,x_n\right]}^{\mathrm{T}}$时那么$x$的$ p$范数为
$$\Vert x{\Vert }_p={\left({\left|x_1\right|}^p+{\left|x_2\right|}^p+\cdots +{\left|x_n\right|}^p\right)}^{\frac{1}{p}}$$
当$p$取0时，表示向量中非零元素的个数（即为其稀疏度）
当$p$取1，2，$\infty$时分别是以下几种最简单的情形
$$\Vert x{\Vert }_1=\left|x_1\right|+\left|x_2\right|+\dots +\left|x_n\right|$$