[cann document express issue 06] first knowledge of tbe DSL operator development

Shengteng developers can be based on TBE（Tensor Boost Engine） Provided python Interface for user-defined operator development ,TBE Operator development has DSL （ Domain-Specific Language ） and TIK （ Tensor Iterator Kernel ） Two ways , The applicable scenarios, advantages and disadvantages of the two methods are as follows ：

In this issue, we mainly introduce the easy to use DSL Operator development mode .

DSL Functional framework

1. Developers call DSL Interface to describe the computing logic , Indicate the calculation method and steps of the operator .
2. After the development of computing logic , Developers can call Auto Schedule Interface starts automatic scheduling , During automatic dispatching TBE Automatically select the appropriate scheduling template according to the calculation type , Complete data segmentation and data flow division , Ensure optimal hardware execution . After the dispatch , Will generate something similar to TVM Of IR（Intermediate Representation） In the middle .
3. IR After generation ,Pass The generated IR Compile and optimize , The optimized way is double buffering （Double Buffer）、 Assembly line （Pipeline） Sync 、 Memory allocation management 、 Instruction mapping, etc .
4. Operator longitude Pass After processing , Will be automatically by CodeGen Generating classes C Temporary files of code , This temporary file generates the binary file of the operator through the compiler , It can be directly loaded and called by the network model .

DSL Operator code implementation

Before implementing the operator code, you need to analyze the mathematical expression of the operator , Input 、 Output , Specify what needs to be called DSL Interface , Then the operator implementation code is developed . be based on DSL The code implementation process of is shown in the following figure ：

The code structure of the operator implementation is as follows ：

Let's implement two inputs shape identical 、 The data type is “float32” Of Add For example, operators , about DSL Operator implementation of the code flow .

1.   First, introduce the development time dependency Python modular

frequently-used Python The modules are as follows ：

• “tbe.dsl”： contain TBE DSL Computing interface of 、 Scheduling interface and compilation interface
• “tbe.tvm”：TBE Is based on TVM Framework extension , Developers can use when implementing operators TVM Interface
• “tbe.common.utils.para_check”：TBE Operator parameter verification interface provided
• “tbe.common.utils.shape_util”：TBE Operator provided shape Processing interface

2.   Declare operator interface

The operator interface definition function contains the input and output information of the operator and the kernel name .

Here is Add The definition of operator ：

In the above example ,add Of the operator type,input_x、input_y Is the input and output of the operator tensor, Define in the form of a dictionary , contain shape、ori_shape、format、ori_format And dtype Information ,kernel_name Is the name of the operator in the kernel , And operator type Keep consistent .

Developers can use when defining operator interface functions TBE Provided parameter verification interface check_input_type Check whether the parameter type of the operator is legal ,check_input_type For decorator function , The usage method is shown below ：

Of course , You can also customize the verification function of related parameters , Basic parameter checking helps to find problems ahead of time in the operator compilation phase .

3.   For input tensor To carry on the placeholder

Get the input data shape、dtype（ This example is float32 Fixed data type of ）, Use TVM Of placeholder Interface pair input tensor To carry on the placeholder , Return to one tensor object , The data in this location is specified only when the program is running .

4.   Implementation of computing logic

The computational logic of an operator can be obtained by TBE Of DSL Computing interface implementation , for example Add Operators can be passed through DSL Of vadd Interface implementation input_x And input_y The addition operation of .

5.   Scheduling and compiling

After the implementation of computing logic , Need to call auto_schedule Interface , Automatically generate corresponding scheduling ; And then call build Interface to compile operators , Compile the operator specific kernel .

among config Configured for compilation parameters map, The configuration information includes whether to print IR、 Operator kernel name and input 、 Output tensor .
Congratulations , So far, you have completed Add Operator implementation code development .

More about

The above is only for DSL The key code of operator development is briefly introduced , For more details and skills of operator implementation, please log in to shengteng community , Read the relevant documents ：https://www.hiascend.com/

Rise CANN The document center is dedicated to providing developers with better content and more convenient development experience , help CANN Developers co build AI ecology . Any comments and suggestions can be fed back in shengteng community , Your every attention is our driving force .