Operator Implementation

The following uses the Matmul+LeakyReLU fused operator as an example to describe the design and implementation process of the Mix fused operator. This example can run only on the Atlas A2 training product / Atlas A2 inference product .

The operator design process consists of operator analysis, data flow analysis, and tiling strategy design.

Operator Analysis

Operator analysis is to specify the mathematical expression, input, output, and kernel function name of an operator.

  1. Specify the mathematical expression and compute logic of an operator. Computation logic of the operator is as follows: A matrix multiplication operation is performed first, and then a LeakyReLU operation is performed on a matrix multiplication result and an alpha parameter. The mathematical expression is as follows: 

    1
    		 
    c = LeakyRelu(a * b + bias, alpha);

  2. Specify the input and output.

    • The inputs of the Matmul+LeakyReLU operator include a, b, and bias, and the output is c. alpha, the coefficient of the activation function LeakyReLU, is a fixed value and can be used as a constant value for compute in operator implementation.
    • In this example, the data type supported by the operator inputs a and b is the half type (float16), the data type supported by the operator input bias is float32, and the data type of the operator output c is float32.
    • The shape of the input matrix a is [M, K], the shape of the input matrix b is [K, N], the shape of the output matrix c is [M, N], and the shape of the input bias is [1, N].
    • The input and output data format of the operator is ND.

  3. Define the kernel function name and parameters.

    • You can customize the kernel function name. In this example, the kernel function is named matmul_leakyrelu_custom.
    • Based on the analysis of the operator input and output, the kernel function has these parameters: a, b, bias, and c. a, b, and bias indicate the memory address of the input in the global memory, and c indicates the memory address of the output in the global memory.

Based on the preceding analysis, the design specifications of the Ascend C Matmul+LeakyReLU operator are as follows:

Data Flow Analysis

Analyze the data flow of the operator. The data flow direction is as follows: After Matmul compute is completed on the Cube Core, data is moved to the Vector Core for LeakyReLU compute. Based on the preceding data flows and the programming paradigm of the fused operator, plan parallel pipeline tasks. The following figure shows the details:

  1. Move input data from the global memory to the Cube Core.
  2. Perform internal Matmul compute. The compute formula is as follows:

    Note that the shape of bias is [1, N]. Each row of the A × B result matrix is biased.
    Figure 1 Matmul matrix multiplication

  3. Move the Matmul compute result to the Vector Core.
  4. Perform the vector compute. In this example, LeakyReLU compute is performed.

    The LeakyReLU (Leaky Rectified Linear Unit) activation function is a commonly used activation function in artificial neural networks. Its mathematical expression and function graph are shown below.

  5. Move the output to the global memory.

The first three steps are encapsulated in the Matmul high-level APIs. In this example, there are only three stages. The following figure shows the details:

According to the preceding analysis, the Matmul high-level APIs, LeakyReLU vector compute API, Mask Operations, EnQue, and DeQue are used during the implementation.

Tiling Strategy Design

Tiling strategy design mainly includes multi-core tiling and intra-core tiling strategies.

The following figure shows the tiling strategy. For details, see Data Tiling.

Operator Implementation

In Cube Programming (High-Level APIs), Ascend C provides a group of Matmul high-level APIs, which encapsulate common algorithm logic for tiling, data movement, and compute, helping you quickly implement Matmul matrix multiplication. The implementation of cube programming in a fused operator is similar. You can call APIs on the host to automatically obtain tiling parameters. After the parameters are passed to the kernel during initialization, matrix multiplication can be completed by using several simple APIs. With reference to the foregoing programming paradigm of the fused operator, the procedure for implementing the fused operator is as follows. For a complete example, see MatmulLeakyRelu.

The following figure shows the code framework implemented on the kernel. After the Matmul object is initialized and the left matrix A, right matrix B, and bias are set, the subsequent Matmul compute, LeakyReLU compute, and CopyOut process are completed by adding a single Iterate loop and a while loop.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
 
