本文分享自华为云社区《基于Ascend C的Matmul算子性能优化最佳实践》，作者：昇腾 CANN。

矩阵乘法是深度学习计算中的基础操作，对于提升模型训练和推理速度至关重要。昇腾 AI 处理器是一款专门面向 AI 领域的 AI 加速器，其 AI Core 采用达芬奇架构，以高性能 Cube 计算引擎为基础，针对矩阵运算进行加速，可大幅提高单位面积下的 AI 算力。Matmul 算子实现的功能是矩阵乘法，通过 Ascend C 算子编程语言优化该算子的实现逻辑，可以使其在昇腾 AI 处理器上获得更优的执行性能。希望通过本案例的讲解，可以为开发者优化昇腾 Cube 类算子性能带来启发。

本案例以矩阵维度 M = 4096，N = 5120，K = 4096，输入数据类型 half，输出数据类型 float，输出格式是 ND 为例，性能验证平台为 Atlas A2 训练系列产品/Atlas 800I A2 推理产品，介绍针对 Matmul 算子的主要优化手段，包括优化分核逻辑、优化基本块、开启大包搬运。

• 优化分核逻辑：开启尽量多的 Cube 核使能并行计算。

• 优化基本块，选择最优的 baseM、baseN、baseK 参数。

• 开启大包搬运：从 Global Memory 搬运数据到 L1 时，对于 A 矩阵，一次搬入 depthA1 个基本块，基本块大小为 baseM * baseK，对于 B 矩阵，一次搬入 depthB1 个基本块，基本块大小为 baseN * baseK。使能大包搬运后，一次搬入的数据量变大，提升 MTE2 搬运效率。

分析主要瓶颈点

借助昇腾 Profiling 性能数据可较方便地分析主要瓶颈点，这里我们重点分析 MTE2，Cube，Scalar pipeline 的流水情况，其中 MTE2（Memory Transfer Engine）pipeline 反映了数据的搬入情况，Cube 和 Scalar pipeline 则反映了 AI Core 中的数据计算及标量的使用情况。

优化前 Profiling 数据如下图所示：

从上图 Profiling 数据来看，aic_mte2_ratio 数值是 0.973，这表明 MTE2 类型指令的 cycle 数在 total cycle 数中的占比过大，这意味着当前性能瓶颈点可能在于 MTE2 流水。此外，从图中的 Block Dim 数值 4 也可以看到，参与计算的 AI 处理器核并没有用满，这里假设当前案例使用的 AI 处理器上共有 20 个核。整体优化思路如下：

• 优化分核逻辑，假设 CurrentCore 是未优化前分核的 Cube 核数，MaxCore 为最大 Cube 核数，当开启全部核并行做当前 shape 数据量的计算时，预估性能收益约为 MaxCore / CurrentCore 的倍数。

• 优化基本块切分将影响搬运数据的效率，算子搬运的总数据量为搬运的左矩阵和右矩阵数据量之和。根据矩阵乘法的算法，搬运左矩阵的次数为 N / baseN，搬运右矩阵的次数为 M / baseM，即搬运总数据量 totalCnt = (N / baseN) * M * K + (M / baseM) * K * N。预估性能收益为搬运数据量的比值，优化前搬运数据量 totalCnt0/优化后搬运数据量 totalCnt1，化简后结果为(1 / baseM0 + 1 / baseN0) / (1 / baseM1 + 1 / baseN1)，其中，baseM0, baseN0 为优化前基本块参数，baseM1, baseN1 为优化后基本块参数。

• 开启大包搬运后，指令条数变化、地址对齐等因素会影响性能，按照经验预估，对于 MTE2 为性能瓶颈的场景，会有 20%以上的 MTE2 性能收益。

优化分核逻辑

由 Profiling 数据看出分核数为 4，启动更多的核同时计算，可以提高计算并行度。在当前案例使用的 AI 处理器上共 20 个核，每个核中包含 1 个 Cube Core 和 2 个 Vector Core。程序中设置 blockDim 为实际使用的核数 20。

