Updated on 2022-03-13 GMT+08:00

Description

To improve the usability of custom operators and the development efficiency, some vector operators are encapsulated in a modular manner. For element-wise operation APIs, you need to use TVM primitives to define the input tensor, call the encapsulated APIs, briefly describe the computation process of the custom operator, and then call the provided auto schedule and build APIs to compile the custom operator into a binary executable file.

TE provides a group of encapsulated APIs, covering vector calculation. The APIs include the element-wise operation APIs, reduction APIs, broadcast API, index operation APIs, concat API, convolution API, 4D/5D conversion APIs, and matrix multiplication API.

You can view the API definition files in the ddk/site-packages/te-0.4.0.egg/te/lang/cce/te_compute, ddk/site-packages/te-0.4.0.egg/te/lang/cce/te_schedule, and ddk/site-packages/topi-0.4.0.egg/topi/generic directories under the installation directory of the DDK package. If you install Mind Studio and the DDK both in boot installation mode, you can log in to the Mind Studio server as the Mind Studio installation user and view the API definition files under the corresponding directory in ~/tools/che/ddk/ddk/site-packages. For details about the mapping between the definition files and the APIs, see the following API description. For details about how to use these APIs to develop custom operators and compile and run code, see TE Custom Operator Development Guide.

Element-wise Operation APIs

The input data is calculated using different operations by element, and the output usually has the same shape as the input.

  • Dual-operand operations. Two tensors are input and calculated by element to obtain the result.

Reduction APIs

Compress the data of a dimension and perform operations such as accumulation and multiplication on the data in the specified direction. The operation output is one dimension less than the input data.

Usage limitation: Due to the data arrangement limitation on the CCE computing platform, the data after the reduction operation needs to be rearranged for subsequent operations. Therefore, when different types of APIs are used, no vector operation can be performed after the reduction operation.

Broadcast API

The broadcast operation is mainly used to process two tensors with different shapes and broadcast an operand with a lower dimension according to a higher-dimension operand. The element-wise calculation is performed after the dimensions of the two operands are the same.

Broadcast a smaller tensor to a larger tensor. For details, see te.lang.cce.broadcast(var, shape, output_dtype=None).

Index Operation APIs

Index operations are used to perform calculations on tensors by segment, such as obtaining the sum value, average value, inner product, maximum value, and minimum value.

Concat API

The concat API is used to concatenate multiple input tensors along an axis.

Concatenate the tensors along an axis. For details, see te.lang.cce.concat(raw_tensors, axis).

Convolution API

Convolution APIs are used to implement convolutional operators.

Implement convolution. For details, see te.lang.cce.conv(*args).

4D/5D Conversion APIs

These APIs are used for conversion between 4D NCHW and 5D NC1HWC0.

Matrix Multiplication API

This API is used to implement matrix multiplication.

For details, see te.lang.cce.matmul(tensor_a, tensor_b, trans_a=False, trans_b=False, alpha_num=1.0, beta_num=0.0, tensor_c=None).

Convergent Operator API

The TE operator convergence function pre-processes information such as the convergence capability and input parameters of a single operator. Therefore, the specified API needs to be called to generate the *.so and *.json files of a single operator.

For details, see bool BuildTeCustomOp(std::string ddkVer, std::string opName, std::string opPath, std::string opFuncName, const char *format, ...).