使用说明

CubeResGroupHandle用于在分离架构中通过软同步控制AIC和AIV之间进行通讯，实现AI Core计算资源分组。CubeResGroupHandle可以将一个AIC与多个消息队列分为一组，每个消息队列和一个AIV进行绑定，在GM侧分配一块内存存储AIV发送到AIC的通讯消息，并按照类似于表格的结构组织被存储的消息。

如下图所示，一个CubeResGroupHandle中可以有一个或多个AIC，例如图中Block0、Block1，每个AIC能与多个不同的消息队列进行通讯，每条消息队列存储同一个AIV发送的消息，不同AIV的消息为不同的消息队列；不同CubeResGroupHandle中的不同AIC可以与同一个AIV进行通讯。消息队列的深度固定为4。下图中，Block 0与Queue 0，Queue 1，Queue 2，Queue3，Queue4进行通信；Block 1与Queue 5，Queue 6，Queue 7，Queue 8，Queue9进行通信。CubeResGroupHandle1和CubeResGroupHandle2的消息队列个数分别为10和12。AIC遍历读取每个消息队列的消息，并执行对应任务，同一条消息队列根据先进先出的顺序处理消息，消息处理顺序如图中箭头所示。

图1 基于CubeResGroupHandle的AI Core计算资源分组通信示意图

基于CubeResGroupHandle实现AI Core计算资源分组步骤如下：

  1. 创建AIC上计算所需要的计算对象类型。
  2. 创建通信区域描述KfcWorkspace，用于记录通信消息Msg的地址分配。
  3. 自定义消息结构体，用于通信。
  4. 自定义回调计算逻辑，根据实际业务场景实现Init函数和Call函数。
  5. 创建CubeResGroupHandle。
  6. 绑定AIV到CubeResGroupHandle。
  7. 收发消息。
  8. AIV退出消息队列。
  1. 创建AIC上计算所需要的计算对象类型。

    用户根据实际需求，自定义AIC所需要的计算对象类型，或者高阶API已提供的Matmul类型。例如，创建Matmul类型如下，其中A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM等含义请参考Matmul模板参数。

    1
2
3
4
5
6
    		 
    // A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM根据实际需求场景构造
using qkType = MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM>;
using dyvType = MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM>;
using dqType = MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM>;
using dkType = MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM>;
using qvType = MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM>;

  2. 创建KfcWorkspace。

    使用KfcWorkspace管理不同CubeResGrouphandle的消息通信区的划分。 
    1
2
    		 
    // 创建KfcWorkspace对象前，需要对该workspace清零
AscendC::KfcWorkspace commMem(workspace);  //创建1个通信区域描述，用于记录通信消息Msg的地址分配。

  3. 自定义消息结构体。

    用户需要自行构造消息结构体CubeMsgBody，用于AIV向AIC发送通信消息。构造的CubeMsgBody必须64字节对齐，该结构体最前面需要定义2字节的CubeGroupMsgHead，使消息收发机制正常运行，CubeGroupMsgHead结构定义请参考表2。除2字节的CubeGroupMsgHead外，其余参数根据业务需求自行构造。
    表1 CubeMsgBody消息结构体

    参数名称

    含义

    CubeMsgBody

    用户自定义的消息结构体。结构体名称可自定义，结构体大小需要64字节对齐。

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
    		 
    // 这里提供64B对齐的结构体示例，用户实际使用时，除CubeGroupMsgHead外，其他参数个数及参数类型可自行构造
struct CubeMsgBody {
   CubeGroupMsgHead head;  // 2B，需放在结构体最前面, 自定义的CubeMsgBody中，CubeGroupMsgHead的变量名需设置为head，否则会编译报错。
   uint8_t funcID;
   uint8_t skipCnt;
   uint32_t value;
   bool isTransA;
   bool isTransB;
   bool isAtomic;
   bool isLast;                 
   int32_t tailM;              
   int32_t tailN;
   int32_t tailK;               
   uint64_t aAddr;
   uint64_t bAddr;
   uint64_t cAddr;
   uint64_t aGap;
   uint64_t bGap;
}
    表2 CubeGroupMsgHead结构体参数定义

    参数名称

    含义

    msgState

    表明该位置的消息状态。参数取值如下：

    • CubeMsgState::FREE：表明该位置还未填写消息，可执行AllocMessage。
    • CubeMsgState::VALID：表明该位置已经含有AIV发送的消息，待AIC接收执行。
    • CubeMsgState::QUIT：表明该位置的消息为通知AIC有AIV将退出流程。
    • CubeMsgState::FAKE：表明该位置的消息为假消息。在消息合并场景，被跳过处理任务的AIV需要发送假消息，消息合并场景请参考PostFakeMsg中的介绍。

    aivID

    发送消息的AIV的序号。

  4. 自定义回调计算逻辑结构体，根据实际业务场景实现Init函数和Call函数。 

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
    		 
    template<class MatmulApiCfg, class CubeMsgBody>
