CANN/catlass单核切K矩阵乘样例-开发者社区

CATLASS Single_core_splitK_Matmul 样例介绍

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass

样例实现

CATLASS34_single_core_splitk_matmul样例算子是基于CATLASS Gemm API实现的昇腾亲和Matmul算子，针对大尺寸矩阵计算场景优化设计，关键算子组件包括以下几部分:

Example组装：single_core_splitk.cpp
Kernel实现：
- 主Kernel文件：single_core_slicek_matmul.hpp
- 复用Padding组件：padding_matmul.hpp
Block组件： block_mmad_single_core_splitk.hpp

Example组装

与通用的模板库开发风格一致，本样例的实现执行过程如上图所示，简要说明如下：

构造输入

生成要计算的左/右矩阵

计算各输入矩阵的尺寸
在Host侧生成各输入矩阵值
构造Device侧输入

组装Padding对象

为便于数据搬运的对齐，组装Padding对象

定义不同的PaddingTag
通过PaddingBuilder组装出对于矩阵A/B的Padding对象
预定义的PaddingC对象
决定是否启用PaddingC对象

组装blockMmad

为便于数据搬运的对齐，先组装相关的Padding对象

定义各输入的Layout特征
声明不同矩阵所对应的类型情况
设置Dispatch策略（用于选取BlockMmad组件）
设置L1和L0的Tile尺寸，用于优化从GM到L1的搬运过程
使用优化后的TileCopy组件
使用上述模板入参组装BlockMmad

组装并执行Kernel

组装Kernel并实例化该对象，完成算子计算

使用新的S型Swizzle策略
使用前述模板组装Kernel
将Kernel传入适配器并实例化
构造入参参数
向Kernel侧传入参数
调用Kernel层Workspace计算
在Device侧申请Workspace(如有必要)
适配器初始化算子
执行算子

精度校验和空间释放

完成最后结果的验证与回收处理

将算子输出结果搬回host侧
计算golden标杆
精度比对
释放输入输出和workspace

Kernel实现

以下介绍Kernel层级的结构体与关键函数，以及AIC/AIV部分的简明计算流程，并说明了所使用的优化策略

Kernel层的主要结构体与函数

以下是在Kernel层所实现的结构体与关键函数

struct Params：运行时算子执行所需的参数
struct Arguments：封装Host侧传入的参数
static size_t GetWorkspaceSize：预先计算对齐所需的空间
static Params ToUnderlyingArguments：将Host侧入参解析为算子侧的Params结构体
void operator()<AscendC::AIV>：AIV(Vector)部分执行代码
void operator()<AscendC::AIC>: AIC(Cube)部分执行代码

AIV/AIC部分计算流程

以下是Kernel层所进行的AIC/AIV操作

AIV上所执行的操作

如果A或B的对齐处理启用：
- 初始化GlobalTensor：gmA， gmWA(或gmB, gmWB)
- 实例化PaddingA对象，完成资源申请，以便在UB(Unified Buffer)上存储对齐前后的数据
- 调用PaddingA对象的operator()方法，即执行数据对齐操作
- 进行核间同步，设置标志位通知AIC对齐操作已完成(CrossCoreSetFlag)
结果矩阵搬出：
- 等待AIC完成计算(CrossCoreWaitFlag)
- 初始化GlobalTensor：gmC(目的地址),gmWC(临时累加地址)
- 将数据搬出到目的地址，进行升精度映射(RemovePaddingNDAndCastC)

AIC上所执行的操作

如果对左/右矩阵有对齐处理：
- 核间同步，等待AIV完成标志位(CrossCoreWaitFlag)
- 初始化GlobalTensor：gmWA,gmWB
初始化GlobalTensor:gmA,gmB,gmC
初始化BlockScheduler和BlockMmad对象
获取当前AIC序号coreIdx、AIC总数coreNum（在Swizzle策略内）以及所需的coreLoops
进入主循环（循环次coreLoops）
- 计算当前A，B矩阵读入的偏移量gmOffsetA,gmOffsetB以及下一块的偏移量gmOffsetNextA,gmOffsetNextB(如果开DoubleBuffer)
注意：优化算法下左矩阵重载，因此gmOffsetNextA实际不会“启用”
- 计算needLoadNextA、needLoadNextB，用以标识是否预加载
- 调用blockMmad进行一次AIC计算（完成L1A与L1B上相对应的分形矩阵计算）
注意：blockMmad会依据K轴方向上切K情况，决定是否启用GM上的原子加
- 设置标志位，通知AIV 计算已完成
- 关闭原子加

单核切K算法

如上图所示，相较于经典的矩阵乘过程，单核切K的Matmul模板采取了复用左矩阵的办法，即对于某个AI Core，其L1A上加载的分形矩阵是驻留的，进而降低对GM的访问量。

