更新时间:2021-03-18 GMT+08:00
分享

TBE简介

概述

TBE(Tensor Boost Engine)是基于TVM(Tensor Virtual Machine)的自定义算子开发框架。TVM是社区的开源项目,旨在将各算子的生成规则进一步抽象,将算子本身分成各个操作原语,在需要的时候加以组合。TVM会根据算子的计算过程的定义,使用Schedule技术和Codegen技术,生成对指定硬件的算子。

由于Schedule是描述在硬件上实现一个算子的计算过程,这需要较强的硬件知识。为了简化用户书写算子的难度,我们在TVM的基础上,简化了书写Schedule的难度,采用“Auto schedule”的概念,提供了一组Tensor Boost Engine API,来组合出算子的计算。用户通过使用API进行适当的组合定义一个算子的计算过程,把Schedule交给Auto schedule去完成。本文介绍在TVM基础上定义的TBE DSL API,用户可通过这些API来开发自己的算子。

TBE DSL API目前主要涵盖向量运算,包括Element-wise类操作接口、Reduction 操作接口、Broadcast操作接口、Index操作接口、Concat操作接口、卷积接口、4D/5D互转接口、矩阵计算接口等。

版本查询

你可以在ATC安装目录下的“python/site-packages/te/te/version.py”中查看当前TBE DSL的版本号。

接口使用说明

使用TBE DSL接口时需要声明环境变量PYTHONPATH:

export install_path=/home/HwHiAiUser/Ascend/ascend-toolkit/latest         
exprot PYTHONPATH=${install_path}/atc/python/site-packages/te:${install_path}/atc/python/site-packages/topi

其中install_path为ATC的安装路径。

分享:

    相关文档

    相关产品

close