// 代码片段 uint32_t blockDim = 20; // 优化前blockDim为4 CHECK_ACL(aclInit(nullptr)); aclrtContext context; int32_t deviceId = 0; CHECK_ACL(aclrtSetDevice(deviceId)); CHECK_ACL(aclrtCreateContext(&context, deviceId)); aclrtStream stream = nullptr; CHECK_ACL(aclrtCreateStream(&stream)); 
uint8_t *aHost; uint8_t *aDevice; CHECK_ACL(aclrtMallocHost((void **)(&aHost), aFileSize)); CHECK_ACL( aclrtMalloc((void **)&aDevice, aFileSize, ACL_MEM_MALLOC_HUGE_FIRST)); ReadFile("./input/x1_gm.bin", aFileSize, aHost, aFileSize); // PrintData(aHost, 16, printDataType::HALF); CHECK_ACL(aclrtMemcpy(aDevice, aFileSize, aHost, aFileSize, ACL_MEMCPY_HOST_TO_DEVICE)); 
uint8_t *bHost; uint8_t *bDevice; CHECK_ACL(aclrtMallocHost((void **)(&bHost), bFileSize)); CHECK_ACL( aclrtMalloc((void **)&bDevice, bFileSize, ACL_MEM_MALLOC_HUGE_FIRST)); ReadFile("./input/x2_gm.bin", bFileSize, bHost, bFileSize); // PrintData(bHost, 16, printDataType::HALF); CHECK_ACL(aclrtMemcpy(bDevice, bFileSize, bHost, bFileSize, ACL_MEMCPY_HOST_TO_DEVICE)); 
uint8_t *workspaceHost; uint8_t *workspaceDevice; CHECK_ACL(aclrtMallocHost((void **)(&workspaceHost), workspaceSize)); CHECK_ACL(aclrtMalloc((void **)&workspaceDevice, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST)); 
uint8_t *tilingHost; uint8_t *tilingDevice; CHECK_ACL(aclrtMallocHost((void **)(&tilingHost), tilingFileSize)); CHECK_ACL(aclrtMalloc((void **)&tilingDevice, tilingFileSize, ACL_MEM_MALLOC_HUGE_FIRST)); CHECK_ACL(aclrtMemcpy(tilingHost, tilingFileSize, GenerateTiling(), tilingFileSize, ACL_MEMCPY_HOST_TO_HOST)); // PrintData(tilingHost, 16, printDataType::UINT32_T); CHECK_ACL(aclrtMemcpy(tilingDevice, tilingFileSize, tilingHost, tilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE)); 
uint8_t *cHost; uint8_t *cDevice; CHECK_ACL(aclrtMallocHost((void **)(&cHost), cFileSize)); CHECK_ACL( aclrtMalloc((void **)&cDevice, cFileSize, ACL_MEM_MALLOC_HUGE_FIRST)); 
matmul_custom_do(blockDim, stream, aDevice, bDevice, cDevice, workspaceDevice, tilingDevice);

由于 Matmul API 都是从 Vector 侧发起的，按照 Cube Core 和 Vector Core 的配比 1：2，在 Matmul tiling 计算中需要按照 2 倍的 blockDim 数切分，因此 Tiling 代码中，设置 Tiling API 按照 40 个核进行数据切分，如下代码所示。

int usedCoreNum = 40; // 优化前usedCoreNum是8 int runMode = 1; int32_t baseM = 64; // 64 int32_t baseN = 64; // 64 optiling::TCubeTiling tilingData; auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion);     MultiCoreMatmulTiling tilingApi(*ascendcPlatform); tilingApi.SetDim(usedCoreNum);

修改代码后，算子执行时间（对应 aicore_time）从 12045us 下降到 2532us，约等于(20 核 / 4 核) = 5 倍的性能提升。

优化分核逻辑后 Profilling 数据如下图所示：

优化基本块

当前 Tiling 中设置的 base 块为 [baseM, baseN, baseK] = [64, 64, 256]，这种基本块 Cube 计算 cycle 少，计算访存比（即计算量与需要数据量的比值）低；搬出一次 Matmul 结果到 Global Memory 的 base 块是 64 * 64，由于输出格式是 ND，数据类型是 float，搬出下一次 Matmul 结果的起始地址需要偏移一个 baseN 的大小，即 64 * 4 = 256 字节，导致 fixpipe 搬出时 Global Memory 地址非 512byte 对齐，那么需要设置更优的基本块。

针对当前 shape 较大的场景，基本块的选择原则为计算访存比最大，即在 Cube 计算量最大的情况下，访存的数据量最小。在输入为 fp16 类型的情况下，Cube 执行单元 1 cycle 能算 16 * 16 * 16 个数。根据经验，[baseM, baseN, baseK] = [128, 256, 64]和[128, 128, 128]两种切分方案均满足搬出时 Global Memory 地址 512Byte 对齐（每搬出一次 Matmul 结果时，地址分别偏移 256 * 4byte 和 128 * 4byte），Cube 计算 cycle 数一致，为(128 * 64 * 256) / (16 * 16 * 16) = (128 * 128 * 128) / (16 * 16 * 16) = 512cycle。

