Scheduling and Building an Operator

As shown in the following code, after the computation logic is defined, the Auto schedule mechanism performs auto scheduling. You can check the computation IR through the print mechanism provided by the TVM. The configuration information includes the print switch status, build switch status, operator name in the kernel, and input and output tensors.

sch = generic.auto_schedule(res)
config = {
    "print_ir": need_print,
    "need_build": need_build,
    "name": kernel_name,
    "tensor_list": [data, res]
}
te.lang.cce.cce_build_code(sch, config)

• Use the auto_schedule interface of generic to perform auto scheduling (defining schedule). The argument of the auto_schedule interface is the output tensor of the operator.

  The schedule object defines how to efficiently execute the described computation process on hardware. That is, the related computations are mapped to the corresponding instructions on a hardware device. A schedule object contains an IR, which uses code similar to pseudo code to describe a computation process. You can print the object by using the parameter need_print.

• The input and output tensors are stored in tensor_list. The input and output tensors must be arranged in the input and output sequences of the operator.

  For example: "tensor_list": [tensor_a, tensor_b, res], where, tensor_a and tensor_b are the input tensors, and res is the output tensor.

• The cce_build_code interface provided by te.lang.cce is used to build the operator based on scheduling and configuration. During operator building when OMG converts the model, a dedicated kernel is built based on the input data shape, type, and operator parameters.
  • sch: schedule object to be calculated by the generated operator.
  • config: map configured by compilation parameters.

  After building, an operator target file *.o (the running target of the operator is AI Core) or *.so (the running target of the operator is AI CPU) and an operator description file *.json are generated.