struct NormalCallbackFuncs {
    __aicore__ inline static void Call(MatmulApiCfg &mm, __gm__ CubeMsgBody *rcvMsg, CubeResGroupHandle<CubeMsgBody> &handle){
      // 用户自己实现逻辑
    };

    __aicore__ inline static void Init(NormalCallbackFuncs<MatmulApiCfg, CubeMsgBody> &foo, MatmulApiCfg &mm, GM_ADDR tilingGM){
       // 用户自己实现逻辑
    };
   
};

    计算逻辑结构体的模板参数请参考表3。

    表3 模板参数说明

    参数

    说明

    MatmulApiCfg

    用户自定义的AIC上计算所需要对象的数据类型，参考步骤1，该模板参数必须填入。

    CubeMsgBody

    用户自定义的消息结构体，该模板参数必须填入。

    用户自定义回调计算逻辑的结构体中需要包含固定的Init函数和Call函数，函数原型如下所示。其中，Init函数的参数说明请参考表4，Call函数的参数说明请参考表5。

    1
2
3
4
    		 
    // 该函数的参数和名称为固定格式，函数实现根据业务逻辑自行实现。
__aicore__ inline static void Init(MyCallbackFunc<MatmulApiCfg, CubeMsgBody> &myCallBack, MatmulApiCfg &mm, GM_ADDR tilingGM){
     // 用户自行实现内部逻辑
    }
    表4 Init函数参数说明

    参数

    输入/输出

    说明

    myCallBack

    输入

    用户自定义的带模板参数的回调计算结构体。

    mm

    输入

    AIC上计算对象，多为Matmul对象。

    tilingGM

    输入

    用户传入的tiling指针。

    		 
    1
2
3
4
    		 
    // 该函数的参数和名称为固定格式，函数实现根据业务逻辑自行实现。
__aicore__ inline static void Call(MatmulApiCfg &mm, __gm__ CubeMsgBody *rcvMsg, CubeResGroupHandle<CubeMsgBody> &handle){
        // 用户自行实现内部逻辑
    }
    表5 Call函数参数说明

    参数

    输入/输出

    说明

    mm

    输入

    AIC上计算对象，多为Matmul对象。

    rcvMsg

    输入

    用户自定义的消息结构体指针。

    handle

    输入

    组消息管理Handle，用户调用其接口进行收发消息，释放消息等。
    某算子的回调计算结构体的代码示例如下。 
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    		 
    // 用户自定义的回调计算逻辑
template<class MatmulApiCfg, class CubeMsgBody>
struct NormalCallbackFuncs {
    template<int32_t funcId>
    __aicore__ inline static typename IsEqual<funcId, static_cast<int32_t>(AscendC::FuncId::MM_NORMAL)>::Type CubeGroupCallBack(
        MatmulApiCfg &mm, __gm__ CubeMsgBody* rcvMsg, CubeResGroupHandle<CubeMsgBody> &handle) {

      // 如果自定义消息结构体超过64B需要调用DataCacheCleanAndInvalid自行刷新标志位，不超过64B则不需要。
      AscendC::GlobalTensor<int64_t> msgGlobal;
      msgGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(rcvMsg) + sizeof(int64_t));
      AscendC::DataCacheCleanAndInvalid<int64_t, AscendC::CacheLine::SINGLE_CACHE_LINE, AscendC::DcciDst::CACHELINE_OUT>(msgGlobal);

      mm.SetOrgShape(rcvMsg->mOrgShape, rcvMsg->nOrgShape, rcvMsg->kaOrgShape, rcvMsg->kbOrgShape, rcvMsg->kcOrgShape);
      mm.SetTail(rcvMsg->tailM, rcvMsg->tailN, rcvMsg->tailK);
      using aType = typename MatmulApiCfg::AType::T;
      using bType = typename MatmulApiCfg::BType::T;
      using cType = typename MatmulApiCfg::CType::T;
      AscendC::GlobalTensor<aType> aGlobal;
      AscendC::GlobalTensor<bType> bGlobal;
      AscendC::GlobalTensor<cType> cGlobal;
      event_t eventId = static_cast<event_t>(AscendC::GetTPipePtr()->FetchEventID(AscendC::HardEvent::FIX_S));

