Wait

Applicability

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	x
Atlas inference product's AI Core	x
Atlas inference product's Vector Core	x
Atlas training products	x

Function

Blocks the AI Core. When fine-grained communication is not used, wait for the communication task corresponding to handleId to complete. During fine-grained communication, wait for the sub-communication task with the step corresponding to handleId to complete. The sequence of calling the Wait API by handleId must be the same as that of calling the Prepare API. By default, this API works on all cores. You can also specify a core by setting GetBlockIdx before calling this API.

Prototype

__aicore__ inline int32_t Wait(HcclHandle handleId)

Parameters

Table 1 API parameters

Parameter

Input/Output

Description

handleId

Input

ID of the corresponding communication task. Only the return value of the API corresponding to the Prepare primitive can be used.

using HcclHandle = int8_t;

Returns

0: success
–1: failure

Restrictions

Before calling this API, ensure that the InitV2 and SetCcTilingV2 APIs have been called.
The input parameter handleId can use only the return value of the Prepare primitive API.
When this API is called on the AIC or AIV core, the calling core must be the same as that of the corresponding Prepare API.
When fine-grained communication is not used, the number of times this API is called must be the same as the number of times the Prepare API is called. In fine-grained communication scenarios, the number of times this API is called must be the same as the total number of steps in the communication task/step size x the number of times the Prepare API is called. The sequence of calling the Wait API by handleId must be the same as that of calling the Prepare API.

Example

REGISTER_TILING_DEFAULT(ReduceScatterCustomTilingData); // ReduceScatterCustomTilingData is a structure defined in the operator header file.
GET_TILING_DATA_WITH_STRUCT(ReduceScatterCustomTilingData, tilingData, tilingGM);
Hccl hccl;
GM_ADDR contextGM = AscendC::GetHcclContext<0>();  // In the kernel of the AscendC custom operator, obtain the HCCL context through this method.
__gm__ void *mc2InitTiling = (__gm__ void *)(&tiling->mc2InitTiling);
__gm__ void *mc2CcTiling = (__gm__ void *)(&(tiling->mc2CcTiling));
hccl.InitV2(contextGM, &tilingData);
auto ret = hccl.SetCcTilingV2(offsetof(ReduceScatterCustomTilingData, mc2CcTiling));
if (ret != HCCL_SUCCESS) {
    return;
}
if (AscendC::g_coreType == AIC) {
    HcclHandle handleId = hccl.ReduceScatter(sendBuf, recvBuf, 100, HcclDataType::HCCL_DATA_TYPE_INT8, HcclReduceOp::HCCL_REDUCE_SUM, 10);
    
    for (uint8_t i=0; i<10; i++) {
hccl.Commit(handleId ); //Notify the server that the ReduceScatter task can be executed.
    }
    for (uint8_t i=0; i<10; i++) {
        hccl.Wait(handleId); // Blocking API. Wait until the ReduceScatter task is complete.
    }
hccl.Finalize(); //If there are no other communication tasks to follow, notify the server to exit after executing the above ReduceScatter task.
}

Parent topic: HCCL Kernel APIs