算子实现

实现流程

上文介绍了Matmul矩阵乘的数据切分方案和数据流。Ascend C提供一组Matmul高阶API，封装了这些常用的切分和数据搬运、计算的算法逻辑，方便用户快速实现Matmul矩阵乘法的运算操作。开发者在host侧通过调用API自动获取Tiling参数，该参数传递到kernel侧后，在初始化操作时传入，通过几个简单的API即可完成矩阵乘操作。完整样例请参考LINK。

图1 矩阵编程流程示意图

host侧自动获取Tiling参数的关键步骤介绍如下：

1. 创建Tiling对象。

   Text

   auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
   matmul_tiling::MultiCoreMatmulTiling tilingApi(*ascendcPlatform);
   

   传入硬件平台信息创建PlatformAscendC对象，然后创建Tiling对象，硬件平台信息可以通过GetPlatformInfo获取。

2. 设置参与Matmul运算的核数，A、B的内存逻辑位置、格式和数据类型。

   Text

   tilingApi.SetDim(ascendcPlatform->GetCoreNumAic());  
   tilingApi.SetAType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16, false);
   tilingApi.SetBType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16, false);
   tilingApi.SetCType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
   

3. 设置矩阵shape信息。

   Text

   tilingApi.SetOrgShape(M, N, K); // 设置原始完整的形状M、N、K
   tilingApi.SetShape(M, N, K);
   

4. 设置可用空间大小信息。

   设置Matmul计算时可用的L1 Buffer/L0C Buffer/Unified Buffer（UB）空间大小，-1表示AI处理器对应Buffer的大小。

   Text

   tilingApi.SetBufferSpace(-1, -1, -1);
   

5. 获取Tiling参数。

   Text

   int64_t res = tilingApi.GetTiling(tilingData);
   if (res == -1) {
       std::cout << "gen tiling failed" << std::endl;
   }
   

6. Tiling参数的序列化保存等其他操作。

   Text

   uint32_t tcubeTilingSize = tilingData.GetDataSize();
   tilingData.SaveToBuffer(tilingBuf, tcubeTilingSize);
   

kernel侧使用Matmul API矩阵乘运算的具体步骤如下：

1. 创建Matmul对象。

   创建Matmul对象的示例如下：

   • 纯Cube模式（只有矩阵计算）场景下，建议在代码中定义ASCENDC_CUBE_ONLY宏，避免额外的性能开销。本节内容以纯Cube模式举例。
   • 默认为MIX模式（包含矩阵计算和矢量计算），该场景下通常不定义ASCENDC_CUBE_ONLY宏，如果在程序中使用了ASCENDC_CUBE_ONLY宏，则必须使用ASCEND_IS_AIC宏和ASCEND_IS_AIV宏将Cube计算和Vector计算隔离开，更多内容请参考融合算子编程。
   Text

   typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half> aType; 
   typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half> bType; 
   typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float> cType; 
   AscendC::Matmul<aType, bType, cType> mm;
   

   创建对象时需要传入A、B、C的参数类型信息，类型信息通过MatmulType来定义，包括：内存逻辑位置、数据格式、数据类型。

2. 初始化操作。

   Text

   REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling); // 初始化
   

   说明 Matmul高阶API内部实现时需要使用系统workspace（即对应本步骤中的GetSysWorkSpacePtr接口），开发者需要自行申请系统workspace的空间：

   • 在host侧Tiling实现时，设置总的workspace的数值大小（包含用户workspace和系统workspace），workspace空间由框架来申请并管理。系统workspace的空间大小通过GetLibApiWorkSpaceSize获取。

   Text

   size_t userWorkspaceSize = 0;
   size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform.GetLibApiWorkSpaceSize());
   size_t workspaceSize = userWorkspaceSize + systemWorkspaceSize;
   