      int64_t skipCount = rcvMsg->skipCnt;
      for (int64_t i = 0; i < skipCount + 1; i++) {
        if (i == skipCount) {
          mm.SetTail(rcvMsg->lastTailM, rcvMsg->tailN, rcvMsg->tailK);
        }
        aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ aType*>(rcvMsg->aAddr + i * rcvMsg->aGap * sizeof(aType)));
        bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ bType*>(rcvMsg->bAddr + i * rcvMsg->bGap * sizeof(bType)));
        cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ cType*>(rcvMsg->cAddr + i * rcvMsg->cGap * sizeof(cType)));
        mm.SetTensorA(aGlobal, rcvMsg->isTransA);
        mm.SetTensorB(bGlobal, rcvMsg->isTransB); 
        mm.IterateAll(cGlobal, rcvMsg->isAtomic);
        mm.End();

        AscendC::SetFlag<AscendC::HardEvent::FIX_S>(eventId);
        AscendC::WaitFlag<AscendC::HardEvent::FIX_S>(eventId);
        AscendC::CubeMsgState waitState = i > 0 ? AscendC::CubeMsgState::FAKE : AscendC::CubeMsgState::VALID;
        handle.FreeMessage(rcvMsg + i, waitState); // AIC相应的计算结束后，需要用户自行将消息释放。
      }
      handle.SetSkipMsg(skipCount);// 在消息合并场景需要调用该接口。
    };

    __aicore__ inline static void Call(MatmulApiCfg &mm, __gm__ CubeMsgBody *rcvMsg, CubeResGroupHandle<CubeMsgBody> &handle){
      CubeGroupCallBack<static_cast<int32_t>(AscendC::FuncId::MM_NORMAL)>(mm, rcvMsg, handle);
    };

    __aicore__ inline static void Init(NormalCallbackFuncs<MatmulApiCfg, CubeMsgBody> &foo, MatmulApiCfg &mm, GM_ADDR tilingGM){
        GET_TILING_DATA_WITH_STRUCT(TilingDataS1s2Gs1Bn2s2, tiling_data_in, tilingGM);
        foo.tilingData = tiling_data_in;
        mm.SetSubBlockIdx(0);
        mm.Init(&(foo.tilingData.mm1TilingData), GetTPipePtr());
    };
    TilingDataS1s2Gs1Bn2s2 tilingData = {};// 用户自定义的在AIC上需要使用到的Tiling结构体。
};

  5. 创建CubeResGroupHandle。

    用户使用CreateCubeResGroup接口创建一个或多个CubeResGroupHandle，同时每个CubeResGroupHandle需要绑定一个AIC上计算所需要的对象。 
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
    		 
    /* 
 * qkType为步骤1定义好的AIC上计算所需要的对象的类型
 * MyCallbackFunc为为步骤4定义好的自定义回调计算逻辑结构体
 * CubeMsgBody为步骤3自定义消息结构体
 * commMem为步骤2用户初始化好的通信区域描述
 * blockStart为0，blockSize为12，msgQueueSize为48，tilingGm为指针，存储了用户在AIC上所需要的tiling信息
/
auto qk =  AscendC::CreateCubeResGroup<1, qkType, MyCallbackFunc, CubeMsgBody>(commMem, 0, 12, 48, tilingGM);
// 创建CubeResGroupHandle后，用户需要自行退出AIC逻辑，可参考下方代码进行实现。
if ASCEND_IS_AIC {
   return；
}

  6. 绑定AIV到CubeResGroupHandle。

    绑定AIV和消息队列序号。注意：消息队列序号queIdx小于该CubeGroupHandle的消息队列总数，每个AIV需要传入不同的queIdx。 
    1
2
    		 
    // qk为步骤5中CreateCubeResGroup创建的CubeResGroupHandle对象
qk.AssignQueue(queIdx);

  7. AIV发消息。

    用户调用AllocMessage, PostMessage等接口进行消息的收发。其中，调用AllocMessage获取消息结构体指针，通过PostMessage发送消息，在消息合并场景调用PostFakeMessage发送假消息，示例如下。 
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
    		 
    AscendC::CubeGroupMsgHead headA = {AscendC::CubeMsgState::VALID, BlockId}; // 真消息需要设置CubeMsgState状态为VALID，并记录当前AIV的核号
AscendC::CubeGroupMsgHead headB = {AscendC::CubeMsgState::VALID, BlockId};
AscendC::CubeMsgBody a {headA,1,0,0, false, false, false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 3, 4, false, 0, 0};
AscendC::CubeMsgBody b {headB,1,0,0, false, false, false, false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 3, 4, false, 0, 0};
if (AscendC::GetBlockIdx() == 0) {
   auto msg = qk.template AllocMessage();
   id = qk.template PostMessage(msg, b);
   bool waitState = qk.template Wait<true>(id);
}else if (AscendC::GetBlockIdx() < 4) {
   auto msg = qk.AllocMessage(); 
   id = qk.PostFakeMsg(msg); // 三个消息被合并，故发假消息
   bool waitState = qk.template Wait<true>(id); 
} else {
   auto msg = qk.template AllocMessage();
   id = qk.template PostMessage(msg, a);
   bool waitState = qk.template Wait<true>(id);
}

  8. AIV退出消息队列。

    调用AllocMessage获取消息结构体指针后，通过SendQuitMsg发送当前消息队列退出。 
    1
2
    		 
    auto msgPtr = a.AllocMessage();        // 获取消息空间指针msgPtr
a.SendQuitMsg(msgPtr);              // 发送退出消息

父主题： CubeResGroupHandle