本案例呈现了在使用Matmul高阶API进行矩阵乘法计算时,使能Matmul Tiling全量常量化对算子性能的提升效果。Matmul API在初始化和迭代过程中有大量Scalar计算,Matmul初始化时的Scalar计算影响指令头开销,Matmul迭代间的Scalar计算可能阻塞MTE2流水。在调用Matmul API实现矩阵乘法时,使用参数替代TCubeTiling变量参数,将Scalar计算提前到编译期进行,以减少运行时的Scalar计算开销,实现算子性能的提升。
Matmul Tiling常量化的适用场景:
- Matmul初始化时的Scalar计算较多,影响指令头开销。
- Matmul迭代之间的Scalar计算较多,阻塞MTE2流水。
Matmul Tiling常量化需要在编译期确定部分Tiling参数,根据确定参数的不同,分为全量常量化和部分常量化两种场景,使用Matmul Tiling常量化需要满足两种场景中任一场景的条件:
[object Object][object Object]全量常量化:能够确定常量singleCore Shape(singleCoreM/singleCoreN/singleCoreK)和常量base Shape(basicM/basicN/basicK,也称baseM/baseN/baseK)。
[object Object][object Object]部分常量化:能够确定常量base Shape(basicM/basicN/basicK,也称baseM/baseN/baseK)。
其中,全量常量化场景比部分常量化场景可以减少更多的Scalar计算开销。
本案例的算子规格如下:
表 1 算子规格
[object Object][object Object]
[object Object]当前案例使用的AI处理器共24个核,每个核中包含1个AIC核和2个AIV核。
Tiling参数如下:
原始shape:M=128, N=30720, K=64。
单核shape:按24个AIC核进行切分,singleCoreM=128,singleCoreN=1280,singleCoreK=64。
对于B矩阵,沿着N轴进行切分,切分成24份的singleCoreN,单核上处理K * singleCoreN大小的数据。对于A矩阵,M轴不进行切分即singleCoreM=M,单核上处理singleCoreM * K大小的数据。总共24个核参与计算。
基本块shape:baseM=128,baseN=256,baseK=64。
L1相关Tiling参数:stepM=1,stepN=1,stepKa=4,stepKb=4,depthA1=8,depthB1=8。
使用msProf工具获取和数据。相较于基础场景,Tiling常量化在编译期期间将部分或全部Tiling参数由变量转化为常数值,在算子执行时直接使用常量化的Tiling参数,可以减少Scalar性能开销,所以重点分析Scalar流水。
优化前的流水图如下,默认不使能Tiling常量化,Tiling参数需要从Host侧拷贝到Kernel侧,导致Matmul初始化时的Scalar计算较多,第一个MTE2指令开始于3.536us左右,MTE2前的指令头开销在算子整个流水中占比较大,因此需要优化Scalar计算。
优化前的Profiling数据如下,从C列的aic_time数据来看,多个核中最大算子执行耗时为10.62us,从G列的aic_scalar_time数据来看,Scalar平均耗时6.32us。
如下图所示,默认不使能Tiling常量化功能时,开发者在host侧创建Tiling对象,通过调用API自动获取Tiling参数。然后将Tiling参数从Host侧传递到Kernel侧,在Kernel侧初始化操作时传入。在算子执行时,使用Tiling变量参数完成矩阵乘操作。
图 1 默认不使能Tiling常量化的Matmul计算流程示意图[object Object][object Object]
如下图所示,使能Tiling常量化功能时,开发者只需要在Kernel侧创建Matmul对象时,调用GetMatmulApiTiling接口在编译期获取常量化Tiling信息,即可完成Tiling常量化。在算子执行时,使用常量化的Tiling参数完成矩阵乘操作,减少Scalar计算开销。
图 2 使能Tiling常量化的Matmul计算流程示意图[object Object][object Object]
Matmul API使能Tiling全量常量化的完整样例请参考。使能Tiling全量常量化功能的步骤如下:
调用获取MatmulConfig模板的接口GetMMConfig时,使用常数值设置MatmulShapeParams,得到带有常量化参数的自定义MatmulConfig模板CUSTOM_CFG。
[object Object]创建Matmul对象。首先调用GetMatmulApiTiling接口,将Tiling信息常量化,得到常量化模板参数CONSTANT_CFG,包括常量化的Matmul Tiling信息和MatmulConfig模板。创建Matmul对象时,使用常量化模板参数CONSTANT_CFG。
[object Object]初始化操作。全量常量化时,可以在REGIST_MATMUL_OBJ接口的入参传递Tiling参数的位置,使用空指针替代。部分常量化时,在Kernel侧使用REGIST_MATMUL_OBJ接口初始化Matmul对象时,仍需要使用Tiling。
[object Object]
优化后的流水图如下,通过使能Tiling全量常量化,无需将Tiling参数从Host侧拷贝到Kernel侧,在编译期完成Tiling常量化,减少了Matmul初始化时的Scalar计算。从0us起到第一个MTE2指令发起,这之间的时间为Matmul初始化时间,Matmul初始化时间从优化前的3.536us减少到2.185us,性能有所提升。
优化后的Profiling数据如下,从C列的aic_time数据来看,多个核中最大算子执行耗时为7.87us,相较于优化前的10.62us提升了25.9%。从G列的aic_scalar_time数据来看,Scalar平均耗时3.38us,相较于优化前的6.32us提升了46.5%。
算子在调用Matmul API完成矩阵乘计算时,若Matmul初始化时的Scalar计算较多,影响了指令头开销,或Matmul迭代间的Scalar计算较多,阻塞了MTE2流水。在这两类场景下,满足上文提及的Tiling常量化使能条件(或),可以考虑使能Tiling常量化,减少Scalar计算开销,提升算子性能。