Function Description

Pay attention to the following two points when using the ATB's multi-stream parallelism in a graph:

Storing all streams used by GraphOp to the context
Setting the running stream of the operation corresponding to each node in GraphOp

Storing all streams used by GraphOp

Call the SetExecuteStreams(const std::vector<aclrtStream> &streams) API of Context to store the required streams to the context.

Setting the running stream of the operation corresponding to each node in GraphOp

Call the SetExecuteStreamId (operation × operation, uint32_t streamId) API in Operation.h to set the stream ID of the corresponding operation. The stream ID is the sequence number in streams when the context is set. (For example, if op1 needs to run on stream2 (streams = {stream1, stream2, stream3}), the SetExecuteStreamId(op1, 1) API is called.)

Note that the operator of the ATB does not need to be modified. However, the user-defined operator that inherits the Operation class needs to inherit the new OperationInfra class and implement the corresponding SetExecuteStreamId, GetExecuteStreamId, and GetExecuteStream methods.
This function supports only ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1, that is, global memory allocation of intermediate tensors is enabled.

Parent topic: Multi-Stream Parallelism in a Graph