3. 设置原始矩阵shape以及当前核使用的左矩阵A、右矩阵B。

   多核场景下，需要根据数据切分方向为每个核设置对应的矩阵起始地址。本样例中沿M轴切分A矩阵和C矩阵：每个核读取A矩阵中singleCoreM行数据，B矩阵不切分，由所有核共同读取完整B矩阵。SetOrgShape需要在SetTensorA、SetTensorB之前调用。

   Text

   mm.SetOrgShape(tiling.M, tiling.N, tiling.Ka, tiling.Kb);
   mm.SetTensorA(aGlobal[GetBlockIdx() * tiling.singleCoreM * tiling.Ka], false);
   mm.SetTensorB(bGlobal[0], false);
   

   其中，GetBlockIdx()表示当前核号。A矩阵按[M, Ka]连续存放，因此当前核的A矩阵起始偏移为GetBlockIdx() * tiling.singleCoreM * tiling.Ka。各核计算C矩阵的不同行，但都需要完整的B矩阵参与K轴累加，因此B矩阵从bGlobal[0]开始读取。SetTensorA和SetTensorB的第二个参数false表示矩阵不转置。

4. 完成矩阵乘操作。

   • 调用Iterate完成单次迭代计算，叠加while循环完成当前核负责的数据计算。C矩阵按[M, N]连续存放，当前核的写回起始偏移为GetBlockIdx() * tiling.singleCoreM * tiling.N。Iterate方式可以自行控制迭代次数，使用方式比较灵活。

     Text

     while (mm.Iterate()) {
         mm.GetTensorC(cGlobal[GetBlockIdx() * tiling.singleCoreM * tiling.N]);
     }
     

   • 调用IterateAll完成当前核负责的全部数据计算并写回C矩阵。IterateAll方式无需循环迭代，使用比较简单。本节样例代码中调用IterateAll实现矩阵乘。

     Text

     mm.IterateAll(cGlobal[GetBlockIdx() * tiling.singleCoreM * tiling.N]);
     

5. 结束矩阵乘操作。

   Text

   mm.End();
   

设置Shape信息

在实现Host Tiling时可以设置Shape信息，用于Tiling计算；kernel侧运行时也可以修改部分Shape信息，用于尾块设置、Matmul复用（多个Matmul计算复用一个Matmul对象）等场景。本节对涉及到的Shape概念进行介绍，并给出host侧和kernel侧设置Tiling信息的指导。

• orgShape：M、N、K
• singleCoreShape：singleCoreM、singleCoreN、singleCoreK
• singleShape：singleM、singleN、singleK
• baseShape：baseM、baseN、baseK

通过数据分块（Tiling）的介绍我们已经了解了orgShape(M、N、K)，singleCoreShape(singleCoreM、singleCoreN、singleCoreK)，baseShape(baseM、baseN、baseK)的概念，如下图所示：

除此之外，单核的Matmul Tiling时，实际参与Matmul计算的shape可以是原始shape中的一部分，singleM, singleN, singleK用于表达实际参与Matmul计算的shape，如下图所示。在单核的情况下，singleM, singleN, singleK会透传给singleCoreM, singleCoreN, singleCoreK。

• 核函数（Kernel）运行时设置

  • SetTail、SetSingleShape都是运行时修改singleCoreM、singleCoreN、singleCoreK，处理尾块时使用SetTail，Matmul复用（多个Matmul计算复用一个Matmul对象）的场景可以使用SetSingleShape重新设置。
  • SetOrgShape是运行时修改M、N、K，Matmul复用的场景可以使用SetOrgShape重新设置。

• 单核Tiling时设置

  • SetOrgShape（必选）：设置M、N、K
  • SetShape（非必选）：设置singleM、singleN、singleK，等同于设置singleCoreM、singleCoreN、singleCoreK
  • SetFixSplit（非必选）：设置baseM、baseN、baseK

• 多核Tiling时设置

  • SetOrgShape（必选）：设置M、N、K

  • SetShape（非必选）：设置singleM、singleN、singleK

  • SetFixSplit（非必选）：设置baseM、baseN、baseK

  • SetSingleShape(非必选)：设置singleCoreM、singleCoreN、singleCoreK

  • SetSingleRange(非必选)：设置singleCoreM、singleCoreN、singleCoreK的范围

设置format格式

创建Matmul对象时需要传入A、B、C的参数类型信息，类型信息通过MatmulType来定义，包括：内存逻辑位置、数据格式、数据类型。示例如下：

Text

typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half> aType; 
typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half> bType; 
typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float> cType; 
AscendC::Matmul<aType, bType, cType> mm;

针对数据格式，包括CubeFormat::ND, CubeFormat::NZ, CubeFormat::ND_ALIGN三种，ND和NZ格式在数据格式章节已经介绍，ND_ALIGN格式的介绍请参考数据排布格式。