【CANN文档速递06期】初识TBE DSL算子开发

昇腾CANN 发表于 2022/06/24 17:17:10 2022/06/24
【摘要】 本期我们主要介绍较容易上手的DSL算子开发方式。

6.png

昇腾开发者可以基于TBE(Tensor Boost Engine)提供的python接口进行自定义算子开发,TBE算子开发有DSL ( Domain-Specific Language )和TIK ( Tensor Iterator Kernel )两种方式,两种方式的适用场景及优缺点如下所示:


本期我们主要介绍较容易上手的DSL算子开发方式。

DSL功能框架


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

DSL算子代码实现

在实现算子代码之前需要分析算子的数学表达式,输入、输出,明确需要调用的DSL接口,然后进行算子实现代码的开发。基于DSL的代码实现流程如下图所示:


算子实现的代码结构如下所示:


下面我们以实现两输入shape相同、数据类型为“float32”的Add算子为例,讲述DSL算子实现的代码流程。

1.  首先引入开发时依赖的Python模块

常用的Python模块如下所示:

  • “tbe.dsl”:包含TBE DSL的计算接口、调度接口以及编译接口
  • “tbe.tvm”:TBE是基于TVM框架扩展而来的,开发者在实现算子的时候可以使用TVM接口
  • “tbe.common.utils.para_check”:TBE提供的算子参数校验接口
  • “tbe.common.utils.shape_util”:TBE提供的算子shape处理接口

2.  声明算子接口

算子接口定义函数中包含算子的输入输出信息以及内核名称。

如下所示为Add算子的定义:

上述示例中,add为算子的type,input_x、input_y为算子的输入输出tensor,采用字典的形式定义,包含shape、ori_shape、format、ori_format与dtype信息,kernel_name为算子在内核中的名称,与算子type保持一致即可。

开发者在定义算子接口函数时可以使用TBE提供的参数校验接口check_input_type校验算子的参数类型是否合法,check_input_type为装饰器函数,使用方法如下所示:


当然,您也可以自定义实现相关参数的校验功能,基本的参数校验有助于在算子编译阶段提前发现问题。

3.  对输入tensor进行占位


获取输入数据的shape、dtype(此示例为float32的固定数据类型),使用TVM的placeholder接口对输入tensor进行占位,返回一个tensor对象,此位置中的数据在程序运行时才被指定。

4.  进行计算逻辑的实现

算子的计算逻辑可以通过TBE的DSL计算接口实现,例如Add算子可以通过DSL的vadd接口实现input_x与input_y的相加操作。


5.  调度与编译

计算逻辑实现完成后,需要调用auto_schedule接口,自动生成相应的调度;然后调用build接口进行算子的编译,编译出算子专用内核。


其中config为编译参数配置的map,配置信息包括是否需要打印IR、算子内核名称以及输入、输出张量。
恭喜您,至此您已经完成了Add算子实现代码的开发。

更多介绍

以上仅对DSL算子开发的关键代码进行了简要介绍,更多算子实现时的细节及技巧可登录昇腾社区,阅读相关文档:https://www.hiascend.com/


昇腾CANN文档中心致力于为开发者提供更优质的内容和更便捷的开发体验,助力CANN开发者共建AI生态。任何意见和建议都可以在昇腾社区反馈,您的每一份关注都是我们前进的动力。

【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区),文章链接,文章作者等基本信息,否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件至:cloudbbs@huaweicloud.com进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容。
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。