Orchestrating Tasks

Procedure

Task orchestration is a process in which ranks involved in collective communication coordinate and perform synchronization and data movement in an orderly manner to complete a collective communication operation. The main purpose of task orchestration is to execute tasks on different communication threads in parallel to maximize resource utilization and improve overall performance.

For AIV communication operators, algorithm orchestration is implemented using data copy and synchronization APIs provided by the Ascend C programming language. You can also encapsulate high-level APIs for communication-specific data transfer and synchronization to facilitate repeated calls.

In the Ascend C programming framework, the external memory and internal memory accessible by the AI Core are referred to as global memory and local memory, respectively. The GlobalTensor data structure is used to store global data in the global memory, and the LocalTensor data structure is used to store data in the local memory of the AI Core. Whether for soft-sync settings, flag checks, or data copy, the underlying implementation relies on Ascend C's DataCopy-related APIs. These APIs support bidirectional data copy between GlobalTensor and LocalTensor, though the data flow direction varies across different scenarios:

• For soft-sync Record, the flag value is first set in the LocalTensor and then copied to the GlobalTensor via DataCopy.
• For soft-sync Wait, the flag value is set in the GlobalTensor. After the flag value is copied to the LocalTensor, a scalar comparison is performed. If the value matches the expected result, execution proceeds; otherwise, the copy and comparison operations are repeated.
• For data transfer where both the source and destination addresses reside in the GlobalTensor, the data must be routed through the internal memory space of the Vector Core. It is first copied from the source GlobalTensor to the LocalTensor via DataCopy, and then copied from the LocalTensor to the destination GlobalTensor.

For a Reduce operator, you can also leverage Ascend C's atomic operation APIs to enable the on-the-fly Reduce feature during data transfer from LocalTensor to GlobalTensor. For instance, the SetAtomicAdd API can be called when the Reduce type is SUM, while the SetAtomicMax API can be called when it is MAX.

Specifically, the task orchestration of an AIV operator involves the following steps:

1. Obtain the local communication memory, which is called HCCL buffer in HCCL.

2. Split the input data and calculate the offset.

   The HCCL Buffer has a default size of 200 MB. If the input data size exceeds this threshold, it must be split into multiple data chunks for separate processing. This requires multiple operator kernel launches on the host, with each task orchestration targeting only one of these data chunks.

3. Copy the operator input data to the HCCL buffer. Commonly used Ascend C APIs include DataCopy and DataCopyPad.
4. Perform pre-sync with other ranks to confirm that the communication is ready. Commonly used Ascend C APIs include DataCopy, PipeBarrier, SetFlag, and WaitFlag.
5. Perform data transfer by copying remote data into the local HCCL Buffer. Commonly used Ascend C APIs include DataCopy and DataCopyPad.
6. Perform post-sync with other ranks to confirm that the data transfer is complete.
7. Copy the result data in the HCCL buffer to the operator output memory.

Sample Code

The following uses the custom AllGather operator as an example to describe its task orchestration code snippet on the Vector Core.

Among them, CpGM2GM encapsulates an on-chip memory-to-memory copy API that supports un-aligned data transfer based on DataCopy and DataCopyPad. Meanwhile, Record1vN and WaitNv1 encapsulate inter-rank synchronization APIs based on DataCopy, SetFlag, and WaitFlag.

if (GetBlockIdx() != rank) {
    // Perform inter-rank synchronization. Each core waits for the data to be ready on other ranks.
    WaitNv1(tag);
    PipeBarrier<PIPE_ALL>();
    // Read the data of the peer rank to the output memory of the local rank.
    CpGM2GM(outputGM + GetBlockIdx() * count, cclGMOther, count);
    // Perform post-sync with the remote rank whose rank ID is GetBlockIdx().
    Record(GetBlockIdx(), tag);
    Wait(GetBlockIdx(), tag);
} else {
    // Copy the data from the input memory of the local rank to the transit memory for other ranks to read.
    CpGM2GM(cclGMSelf, inputGM, count);
    PipeBarrier<PIPE_ALL>();
    // Perform inter-rank synchronization to notify other ranks to read data.
    Record1vN(tag);
    // Copy the data from the transit memory of the current rank to the user output memory.
    CpGM2GM(outputGM + count * rank, cclGMSelf, count);
}