TBE实战算子开发—DSL方式

前言

TBE算子详解

Hello World

• 目标：
  使用DSL语言实现平方根功能的TBE算子。
• 接口命名：
  sqrt_cce()。
• 算子分析：
  Sqrt算子功能是对Tensor中每个原子值求开方，数学表达式子为y=√ 。
  根据当前TBE框架可支持的计算描述API，采用如下公式来表达Sqrt算子的计算过程：
  

TBE算子代码结构

• TBE算子入参

def sqrt_cce  (shape, dtype, kernel_name="sqrt_cce", need_build=False, need_print=False):

shape : Tensor的属性，表示Tensor的形状，用list或tuple类型表示。
例如（3, 2, 3）、（4, 10）。
- dtype : Tensor的数据类型，用字符串类型表示。
例如“float32”、“float16”、“int8”等。

• TBE算子输入占位符
  样例：data = tvm.placeholder(shape, name=“data”, dtype=dtype)
  tvm.placeholder()是tvm框架的API，用来为算子执行时接收的数据占位，通俗理解与C语言中%d、%s一样，返回的是一个Tensor对象，上例中使用data表示；入参为shape，name，dtype,是为Tensor对象的属性。
  这里的输入是指算子执行时的输入数据，与编译时期入参不同，编译时期入参（shape,type）是为了得到算子执行文件的入参。
• TBE算子中定义计算过程
  定义计算过程是指使用DSL语言，根据数学算式，描述出实现算子功能的计算步骤，以算子sqrt为例：
  
  根据sqrt算子的函数表达式，描述sqrt的计算过程如上所示，其中te.lang.cce.name()皆为TBE框架的API.
  DSL语言简化了计算过程描述
  
• TBE算子中调度（schedule）操作
  样例：sch = generic.auto_schedule(res)
  计算过程描述完之后，就会做调度；调度是与硬件相关的，功能主要是调整计算过程的逻辑，意图优化计算过程，使计算过程更高效，以及保证计算过程中占用硬件存储空间不会超过上限。
  
• TBE算子中的构建操作
  样例 ： te.lang.cce.cce_build_code(sch, config)
  TBE框架提供了cce_build_code()API，传入schedule以及相关的配置项，即可完成编译，生成最终硬件可执行文件。

TBE框架ComputeAPI

• TBE框架提供描述计算过程的API：
  TBE算子都是调用框架提供的computeAPI来描述计算过程，接口皆为te.lang.cce.name的形式；compute API 现根据功能类型可分为以下几类：
  
• elewise_compute ：
  对Tensor中每个原子值分别做相同操作，elewise_compute分为singleElewise 和 binaryElewise以及multipleElewise; singleElewise是单独对一个Tensor中每个值做操作，如te.lang.cce.vabs 即对每个数值x求绝对值：
  binaryElewise是输入两个shape相同的tensor,对应位置上的数值做操作，如te.lang.cce.vadd:
  
  multipleElewise的API与binaryElewiseAPI的区别在于接收多个Tensor,计算逻辑依然是对应位置上的数值做计算。
• cast_compute：
  cast_computeAPI的功能是对Tensor做类型转换，满足一些计算过程描述的需要，比如样例中，把“float16”的Tensor转换为“float32”的Tensor,用以保证精度的要求。
• segment_compute
  对Tensor进行分段操作，举例如下：
  
• reduction_compute
  对Tensor按轴进行操作，例如te.lang.cce.sum(data,axis)表示对data按axis进行累加,举例如下：
  
• broadcast_compute
  对Tensor按照目标shape进行广播，如shape为（2,1,2）的Tensor广播成(2,2,2)的Tensor ，举例如下：
  
• 除以上经常被使用的API，还有一些computeAPI如下所示：
  mmad_compute 支持矩阵乘法
  卷积相关compute 一些专门针对卷积的computeAPI

TBE算子编译过程

• TBE算子编译与应用
  
  TBE算子经过编译，生成能够在硬件上运行的xxx.o形式的可执行文件，此为TBE算子在TBE框架下执行后的输出件。
• TBE算子编译过程
  TBE算子编译过程分为DSL->Schedule->pass->codegen四步
  • Schedule
    经过Schedule，计算过程描述逻辑发生了转变，针对硬件的存储上限做了切分、合并等操作。
  • Pass
    Pass阶段会对计算过程描述进行指令替换，将数学方式表示的计算描述，映射为硬件可以读懂的指令，
    且会对指令进行优化，以获得更高效的性能，经过pass后，计算过程变为IR表示。
  • Codegen
    Codegen是TBE框架执行的最后一步，将Pass产生的IR构建为cce代码，进而经过compile生成二进制的可执行文件。

以计算过程中的一条DSL语句举例，
te.lang.cce.vexp(mul_val)，含义为对Tensor
（mul_val）求e^(x)
1.经过TBE框架的schedule过程，DSL描述转变为
Produce对象来表示计算过程，此时具有数据表示
含义；
2.经过pass层，produce对象中数学算式替换为指令
描述；
3.经过build cce 过程，produce 对象表示的计算过程
完全转变为cce语句；
4.最后经过compile,cce代码编译输出二进制可执行文
件。

总结

• TBE算子代码结构
  主要部分：TBE算子代码结构和开发方式，TBE框架提供的computeAPI
• TBE算子编译过程：
  编译期输入输出：TBE框架输出件是二进制的可执行文件。
  框架编译过程：DSL->Schedule->Pass->Codegen。