Auto Schedule

To build an operator based on the TBE DSL is to use the fused DSL API to express the compute logic of the operator, and then call the Auto Schedule function to automatically schedule the operator to complete data tiling and divide data flow directions. The Auto Schedule mechanism is the default schedule optimization mechanism at the low level of the TBE. Developers cannot control the mechanism during operator development.

The following is an example of developing an operator based on the DSL to obtain the exponent of x, accumulate and reduce along axis 0, and then obtain the reciprocal.

x = tvm.placeholder((512, 1024), "float16")
exp_x = te.lang.cce.vexp(x)
reduce_exp_x = te.lang.cce.sum(exp_x, axis = 0)
res = te.lang.cce.vrec(reduce_exp_x)

with tvm.target.cce():
    sch = topi.generic.auto_schedule(res)

Developers call the topi.generic.auto_schedule() API to enable automatic scheduling of TBE. Figure 3-6 shows the automatic scheduling process.

Figure 3-6 Overall process of the Auto Schedule mechanism

1. When the auto schedule API is called, a compute syntax tree is transferred. Each compute statement in the TBE is added with the tag_scope tag during compilation.

   Add the tag_scope tag as follows:

   with tvm.tag_scope(op):
       tmp = tvm.compute(shape, lambda_func, name=name)

   As shown in Figure 3-7, the compute syntax tree on the left is also called abstract syntax tree (AST). During compilation, tag_scope is added to each compute statement.

   Figure 3-7 Example of mapping between compute statements and tag_scope
   
2. The corresponding pattern is identified based on the scope tag. The pattern types supported by the TBE include elewise, reduce, segment, concat, conv, depthwise, and pooling2d. The TBE splits the AST based on the pattern rule. For example, the simplest pattern rule is that the elewise pattern can be connected to other patterns, and reduce, segment, and concat cannot be in the same AST subgraph.
3. After the AST subgraph is split, the TBE creates and initializes the schedule object.
4. During the scheduling process, the system finds the boundary of the AST subgraph, and then selects a proper schedule template for each subgraph according to the pattern. The scheduling process includes data flow management, tiling, and instruction mapping.