TBE DSL Introduction

Concepts

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 community project. It aims to further abstract the generation rules of operators by dividing the operators into operation primitives and combining them when necessary. According to the definition of the computation process of operators, the TVM uses Schedule and CodeGen technologies to generate the operators for the specified hardware.

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

Currently, TBE DSL APIs mainly include math compute APIs, Cast compute APIs, NN compute APIs, and build APIs.

API Overview

You can view the API definition in the python/site-packages/tbe/dsl/api.py under the CANN component directory.

Table 1 TBE DSL APIs

API Category

Description

Math

Math compute APIs.

Cast

Rounds the input tensor element-wise based on rules.

NN

Computes NN-related operations.

Tensor

Reconcatenates or splits input tensors.

Build APIs

Build and schedule APIs.

API Usage

Configure the basic CANN environment variables by referring to Environment Setup before using TBE DSL APIs.