TBE DSL简介

概述

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目前主要涵盖向量运算，包括Math、NN、Reduce、卷积、矩阵计算等接口。

接口总览

您可以在CANN软件安装后文件存储路径下的“python/site-packages/tbe/dsl/api.py”文件中查看接口定义。

表1 TBE DSL接口总览
接口分类	简介
Math	对Tensor中每个原子值分别做相同操作的计算接口。
NN	神经网络相关计算接口。
Cast	取整计算接口，对输入Tensor中的每个元素按照一定的规则进行取整操作。
Inplace	对Tensor进行按行相关计算。
Reduce	对Tensor按轴进行相关操作的计算接口。
Matmul	矩阵乘计算。
Gemm	通用矩阵乘计算接口。
卷积	包含2D和3D卷积运算及其反向运算的相关接口。
Pooling2d	2D池化接口。
Pooling3d	3D池化接口。
Array	在指定轴上对输入Tensor进行重新连接、切分或者排序的接口。
auto schedule	调度接口。
build	编译接口。

接口使用说明

使用TBE DSL接口前需要参见环境准备完成CANN基础环境变量的配置。

父主题： TBE DSL API