Auto Tiling

Core Process

Figure 1 Core process of Auto Tiling

The execution logic of the kernel is expressed in AscIR. Auto Tiling extracts key information based on the expressions of the AscIR graph, including:

1. LocalBuffer usage modeling

   Auto Tiling must ensure that the local buffer usage of each level is within the range allowed by hardware. For example, the sum of the sizes of TQue/TBuf requested by the kernel and the size of the temporary buffer cannot exceed the Unified Buffer size limit of hardware. AscIR expresses the location of each tensor on the global memory/Unified Buffer and the reuse relationship between tensors. Based on the information, Auto Tiling symbolizes the constraints of the local buffer at each level.

2. Time consumption formula modeling

   Auto Tiling conducts performance modeling for each API, and uses symbols to express the performance of these APIs in each pipeline. The number of API calls is determined based on the loop axis expressed by AscIR. In this way, the total time consumed by all APIs of AscIR in each pipeline is derived. Assume that the execution of the bottleneck pipeline can well hide the execution of the non-bottleneck pipeline. Therefore, the total time consumed for task execution is mainly reflected in the time of the bottleneck pipeline, as shown in the following figure.

   The following figure shows three pipelines. The bottleneck pipeline is AIV_MTE2. When AIV_MTE2 is executed, the execution of AIV_MTE3 and AIV_VEC is hidden. The total time consumed by the task is reflected by the time consumed by AIV_MTE2.

   Figure 2 Pipeline diagram
   
3. Tiling Solver

   Auto Tiling aims to minimize the time consumption of bottleneck execution units and solve the problem by optimizing the model through modeling data. The prerequisite is that the storage usage for operator implementation does not exceed the physical storage size of each NPU level. For example, the heuristic solution method, starting from the initial solution, is used to search for the solution space that meets the memory requirements in the feasible domain meeting the constraints based on the gradient descent direction of the performance formula modeling, until the value of the formula modeling reaches the minimum value. In this way, the optimization process is exited and the optimal tiling result is returned.

Key Technologies

• Performance formula modeling
  • Objective: The simulation accuracy of the performance formula determines the accuracy of template selection and heuristic solution. Therefore, accurate performance modeling is required for APIs.
  • Implementation:
    • For stable basic Ascend C APIs, the Auto Tiling module collects performance statistics based on different inputs, obtains the relationship between performance and inputs, establishes a performance model, and expresses the performance model using the symbolic technology.
    • For variable Ascend C APIs, the Auto Tiling module computes the input parameters and number of calls based on the API calling logic, and generates a complete performance model for the APIs to obtain relatively accurate performance modeling.
• Axis sorting solver
  • Objective: A performance-formula-driven gradient descent solver relies heavily on the accuracy of its performance model. In addition, the tiling solution time is uncertain, and it is likely that the host bound problem may occur. As the APIs have a specific preference for axis tiling, an axis sorting solver is designed as an alternative solution for iteration solving.
  • Implementation:
    1. Determine the priority of the tiling axis based on the API. The sequence is as follows:
       1. The priority of the parent axis is higher than that of the child axis.
       2. The priority of the reduction axis is higher than that of the non-reduction axis.
       3. The priority of the broadcast axis is higher than that of the non-broadcast axis.
       4. The priority of the non-innermost axis is higher than that of the innermost axis.
    2. Then, divide the tiling operation into two parts.
       • Intra-core tiling: The variables are traversed in reverse order of axis sorting. The variables are preferentially adjusted to the maximum values to determine whether the hardware constraint conditions are met. If the hardware constraints are not met, the variables are adjusted using the dichotomy until these constraints are met. Then, the next intra-core tiling variable is adjusted until all variables meet the hardware constraints. For example, s1tt2, s1tt, s1t and s2t are tiling-related axes. The tiling process is as follows:

         s2t is traversed first and adjusted to the maximum value 1024, which meets the hardware constraints. Then, s1t is traversed and adjusted to the maximum value 256, which meets the hardware constraints. Then, the next variable s1tt>s1tt2 is adjusted until all variables meet the hardware constraints.

         Figure 3 Intra-core tiling
         
       • Multi-core tiling: Variables related to multiple cores are identified, these variables are traversed in descending order, and the record with a larger number of occupied cores is found. If the found number exceeds the number of physical cores (for example, the number of NPU cores is 48 for Atlas A2 training product / Atlas A2 inference product ), this result is returned.

         The following figure shows the tiling process, in which bngs1T is the axis for multi-core tiling. If the number of occupied cores is different, the record with a larger number of occupied cores is preferentially selected. According to the preceding policy, the number of occupied cores 47 is selected.

         Figure 4 Multi-core tiling