昇腾社区首页
中文
注册

Matmul使能AtomicAdd选项

【优先级】中

【描述】对于Matmul得到的结果矩阵C(m, n),若后续需要和GM上的矩阵D(m, n)进行Add操作,则可以在GetTensorC接口或者IterateAll接口的GM通路上,将enAtomic参数设为1,开启AtomicAdd累加操作,在搬出矩阵C到GM时,矩阵C的结果将直接累加到矩阵D的GM地址上,从而实现与矩阵D的Add操作。

【反例】

将Matmul的结果矩阵C和GM上的矩阵D分别搬到UB上,做完Add操作后,结果再搬出到GM。这样至少要多分配一块UB内存给矩阵D,假设在分离架构的处理器上执行,将多做三次搬运操作(矩阵C从GM搬到UB、矩阵D从GM搬到UB、Add结果从UB搬出到GM)。

 template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE>
 __aicore__ inline void MatMulKernel(...)
 {
    ...
    Matmul<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_MDL> mm;
    TPipe pipe;
    REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm);

    mm.SetTensorA(gm_a);
    mm.SetTensorB(gm_b);
    mm.SetBias(gm_bias);

    mm.IterateAll(local_c);

    // while (mm.Iterate()) {
        // mm.GetTensorC(local_c);
    // }
    
    DataCopy(local_d, gm_d, d_size);
    event_t eventIdMTE2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
    SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
    WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
    Add(local_d, local_d, local_c, d_size);
    DataCopy(gm_d, local_d, d_size);
    ...
 }

 extern "C" __global__ __aicore__ void example_kernel(...)
 {
     ...
     typedef MatmulType<TPosition::GM, CubeFormat::ND, half> aType; 
     typedef MatmulType<TPosition::GM, CubeFormat::ND, half> bType; 
     typedef MatmulType<TPosition::GM, CubeFormat::ND, float> cType; 
     typedef MatmulType<TPosition::GM, CubeFormat::ND, float> biasType;
     MatMulKernel<aType, bType, cType, biasType)(...);
     ...
 }

【正例】

计算Matmul结果时,调用IterateAll接口或者GetTensorC接口搬运到矩阵D的GM地址上,同时将接口中enAtomic参数设为1,搬出到GM时,Matmul结果矩阵C会累加到矩阵D上,从而得到两个矩阵Add后的结果。

 template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE>
 __aicore__ inline void MatMulKernel(...)
 {
    ...
    Matmul<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_MDL> mm;
    TPipe pipe;
    REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm);

    mm.SetTensorA(gm_a);
    mm.SetTensorB(gm_b);
    mm.SetBias(gm_bias);

    mm.IterateAll(gm_d, 1); // IterateAll接口中的enAtomic设为1
    
    // while (mm. Iterate ()) {
        // mm.GetTensorC(gm_d, 1);     // GetTensorC接口中的enAtomic设为1
    // }
    ...
 }

 extern "C" __global__ __aicore__ void example_kernel(...)
 {
     ...
     typedef MatmulType<TPosition::GM, CubeFormat::ND, half> aType; 
     typedef MatmulType<TPosition::GM, CubeFormat::ND, half> bType; 
     typedef MatmulType<TPosition::GM, CubeFormat::ND, float> cType; 
     typedef MatmulType<TPosition::GM, CubeFormat::ND, float> biasType;
     MatMulKernel<aType, bType, cType, biasType)(...);
     ...
 }

【性能对比】

图1 Matmul使能AtomicAdd选项前后性能对比

以矩阵维度M=64,N=256,K=256,矩阵D为(64, 256)为例,Matmul使能AtomicAdd选项前后的性能对比如上图所示,平均cycle数从开启AtomicAdd选项前的154181变为开启后的135054,性能优化12.4%。因此在这种场景下,使能AtomicAdd选项能获取更优的性能。