TIK Operator Implementation Workflow

The following figure shows the general workflow for building a Python program of a TIK operator using TIK APIs.

Figure 1 Operator Implementation Workflow (AI Core)

If an operator is too large (for example, the number of nested variable layers in a scope of the operator is too large), the occupied space exceeds the threshold. In this case, an error may be reported due to insufficient system resources during compilation. You can run the ulimit -s command to temporarily increase the stack space. For example, if the default stack space is 8 MB, run the ulimit -s 16384 command to increase the stack space to 16 MB.

The code structure is as follows:

# Import the dependent Python modules.
from tbe import tik
import tbe.common.platform as tbe_platform
from tbe.common.utils import para_check
import numpy as np

# Entry point function for operator definition, that is, the function called when the operator is running. For details about the definition, see TIK Entry Point Function.
@para_check.check_input_type(dict, dict, dict, str)
def operation_name(input_x, input_y, output_z, kernel_name):

    # Define the target machine, that is, set the version of the Ascend AI Processor. The default running target is AI Core.
    # Set soc_version to the Ascend AI Processor version in use.
    tbe_platform.set_current_compile_soc_info(soc_version)
    # Create a TIK DSL container.
    tik_instance = tik.Tik(disable_debug=False)

    # Define data.
    data_x = tik_instance.Tensor("dtype", shape, name="data_x", scope=tik.scope_gm)
    ... ...
    data_x_ub = tik_instance.Tensor("dtype", shape, name="data_x_ub", scope=tik.scope_ubuf)
    # Move data to the specified core.
    tik_instance.data_move(data_x_ub, data_x, ...)

    # Compute data.
    tik_instance.vec_xxx(...)
    # Move data from the specified core.
    tik_instance.data_move(data_x, data_x_ub, ...)
    # Compile the operator into a target file. For details about the TIK compilation function definition, see TIK Building Function.
    tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
    # Return a TIK instance.
    return tik_instance