TBE DSL Introduction

TBE Overview

Tensor Boost Engine, abbreviated as TBE, is a framework for developing custom operators based on the Tensor Virtual Machine (TVM). The TVM is an open-source project of the community. It aims to further abstract the generation rules of operators by dividing the operators into operation primitives and combining the operators when necessary. According to the definition of the computation process of operators, the TVM uses the Schedule technology and the CodeGen technology to generate the operators for the specified hardware.

Schedule is used to describe the computation process for implementing an operator on hardware, which requires profound hardware knowledge. To reduce the difficulty in writing operators, the writing of Schedule is simplified based on the TVM. Guided by the concept of "Auto_Schedule", a collection of TBE APIs are provided to combine the computation of operators. By combination using APIs, you can define the computation process of an operator and hand over the schedule to "Auto_Schedule". This document describes the TBE domain-specific language (DSL) APIs defined based on the TVM. You can use these APIs to develop operators.

The TBE DSL APIs mainly cover vector operations, including element-wise operation APIs, reduction APIs, broadcast APIs, index operation APIs, concat APIs, convolution APIs, 4D to 5D conversion APIs, and matrix computation APIs.

Version Query

You can view the version number of the current TBE-DSL in the python/site-packages/te/version.py directory in the ATC installation path.

API Usage

Before calling TBE DSL APIs, declare the environment variable PYTHONPATH.

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

Where, install_path indicates the ATC installation path.