IterateBatch

功能说明

单次Matmul计算处理的shape比较小时，由于每次计算均涉及到内部的通信，可能会影响性能，该接口提供批量处理Matmul的功能，调用一次IterateBatch，可以计算出多个singleCoreM * singleCoreN大小的C矩阵。

在使用该接口前，需要了解一些必备的数据排布格式：

• 通用数据格式：BMNK的数据排布格式

  

• BSH/SBH：B：Batch，批处理的大小； S：sequence length，序列长度；H = N * D，其中，N为head的数量，D为head的大小。Layout格式如下图所示：

  

  

• BSNGD：为原始BSH shape做reshape后的shape，S和D为单Batch的矩阵乘的M轴（或N轴）和K轴，一个SD为一个batch的计算数据，Layout格式如下图所示：

  

• SBNGD：为原始SBH shape做reshape后shape，S和D为的矩阵乘的M轴（或N轴）和K轴，一个SD为一个Batch的计算数据，Layout格式如下图所示：

  

• BNGS1S2：一般为前两种Layout进行矩阵乘的输出，S1S2数据连续存放，一个S1S2为一个Batch的计算数据，Layout格式如下图所示：

  

实例化Matmul时，需要通过MatmulType设置输入输出的Layout格式，当前支持4种Layout类型：BSNGD、SBNGD、BNGS1S2、NORMAL（BMNK的数据排布格式使用NORMAL表示）。

对于BSNGD、SBNGD、BNGS1S2 Layout格式，调用该接口之前需要在host Tiling实现中使用SetALayout、SetBLayout、SetCLayout、SetBatchNum设置A/B/C的Layout轴信息和最大BatchNum数；对于NORMAL Layout格式则需使用SetBatchInfoForNormal设置A/B/C的M/N/K轴信息和A/B矩阵的BatchNum数。

单个矩阵乘迭代顺序可通过tiling参数iterateOrder调整。

更多矩阵编程batch场景的相关内容请参考batch场景。

函数原型

• mix模式
  • 输出至GM

    1 2

    template <bool sync = true, bool waitIterateBatch = false> __aicore__ inline void IterateBatch(const GlobalTensor<DstT>& gm, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0, const bool enPartialSum = false, const uint8_t enAtomic = 0)

  • 输出至VECIN

    1 2

    template <bool sync = true> __aicore__ inline void IterateBatch(const LocalTensor<DstT>& ubCmatrix, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0, const bool enPartialSum = false, const uint8_t enAtomic = 0)

    使用前需先调用SetBatchNum接口设置batchA和batchB的大小。

• 纯cube模式
  • 输出至GM

    1	__aicore__ inline void IterateBatch(const GlobalTensor<DstT>& gm, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0)

  • 输出至VECIN

    1	__aicore__ inline void IterateBatch(const LocalTensor<DstT>& ubCmatrix, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0)

参数说明

返回值

无

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

Atlas推理系列产品AI Core

约束说明

• 该接口接口只支持Norm模板，即BatchMatmul只支持Norm模板。
• 对于BSNGD、SBNGD、BNGS1S2 Layout格式，输入A、B矩阵按分形对齐后的多Batch数据总和应小于L1 Buffer的大小；对于NORMAL Layout格式没有这种限制，但需通过MatmulConfig配置输入A、B矩阵多Batch数据大小与L1 Buffer的大小关系。
• 如果接口输出到Unified Buffer上，输出C矩阵大小BaseM*BaseN应小于分配的Unified Buffer内存大小。
• 对于BSNGD、SBNGD、BNGS1S2 Layout格式，输入输出只支持ND格式数据。对于NORMAL Layout格式， 输入支持ND/NZ格式数据。
• 该接口不支持量化模式，即不支持SetQuantScalar、SetQuantVector接口。
• BSNGD场景，不支持一次计算多行SD，需要算子程序中循环计算。
• 异步模式不支持IterateBatch搬运到UB上。
• Atlas推理系列产品AI Core，只支持NORMAL Layout格式。
• 在Atlas推理系列产品AI Core上，不支持A、B矩阵内存逻辑位置为TPosition::TSCM的输入。
• Atlas推理系列产品AI Core上，Bias不支持复用，Bias的shape大小必须为Batch * N。

调用示例