举例而言，对于经典计算过程，单核所进行的数据搬运过程包括：1. 从GM搬运A0到L1A，B0到L1B；2. 从GM搬运A1到L1A，B1到L1B；..., N. 从GM搬运A(n-1)到L1A， B(n-1)到L1B。

在本单核切K策略下，单个AI Core上的流程图如下所示：

显然，该办法能有效减轻MTE2的搬运负担，但是会加重L0C计算结果搬回到GM的负担，需要在GM上开启原子加。

// Atomic set once in mmad level if (atomicAdd) { AscendC::SetAtomicAdd<ElementAccumulator>(); } else { AscendC::SetAtomicNone(); }

增大L1利用空间简单建模可知（见下述补充），从L0C搬运回GM的数据量正比于$MNK/k_{\text{L1}}$，其中$M$, $N$, $K$分别为输入的矩阵尺寸，$k_{\text{L1}}$是K轴上L1Tile的尺寸大小。在符合L1A, L1B物理尺寸限制的条件下，$k_{\text{L1}}$越大可减少写出次数，推荐L1TileShape按下表进行配置。

数据搬运理论建模

我们首先分析下基础Matmul（详见本仓00_basic_matmul样例）的内存访问量。设矩阵乘的尺寸为$(M, N, K)$，L1上的Tile尺寸相应为$m$, $n$, $k$。假设二者对齐。则单个分形（结果）矩阵从GM（Global Memory）至L1所需的搬运量是$K(mk+kn)/k$，加之分形矩阵个数为$MN/mn$，因此总的访问量为$MNK(1/m+1/n)$。另一方面，由于在L0C上进行累加，因此搬出的单位数据量即为$MN$。再考虑使用本优化策略后的内存访问，设复用左矩阵，则单核总的数据搬运量为$mK + KN$，总的访问量是$(mK + KN)M/m$，也即$MNK(1/m+1/N)$，小于基础Matmul的情形。另外由于计算结果并不能在L0C上进行累加，需“随算随搬”，因此搬出的数据量为$MNK/k$，较基础有一定升高。

类别	内存访问操作总量
基础矩阵乘	$MNK(1/m+1/n) + MN$
单核切K	$MNK(1/m+1/N) + MNK/k$

数据类型	`L1TileShape::M`	`L1TileShape::N`	`L1TileShape::K`
FP16/BF16	256	128	512
FP32	256	128	256

数据对齐如需对数据进行类型转换，需采取数据对齐，以提高处理流程中的数据搬运带宽。

// RemovePaddingNDAndCast的核心处理逻辑 uint32_t loopsPerTile = RoundUp(tileLen, COMPUTE_LENGTH); uint32_t coreLoops = tilesPerAiv * loopsPerTile; for (uint32_t loopIdx = 0; loopIdx < coreLoops; ++loopIdx) { // 计算gmWC搬运的offset // 将GmWC(src)的数据搬运至inputBuffer上 copyGm2Ub(inputBuffer, src, ubLayout, srcLayout); // ... // Cast到half AscendC::Cast(outputBuffer, inputBuffer,AscendC::RoundMode::CAST_RINT, actualDataNum); // ... // 将outputBuffer搬出到gmC(dst) copyUb2Gm(dst, outputBuffer[bufferIndex], dstLayout, ubLayout); // ... }

Swizzle排布方案

在矩阵C上的swizzle策略采用S型特征，相较于Z型的swizzle排布，在换行处可节约一次搬运。如下图所示，按该S型Swizzle策略，前一组需要加载到L1A, L1B上的分形矩阵为<A02, <B20, B21, B22, B23>>，后一次需要加载的是<A12, <B20, B21, B22, B23>>(以SwizzleOffset为4为例)。“换行”过程下，可以固定L1B上的数据，仅从GM上重读入左矩阵，可减轻载入负担。

说明：该Swizzle策略在大尺寸下有效，排布策略会优先将任务分配至不同的AI Core上。

性能收益

经过实测，应用上述单核切K算法在大尺寸场景下较基础Matmul有正向收益，且随着K轴的尺寸增大，GM至L1节省的正向收益高过重复写入GM的负向收益，可参考下表。

不过需要说明的是，在M，N过小，或者K偏低的情况下，会出现分核负载不均衡的情况。

| M | N | K | 耗时(us) | 耗时[标杆]`(us) | 加速比 | | -- | -- | ---- | ----- | ----- | ----- | | 2048 | 4096 | 4000 | 261 | 445 | 1.7049 | | 4096 | 4096 | 8000 | 917 | 1231 | 1.3424 | | 4096 | 4096 | 40000 | 4669 | 7775 | 1.6652 | | 2048 | 4096 | 80000 | 16850 | 57144 | 1.7499 |

说明：