AllGather

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

Task delivery API of AllGather, a collective communication operator. It returns handleId of the task to users. The function of AllGather is as follows: Re-sort the inputs of all nodes in a communicator based on the rank ID, combine the inputs, and send the result to the outputs of all nodes.

Prototype

      
           template <bool commit = false>
__aicore__ inline HcclHandle AllGather(GM_ADDR sendBuf, GM_ADDR recvBuf, uint64_t sendCount, HcclDataType dataType, uint64_t strideCount, uint8_t repeat = 1)

Parameters

**Table 1** Template parameters
Parameter	Input/Output	Description
commit	Input	Bool. Values: true: When the Prepare API is called, the Commit API instructs the server to execute the communication task. false: When the Prepare API is called, the server is not instructed to execute the communication task.

**Table 2** API parameters
Parameter	Input/Output	Description
sendBuf	Input	Address of the source data buffer.
recvBuf	Output	Address of the destination data buffer to receive collective communication result.
sendCount	Input	Number of sendBuf data elements involved in the AllGather operation. The data count of the recvBuf is computed as: sendCount × rank size, that is sendCount × number of cards.
dataType	Input	Data type of the AllGather operation. Currently, all data types of HcclDataType are supported. For details about HcclDataType, see Table 1.
strideCount	Input	If the value of strideCount is 0, the addresses of adjacent data blocks are contiguous when data on multiple cards is combined to the recvBuf of one card. The data blocks on card rank[i] are stored in the recvBuf, and the offset data count is i x sendCount. In non-multi-round tiling scenarios, you are advised to set this parameter to 0. If the value of strideCount is greater than 0, the offset data count of start addresses of adjacent data blocks in the recvBuf is strideCount when data on multiple cards is combined to the recvBuf of one card. The data blocks on card rank[i] are stored in the recvBuf, and the offset data count is i x strideCount. Note: The preceding offset data amount is the number of data elements. The unit is sizeof(dataType).
repeat	Input	Number of AllGather communication tasks delivered at a time. The value of repeat is greater than or equal to 1, and the default value is 1. When repeat is greater than 1, the server automatically computes the sendBuf and recvBuf addresses of each AllGather task. The calculation formulas are as follows: sendBuf[i] = sendBuf + sendCount x sizeof(datatype) x i, i∈[0, repeat) recvBuf[i] = recvBuf + sendCount x sizeof(datatype) x i, i∈[0, repeat) Note: When the value of repeat is greater than 1, the strideCount parameter must be used together to plan the communication data address.

Figure 1 AllGather communication example

Returns

The task ID handleId is returned. The value of handleId is greater than or equal to 0. If the API fails to be called, the value -1 is returned.

Restrictions

Before calling this API, ensure that the InitV2 and SetCcTilingV2 APIs have been called.
If the core for delivering the communication task is not specified in the config template parameters of the HCCL object, this API can be called only on the AIC or AIV core. If the core for delivering the communication task is specified in the config template parameters of the HCCL object, this API can be called on both the AIC and AIV cores. The API delivers the communication task only on the AIC or AIV core based on the specified core type.
For the Atlas A2 training products / Atlas A2 inference products , the total number of times that Prepare APIs are called in a communicator cannot exceed 63.
For the Atlas A3 training products / Atlas A3 inference products , the total number of times that all Prepare and InterHcclGroupSync APIs are called in a communicator cannot exceed 63.

Example

Non-multi-round tiling scenario

As shown in the following figure, each of the four cards has 300 float16 data records (sendCount = 300). Each card obtains data from the xGM memory. After data of each card is gathered, the result is output to the yGM of each card.

Figure 2 AllGather communication between four cards in non-multi-round tiling scenario

        
         
           
           
             extern "C" __global__ __aicore__ void all_gather_custom(GM_ADDR xGM, GM_ADDR yGM, GM_ADDR workspaceGM, GM_ADDR tilingGM) {
    auto sendBuf = xGM;  // xGM is the input GM address of AllGather.
    auto recvBuf = yGM;  // yGM is the output GM address of AllGather.
    uint64_t sendCount = 300;  // Each card has 300 float16 data records.
    uint64_t strideCount = 0;  // strideCount can be set to 0 in non-tiling scenarios.
    REGISTER_TILING_DEFAULT(AllGatherCustomTilingData); // AllGatherCustomTilingData is a structure defined in the operator header file.
    GET_TILING_DATA_WITH_STRUCT(AllGatherCustomTilingData, tilingData, tilingGM);

    Hccl hccl;
    GM_ADDR contextGM = AscendC::GetHcclContext<0>();  // Obtain the HCCL context on the kernel of the Ascend C custom operator.

    if (AscendC::g_coreType == AIV) {  // Specify AIV Cores for communication.
        hccl.InitV2(contextGM, &tilingData);
        auto ret = hccl.SetCcTilingV2(offsetof(AllGatherCustomTilingData, allGatherCcTiling));
        if (ret != HCCL_SUCCESS) {
          return;
        }
        HcclHandle handleId1 = hccl.AllGather<true>(sendBuf, recvBuf, sendCount, HcclDataType::HCCL_DATA_TYPE_FP16, strideCount);
        hccl.Wait(handleId1);    
        AscendC::SyncAll<true>();  // All AIV Cores are synchronized to prevent too fast execution on core 0. Calling the hccl.Finalize() API prematurely can cause suspension of other cores during the Wait operation.
        hccl.Finalize();
    }
}

            

          

        
       