• 该示例完成aGM、bGM矩阵乘，结果保存到cGm上，其中aGM、bGM、cGM数据的layout格式均为NORMAL，左矩阵每次计算batchA个MK数据，右矩阵每次计算batchB个KN数据。

  #include "kernel_operator.h"
  #include "lib/matmul_intf.h"
  
  extern "C" __global__  __aicore__ void kernel_matmul_rpc_batch(GM_ADDR aGM, GM_ADDR bGM, GM_ADDR cGM, GM_ADDR biasGM, GM_ADDR tilingGM, GM_ADDR workspaceGM, uint32_t isTransposeAIn, uint32_t isTransposeBIn, int32_t batchA,  int32_t batchB)
  {
      // 定义matmul type
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, half, false, LayoutMode::NORMAL> aType;
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, half, true, LayoutMode::NORMAL> bType;
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, float, false, LayoutMode::NORMAL> cType;
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, float> biasType;
  
      // 初始化tiling数据
      TCubeTiling tiling;
      auto tempTilingGM = (__gm__ uint32_t*)tilingGM;
      auto tempTiling = (uint32_t*)&tiling;
      for (int i = 0; i < sizeof(TCubeTiling) / sizeof(int32_t); ++i, ++tempTilingGM, ++tempTiling) {
          *tempTiling = *tempTilingGM;
      }
  
      // 初始化gm数据
      AscendC::GlobalTensor<half> aGlobal;
      AscendC::GlobalTensor<half> bGlobal;
      AscendC::GlobalTensor<float> cGlobal;
      AscendC::GlobalTensor<float> biasGlobal;
      int32_t sizeA = tiling.ALayoutInfoB * tiling.singleCoreM * tiling.singleCoreK * sizeof(A_T);
      int32_t sizeB = tiling.BLayoutInfoB * tiling.singleCoreK * tiling.singleCoreN * sizeof(B_T);
      int32_t sizeC = tiling.CLayoutInfoB * tiling.singleCoreM * tiling.singleCoreN * sizeof(C_T);
      int32_t sizebias = tiling.CLayoutInfoB * tiling.singleCoreN * sizeof(C_T);
      aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(aGM), sizeA);
      bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(bGM), sizeB);
      cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(cGM), sizeC);
      biasGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(biasGM), sizebias);
      tiling.shareMode = 0;            
      tiling.shareL1Size = 512 * 1024;
      tiling.shareL0CSize = 128 * 1024;
      tiling.shareUbSize = 0;
      int offset_a=0, offset_b=0, offset_c=0, offset_bias=0;
      AscendC::GlobalTensor<A_T> gm_a;
      gm_a.SetGlobalBuffer(const_cast<__gm__ A_T*>(aGlobal[offset_a].GetPhyAddr()), tiling.singleCoreM * tiling.singleCoreK);
      AscendC::GlobalTensor<B_T> gm_b;
      gm_b.SetGlobalBuffer(const_cast<__gm__ B_T*>(bGlobal[offset_b].GetPhyAddr()), tiling.singleCoreK * tiling.singleCoreN);
      AscendC::GlobalTensor<C_T> gm_c;
      gm_c.SetGlobalBuffer(const_cast<__gm__ C_T*>(cGlobal[offset_c].GetPhyAddr()), tiling.singleCoreM * tiling.singleCoreN) ;
      AscendC::GlobalTensor<BiasT> gm_bias;
      gm_bias.SetGlobalBuffer(const_cast<__gm__ BiasT*>(biasGlobal[offset_bias].GetPhyAddr()), tiling.singleCoreN);
      // 创建Matmul实例
      constexpr static MatmulConfig MM_CFG = GetNormalConfig(false, false, false, BatchMode::BATCH_LESS_THAN_L1);
      matmul::Matmul<aType, bType, cType, biasType, MM_CFG> mm1;
      TPipe pipe;
  
      REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm1);
      mm1.Init(&tiling);
      mm1.SetTensorA(gm_a, isTransposeAIn);
      mm1.SetTensorB(gm_b, isTransposeBIn);
      mm1.SetWorkspace(workspaceGM, 0);
      if(tiling.isBias) {
          mm1.SetBias(gm_bias);
      }
      // 多batch Matmul计算
      mm1.IterateBatch(gm_c, batchA, batchB, false);
  }