template<typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulLeakyKernel<aType, bType, cType, biasType>::Process(){
    uint32_t computeRound = 0;
    
    // Set the input of Matmul (including the left matrix, right matrix, and bias).
    matmulObj.SetTensorA(aGlobal);
    matmulObj.SetTensorB(bGlobal);
    matmulObj.SetBias(biasGlobal);
    // Call Matmul iterate to obtain the compute result of [baseM, baseN].
    while (matmulObj.template Iterate<true>())
    {
        MatmulCompute();
        LeakyReluCompute();
        CopyOut(computeRound);
        computeRound++;
    }
    matmulObj.End();
}
// Kernel entrypoint function, Mix scenario, and AIC:AIV = 1:2
__global__ __mix__(1, 2) void matmul_leakyrelu_custom(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
                                                      __kfc_workspace__ GM_ADDR workspace, AscendC::tiling::TCubeTiling tiling)
{
    AscendC::TPipe pipe;
    MatmulLeakyKernel<half, half, float, float> matmulLeakyKernel;
    matmulLeakyKernel.Init(a, b, bias, c, workspace, tiling, &pipe);
    // Initialize the Matmul object.
    REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), matmulLeakyKernel.matmulObj, &matmulLeakyKernel.tiling);
    matmulLeakyKernel.Process(&pipe);
}

The implementation code of Matmul compute, LeakyReLU compute, and CopyOut is as follows:

  1. Matmul compute:

    1. Perform the internal Matmul compute on the Cube Core.
    2. Move the Matmul compute result to the Vector Core. 
       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
      		 
      template<typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulLeakyKernel<aType, bType, cType, biasType>::Process(){
    uint32_t computeRound = 0;
    // ...
    // Call Matmul iterate to obtain the compute result of [baseM, baseN].
    while (matmulObj.template Iterate<true>())
    {
        MatmulCompute();
        // ...
        computeRound++;

    }
    matmulObj.End();
}

template<typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulLeakyKernel<aType, bType, cType, biasType>::MatmulCompute(){
    reluOutLocal = reluOutQueue_.AllocTensor<cType>();
    // Call GetTensorC to move the Matmul compute result to the Vector Core.
    matmulObj.template GetTensorC<true>(reluOutLocal, false, true);
}

  2. Perform LeakyReLU compute. 

    1
2
3
4
5
6
    		 
    // Call the LeakyRule API to compute.
template<typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulLeakyKernel<aType, bType, cType, biasType>::LeakyReluCompute(){
    AscendC::LeakyRelu(reluOutLocal, reluOutLocal, (cType)alpha, tiling.baseM * tiling.baseN);
    reluOutQueue_.EnQue(reluOutLocal);
}

  3. CopyOut, move the output to the global memory. 

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
    		 
    // Move the result to the global memory.
template<typename aType, typename bType, typename cType, typename biasType>
__aicore__ inline void MatmulLeakyKernel<aType, bType, cType, biasType>::CopyOut(uint32_t count){
    reluOutQueue_.DeQue<cType>();
    const uint32_t roundM = tiling.singleCoreM / tiling.baseM;
    const uint32_t roundN = tiling.singleCoreN / tiling.baseN;
    uint32_t startOffset = (count % roundM * tiling.baseM * tiling.N + count / roundM * tiling.baseN);
    AscendC::DataCopyParams copyParam = {(uint16_t)tiling.baseM,
        (uint16_t)(tiling.baseN * sizeof(cType) / DEFAULT_C0_SIZE), 0,
        (uint16_t)((tiling.N - tiling.baseN) * sizeof(cType) / DEFAULT_C0_SIZE)};
    AscendC::DataCopy(cGlobal[startOffset], reluOutLocal, copyParam);
    reluOutQueue_.FreeTensor(reluOutLocal);
}

Implement the GenerateTiling function on the host to automatically obtain tiling parameters. The key procedure is as follows:

  1. Create a tiling object. 

    1
2
    		 
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
matmul_tiling::MultiCoreMatmulTiling cubeTiling(ascendcPlatform);

    When creating an object, you need to pass the hardware platform information, which can be obtained by calling GetPlatformInfo.

  2. Set the data types and formats of A, B, and bias.

    The following is a setting example. TPosition::LCM is the logical location on Unified Buffer, which is equivalent to TPosition::VECCALC. For details about TPosition, see TPosition. 
    1
2
3
4
    		 
    cubeTiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
cubeTiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
cubeTiling.SetCType(matmul_tiling::TPosition::LCM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
cubeTiling.SetBiasType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

  3. Set the matrix shape. 

    1
2
    		 
    cubeTiling.SetShape(M, N, K);
cubeTiling.SetOrgShape(M, N, K);

  4. Set the size of the available space.

    Set the size of the available space on L1 Buffer, L0C Buffer, or Unified Buffer for Matmul compute. The value –1 indicates the size of the buffer corresponding to the AI processor. 
    1
    		 
    cubeTiling.SetBufferSpace(-1, -1, -1);

  5. Set other parameters as required, for example, bias that will participate in the compute. 

    1
    		 
    cubeTiling.SetBias(true);

  6. Obtain tiling parameters. 

    1
2
3
4
    		 
    MatmulLeakyreluCustomTilingData tiling;
if (cubeTiling.GetTiling(tiling.cubeTilingData) == -1){
    return ge::GRAPH_FAILED;
}

  7. Perform other operations such as serialization and saving of tiling parameters.
  • In multi-core scenarios, you need to set the number of cores used for Matmul compute through SetDim. The rules for setting the MIX mode (including cube compute and vector compute) are as follows:
    • Separated mode: The Matmul API is initiated from the AIV. When Iterate is called for compute, the AIV only notifies that the AIC needs to perform cube compute. After that, the AIC notifies the AIV that the compute is complete. In this architecture, SetBlockDim is set to the number of AI Cores (AIC + AIV) used for compute, and SetDim is set to the number of AIVs used for compute. For example, when SetBlockDim is set to 20, 20 AI Cores (AIC + AIV) are started, and when SetDim is set to 40, 40 AIVs are tiled.
    • Coupled mode: The number of cores loaded by SetBlockDim is the same as that used by the Matmul API for compute. The values of SetDim and SetBlockDim are the same.
  • The system workspace is required for internal implementation of Matmul high-level APIs. You need to:
    • Set the total workspace size (including the user workspace and system workspace) when implementing tiling on the host. The workspace space is allocated and managed by the framework. The size of the system workspace can be obtained by calling GetLibApiWorkSpaceSize. 
      1
2
3
4
      		 
      size_t userWorkspaceSize = 0;
size_t systemWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
size_t *currentWorkspace = context->GetWorkspaceSizes(1);
currentWorkspace[0] = userWorkspaceSize + systemWorkspaceSize;
    • If the operator project is neither a custom operator project nor a kernel launch operator project with the HAVE_WORKSPACE compilation macro, the kernel needs to set the system workspace through SetSysWorkSpace before Matmul initialization. 
      1
2
3
4
5
      		 
      // The workspace space must be set when Matmul is used.
SetSysWorkspace(workspace);
if (GetSysWorkSpacePtr() == nullptr) {
    return;
}
  • In the preceding implementation method, the code isolation and inter-core synchronization on the AIC and AIV are implemented by the framework. In addition to this method, you can also choose to implement fusion operators in separated mode through underlying coding, which is more flexible. When using the underlying coding method, pay attention to the following points:
    • Implement synchronization between AIC Cores and AIV Cores. For example, in the Matmul+LeakyReLU operator sample, ensure that the AIV performs LeakyReLU compute after the AIC completes cube compute.
    • When using Matmul high-level APIs, you need to set ASCENDC_CUBE_ONLY, indicating that Matmul APIs are called only on the AIC.
    • Use the API in Setting Kernel Type to set the kernel type to KERNEL_TYPE_MIX_xxx and enable both AIV Cores and AIC Cores.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
    		 
    #define ASCENDC_CUBE_ONLY // Specify that Matmul operators run on AICs.
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); // Set the kernel type to KERNEL_TYPE_MIX_xxx.
if ASCEND_IS_AIC {
    ...
    // The AIC Core performs Matmul compute.
    // After the AIC Core completes the compute, the synchronization flag is sent by using AscendC::CrossCoreSetFlag<modeId, pipe>(flagId).
}
if ASCEND_IS_AIV {
    ...
    // The AIV Core receives the synchronization flag by using AscendC::CrossCoreWaitFlag(flagId).
    // The AIV Core performs LeakyReLU compute.
}

    For a complete example, see BareMix sample.

Parent topic: CV Fusion