针对[baseM, baseN, baseK] = [128, 256, 64]，计算访存比为 512cycle / (128 * 64 * 2 + 256 * 64 * 2) = 512cycle / 48KB；针对[baseM, baseN, baseK] = [128, 128, 128]，计算访存比为 512cycle / (128 * 128 * 2 + 128 * 128 * 2) = 512cycle / 64KB。可见，[128, 256, 64]基本块方案的计算访存比更高，计算密度更大，同样的计算量，需要的数据量最小，可最大限度地提高 Cube 单元计算量。

修改 Tiling 代码，通过 SetFixSplit()接口设置 baseM 和 baseN，tiling 函数会自动计算出最优 baseK，这里得到 64。

int32_t baseM = 128; // 优化前baseM是64 int32_t baseN = 256; // 优化前baseN是64 
optiling::TCubeTiling tilingData; auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion);     MultiCoreMatmulTiling tilingApi(*ascendcPlatform); tilingApi.SetDim(usedCoreNum); tilingApi.SetAType(leftPos, leftFormat, leftDtype, bool(transposeA)); tilingApi.SetBType(rightPos, rightFormat, rightDtype, bool(transposeB)); tilingApi.SetCType(resPos, resFormat, resDtype); tilingApi.SetBiasType(biasPos, biasFormat, biasDtype); 
tilingApi.SetOrgShape(M, N, K); tilingApi.SetShape(M, N, K); tilingApi.SetFixSplit(baseM, baseN, -1);

从下图可以看到，使能这组基本块后，MTE2 耗时（对应 aic_mte2_time）从 2452us 降低到 808us，MTE2 性能提升 3 倍。

优化基本块后 Profilling 数据如下图所示：

使能大包搬运

当前带宽利用率为：totalSize / mte2Time = totalCnt * dtype / mte2Time，代入数据计算为 2491GB/s。未使能大包搬运的情况下，矩阵从 Global Memory 搬运到 L1 一次只搬运 1 个基本块。通过模板参数使能大包搬运，一次搬运多个基本块，提高 MTE2 带宽利用率。

// 原始matmul对象定义: Matmul<MatmulType<TPosition::GM, CubeFormat::ND, A_T>, MatmulType<TPosition::GM, CubeFormat::ND, B_T>, MatmulType<TPosition::GM, CubeFormat::ND, C_T>, MatmulType<TPosition::GM, CubeFormat::ND, BiasT>>> mm; // 通过在定义matmul对象的模板参数里加上CFG_MDL参数使能大包搬运功能： Matmul<MatmulType<TPosition::GM, CubeFormat::ND, A_T>, MatmulType<TPosition::GM, CubeFormat::ND, B_T>, MatmulType<TPosition::GM, CubeFormat::ND, C_T>, MatmulType<TPosition::GM, CubeFormat::ND, BiasT>, CFG_MDL>> mm;

从下图可以看到，使能大包搬运后，MTE2 耗时从 808us 下降到 591us，带宽利用率代入数据计算为 3406GB/s，利用率提升 36%+，Cube 利用率（对应 aic_mac_ratio）达到 80%+。

使能大包搬运后 Profilling 数据如下图所示：

验证优化方案性能收益

• 优化分核逻辑，实际收益 75 倍，约等于(20 核 / 4 核) = 5 倍收益，并且考虑到核的启动开销，可以认为两者基本一致。

• 优化基本块，实际收益约 3 倍，理论评估带入上述分析公式，收益为(1 / 64 + 1 / 64) / (1 / 128 + 1 / 256)，约等于 7 倍，考虑到 cache 缓存的影响，可以认为两者基本一致。

• 大包搬运，大包搬运实际收益 25%+，与经验值基本一致。

但需要注意的是，优化分核逻辑和基本块一般在输入数据 shape 足够大、数据量足够多时，才能分满核和使能最优的基本块。因此，大 shape 场景下 MTE2 Bound 算子可参考此案例的优化手段。

更多学习资源

了解更多 Ascend C 算子性能优化手段和实践案例，请访问昇腾社区 Ascend C 信息专区：https://www.hiascend.com/ascend-c

相关推荐阅读：

《基于Ascend C的FlashAttention算子性能优化最佳实践》

华为开发者空间，汇聚鸿蒙、昇腾、鲲鹏、GaussDB、欧拉等各项根技术的开发资源及工具致力于为每位开发者提供一台云主机、一套开发工具及云上存储空间，让开发者基于华为根生态创新。

点击链接，免费领取您的专属云主机