• 该示例完成aGM、bGM矩阵乘，结果保存到cGm上，其中aGM数据的layout格式为BSNGD，bGM数据的layout格式为BSNGD，cGM的layout格式为BNGS1S2，左矩阵每次计算batchA个SD数据，右矩阵每次计算batchB个SD数据。

  #include "kernel_operator.h"
  #include "lib/matmul_intf.h"
  
  extern "C" __global__  __aicore__ void kernel_matmul_rpc_batch(GM_ADDR aGM, GM_ADDR bGM, GM_ADDR cGM, GM_ADDR biasGM, GM_ADDR tilingGM, GM_ADDR workspaceGM, uint32_t isTransposeAIn, uint32_t isTransposeBIn, int32_t batchA,  int32_t batchB)
  {
      // 定义matmul type
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, half, false, LayoutMode::BSNGD> aType;
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, half, true, LayoutMode::BSNGD> bType;
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, float, false, LayoutMode::BNGS1S2> cType;
      typedef matmul::MatmulType <AscendC::TPosition::GM, CubeFormat::ND, float> biasType;
  
      // 初始化tiling数据
      TCubeTiling tiling;
      auto tempTilingGM = (__gm__ uint32_t*)tilingGM;
      auto tempTiling = (uint32_t*)&tiling;
      for (int i = 0; i < sizeof(TCubeTiling) / sizeof(int32_t); ++i, ++tempTilingGM, ++tempTiling) {
          *tempTiling = *tempTilingGM;
      }
  
      // 初始化gm数据
      AscendC::GlobalTensor<half> aGlobal;
      AscendC::GlobalTensor<half> bGlobal;
      AscendC::GlobalTensor<float> cGlobal;
      AscendC::GlobalTensor<float> biasGlobal;
      int32_t sizeA = tiling.ALayoutInfoB * tiling.ALayoutInfoS * tiling.ALayoutInfoN * tiling.ALayoutInfoG * tiling.ALayoutInfoD * sizeof(half);
      int32_t sizeB = tiling.BLayoutInfoB * tiling.BLayoutInfoS * tiling.BLayoutInfoN * tiling.BLayoutInfoG * tiling.BLayoutInfoD * sizeof(half);
      int32_t sizeC = tiling.CLayoutInfoB * tiling.CLayoutInfoS1 * tiling.CLayoutInfoN * tiling.CLayoutInfoG * tiling.CLayoutInfoS2 * sizeof(float);
      int32_t sizebias = tiling.CLayoutInfoB * tiling.CLayoutInfoN * tiling.CLayoutInfoG * tiling.CLayoutInfoS2 * sizeof(float);
      aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(aGM), sizeA);
      bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ half*>(bGM), sizeB);
      cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(cGM), sizeC);
      biasGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(biasGM), sizebias);
      tiling.shareMode = 0;            
      tiling.shareL1Size = 512 * 1024;
      tiling.shareL0CSize = 128 * 1024;
      tiling.shareUbSize = 0;
      int offset_a=0, offset_b=0, offset_c=0, offset_bias=0;
      AscendC::GlobalTensor<A_T> gm_a;
      gm_a.SetGlobalBuffer(const_cast<__gm__ half*>(aGlobal[offset_a].GetPhyAddr()), tiling.ALayoutInfoS * tiling.ALayoutInfoN * tiling.ALayoutInfoG * tiling.ALayoutInfoD);
      AscendC::GlobalTensor<B_T> gm_b;
      gm_b.SetGlobalBuffer(const_cast<__gm__ half*>(bGlobal[offset_b].GetPhyAddr()), tiling.BLayoutInfoS * tiling.BLayoutInfoN * tiling.BLayoutInfoG * tiling.BLayoutInfoD);
      AscendC::GlobalTensor<C_T> gm_c;
      gm_c.SetGlobalBuffer(const_cast<__gm__ float*>(cGlobal[offset_c].GetPhyAddr()), tiling.CLayoutInfoS1 * tiling.CLayoutInfoN * tiling.CLayoutInfoG * tiling.CLayoutInfoS2) ;
      AscendC::GlobalTensor<BiasT> gm_bias;
      gm_bias.SetGlobalBuffer(const_cast<__gm__ float*>(biasGlobal[offset_bias].GetPhyAddr()), tiling.CLayoutInfoN * tiling.CLayoutInfoG * tiling.CLayoutInfoS2);
      // 创建Matmul实例
      matmul::Matmul<aType, bType, cType, biasType> mm1;
      AscendC::TPipe pipe;
  
      REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm1);
      mm1.Init(&tiling);
      mm1.SetWorkspace(workspaceGM, 0);
      int batchC = batchA > batchB ? batchA : batchB;
      int g_lay = tiling.ALayoutInfoG > tiling.BLayoutInfoG ? tiling.ALayoutInfoG : tiling.BLayoutInfoG;
      // 计算需要多Batch计算循环次数
      int for_exent = tiling.ALayoutInfoB * tiling.ALayoutInfoN * g_lay / tiling.BatchNum;
      for(int i=0; i<for_exent; ++i) {
          // 计算每次多batch计算A/B矩阵的起始地址
          int batchOffsetA = i * tiling.ALayoutInfoD * batchA;
          int batchOffsetB = i * tiling.BLayoutInfoD * batchB;
          mm1.SetTensorA(gm_a[batchOffsetA], isTransposeAIn);
          mm1.SetTensorB(gm_b[batchOffsetB], isTransposeBIn);
          int idx_c = i * batchC;
          if (tiling.CLayoutInfoG == 1 && (tiling.BLayoutInfoG != 1 || tiling.ALayoutInfoG != 1)) {
              idx_c = idx_c / (tiling.BLayoutInfoG > tiling.ALayoutInfoG ? tiling.BLayoutInfoG : tiling.ALayoutInfoG);
          }
          if(tiling.isBias) {
              int batchOffsetBias = idx_c * tiling.CLayoutInfoS2;
              mm1.SetBias(gm_bias[batchOffsetBias]);
          }
          int batchOffsetC = idx_c * tiling.CLayoutInfoS2;
          if (C_TYPE::layout == LayoutMode::BNGS1S2) {
              batchOffsetC = idx_c * tiling.CLayoutInfoS2 * tiling.CLayoutInfoS1;
          }
          // 多batch Matmul计算
          mm1.IterateBatch(gm_c[batchOffsetC], batchA, batchB, false);
      }
  }