Multi-round tiling scenario

Multi-round tiling is enabled to process communication equivalently to the foregoing non-multi-round tiling example. As shown in the following figure, the 300 float16 data records on each card are tiled into two first data blocks and one tail data block. Each first block includes 128 float16 data records (tileLen = 128), and each tail block includes 44 float16 data records (tailLen = 44). During implementation in the operator, the tiled data is divided into three rounds for AllGather communication, achieving communication result equivalent to that of non-multi-round tiling.

Figure 3 Data tiling of each card

Specifically, in the first round of communication, AllGather operation is performed on 0-0\1-0\2-0\3-0 data blocks on each rank. In the second round of communication, AllGather operation is performed on 0-1\1-1\2-1\3-1 data blocks on each rank. In the third round of communication, AllGather operation is performed on 0-2\1-2\2-2\3-2 data blocks on each rank. In the result of each round of communication, the number of data records between start addresses of adjacent data blocks on each card is strideCount. Take the result of the first round of communication as an example. The number of data records between start addresses of blocks 0-0 and 1-0 on rank0 is 300, computed by the formula strideCount = 2 x tileLen + 1 x tailLen.

Figure 4 First round of AllGather communication between four cards

         
          
            
            
              extern "C" __global__ __aicore__ void all_gather_custom(GM_ADDR xGM, GM_ADDR yGM, GM_ADDR workspaceGM, GM_ADDR tilingGM) {
    constexpr uint32_t tileNum = 2U;   // Number of first blocks
    constexpr uint64_t tileLen = 128U; // Number of data elements in first blocks
    constexpr uint32_t tailNum = 1U;   // Number of tail blocks
    constexpr uint64_t tailLen = 44U;  // Number of data elements in tail blocks
    auto sendBuf = xGM;  // xGM is the input GM address of AllGather.
    auto recvBuf = yGM;  // yGM is the output GM address of AllGather.
    REGISTER_TILING_DEFAULT(AllGatherCustomTilingData); // AllGatherCustomTilingData is a structure defined in the operator header file.
    GET_TILING_DATA_WITH_STRUCT(AllGatherCustomTilingData, tilingData, tilingGM);

    Hccl hccl;
    GM_ADDR contextGM = AscendC::GetHcclContext<0>();  // Obtain the HCCL context on the kernel of the Ascend C custom operator.
    if (AscendC::g_coreType == AIV) {  // Specify AIV Cores for communication.
        hccl.InitV2(contextGM, &tilingData);
        auto ret = hccl.SetCcTilingV2(offsetof(AllGatherCustomTilingData, allGatherCcTiling));
        if (ret != HCCL_SUCCESS) {
          return;
        }
        uint64_t strideCount = tileLen * tileNum + tailLen * tailNum;
        // Process two first blocks.
        constexpr uint32_t tileRepeat = tileNum; 
        // Except the input parameters sendBuf and recvBuf, other parameters for processing the two head blocks are the same. Therefore, repeat is set to 2, and sendBuf and recvBuf of the AllGather task of the second head block are updated in the API.
        HcclHandle handleId1 = hccl.AllGather<true>(sendBuf, recvBuf, tileLen, HcclDataType::HCCL_DATA_TYPE_FP16, strideCount, tileRepeat); 
        // Process one tail block.
        constexpr uint32_t kSizeOfFloat16 = 2U;
        sendBuf += tileLen * tileNum * kSizeOfFloat16;
        recvBuf += tileLen * tileNum * kSizeOfFloat16;
        constexpr uint32_t tailRepeat = tailNum; 
        HcclHandle handleId2 = hccl.AllGather<true>(sendBuf, recvBuf, tailLen, HcclDataType::HCCL_DATA_TYPE_FP16, strideCount, tailRepeat);
           
        for (uint8_t i=0; i<tileRepeat; i++) {
            hccl.Wait(handleId1);
        }
        hccl.Wait(handleId2);  
        AscendC::SyncAll<true>();  // All AIV Cores are synchronized to prevent too fast execution on core 0. Calling the hccl.Finalize() API prematurely can cause suspension of other cores during the Wait operation.
        hccl.Finalize();
    }
}

             

           

         
        

Parent topic: HCCL Kernel APIs