Matmul High-Level API Enabling CUBE_ONLY

Case Study

This case shows the optimized performance when CUBE_ONLY of the high-level Matmul API is enabled in the matrix multiplication operator scenario. As shown in the following figure, the Matmul API uses the MIX mode by default. That is, the user initiates a message from the AIV, and the message is forwarded through the message communication framework and then Matmul computation is performed on the AIC. This message processing mechanism causes extra Scalar performance overhead. Compared with the MIX mode, CUBE_ONLY can directly skip the message communication framework to complete MatMul computation, improving operator performance.

Figure 1 Matmul process in MIX mode by default
Table 1 Operator case specifications

Input

Shape

Data Type

Format

a

128, 64

float16

ND

b

64, 30720

float16

ND

The AI Processor used in this case has 24 cores, each of which contains one AIC and two AIVs.

The tiling parameters are as follows:

Obtaining Profile Data

Use the msProf tool to obtain the Operator Simulation Pipeline and On-board Profiling data. The CUBE_ONLY mode mainly optimizes the Scalar pipeline performance. Therefore, you can focus analyzing the Scalar pipeline.

Analyzing Main Bottlenecks

Developing Optimization Solutions

In the default MIX mode, the user initiates a message on the AIV side. After the message is forwarded by the message communication framework, the Matmul computation is performed on the AIC side. Based on this process, when using the high-level Matmul API to compile operator code, you can use the REGIST_MATMUL_OBJ macro, without distinguishing the AIV and AIC. However, this message processing mechanism causes additional performance overhead, as shown in Figure 1 Matmul process in the default MIX mode.

The procedure for implementing the default MIX mode is as follows:

  1. On the kernel side, define the Matmul object. 

    1
2
3
4
5
6
7
    		 
    #include "lib/matmul_intf.h"

using A_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType>;
using B_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType>;
using C_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, CType>;
using BIAS_TYPE =  AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType>;
AscendC::Matmul<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM> matmulObj;

  2. On the host side, the Matmul multi-core tiling object calls the SetDim API to set the number of cores involved in the computation. 

    1
2
3
4
    		 
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
matmul_tiling::MultiCoreMatmulTiling cubeTiling(*ascendcPlatform);
int32_t numBlocks = ascendcPlatform->GetCoreNumAiv(); // In MIX mode, use GetCoreNumAiv to obtain the number of available cores of the AI Processor.
cubeTiling.SetDim(numBlocks);

  3. Call the kernel function and set the numBlocks parameter of the kernel function by referring to Kernel Function Definition and Calling. 

    1
    		 
    In matmul_custom_do(ascendcPlatform->GetCoreNumAic(), stream, x1, x2, bias, y, workspaceDevice, tilingDevice); // In the MIX mode, the startup is performed by groups of AIVs and AICs. numBlocks is used to set the number of AI Cores to be started.

In the operator scenario without Vector computation, you can skip the message communication framework mechanism and enable the CUBE_ONLY mode to complete Matmul computation, reducing the performance overhead of message communication and improving the operator performance.

Figure 2 Matmul process in CUBE_ONLY mode

For details about the complete example of enabling the CUBE_ONLY mode using the Matmul API, see Matmul API performance optimization sample. The procedure for enabling the CUBE_ONLY mode is as follows:

  1. On the kernel side, set the ASCENDC_CUBE_ONLY macro before including the matmul_intf.h header file in the code for defining the Matmul object. 

    1
2
3
4
5
6
7
8
    		 
    #define ASCENDC_CUBE_ONLY // Set the ASCENDC_CUBE_ONLY macro before #include "lib/matmul_intf.h".
#include "lib/matmul_intf.h"

using A_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType>;
using B_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType>;
using C_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, CType>;
using BIAS_TYPE =  AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType>;
AscendC::Matmul<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM> matmulObj;

  2. On the host side, the Matmul multi-core tiling object calls the SetDim API to set the number of cores involved in the computation. 

    1
2
3
4
    		 
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
matmul_tiling::MultiCoreMatmulTiling cubeTiling(*ascendcPlatform);
int32_t numBlocks = ascendcPlatform->GetCoreNumAic(); // In CUBE_ONLY mode, use the GetCoreNumAic API to obtain the number of available cores of the AI Processor.
cubeTiling.SetDim(numBlocks);

  3. Call the kernel function and set the numBlocks parameter of the kernel function by referring to Kernel Function Definition and Calling. 

    1
    		 
    matmul_custom_do(ascendcPlatform->GetCoreNumAic(), stream, x1, x2, bias, y, workspaceDevice, tilingDevice); // For operators that contain only Cube Units, numBlocks is used to set the number of AICs to be started.

  4. On the kernel side, add the return branch from the AIV to the kernel function implementation. 

    1
2
3
4
5
6
7
8
9
    		 
    extern "C" __global__ __aicore__ void matmul_custom(GM_ADDR a, GM_ADDR b,
    GM_ADDR bias, GM_ADDR c, GM_ADDR workspace, GM_ADDR tilingGm)
{
    if (g_coreType == AscendC::AIV) { // CUBE_ONLY mode, with AIV directly returning
        return;
    }
    ...
    // Other code
}

Verifying Optimization Benefits

Summary

In scenarios where only Cube computation is involved and Vector computation is not involved, you can enable the CUBE_ONLY mode to optimize the message communication performance overhead in Matmul computation and improve operator performance.

Parent topic: Matmul Performance Tuning Cases