功能介绍
ATB加速库图内多流并行功能需要注意两个部分:
- 存储GraphOp所用到的所有流到context中。
- 设置GraphOp中各个node对应Op的运行流。
存储GraphOp所用到的所有流
通过调用Context的接口SetExecuteStreams(const std::vector<aclrtStream> &streams)将需要用到的多个streams存放到context中。
设置GraphOp中各个node对应Op的运行流
通过调用Operation.h中的接口SetExecuteStreamId(Operation *operation, uint32_t streamId)设置对应Operation所运行的streamId,这里的streamId为上述设置Context时streams中的序号。(例如如果需要op1在stream2上运行(streams={stream1,stream2,stream3}),那么这里调用为SetExecuteStreamId(op1, 1))。

- 需要注意的是,对于加速库本身的Op不需要修改,但是对于继承Operation类的用户自定义Op,需要重新继承新的OperationInfra类,并且实现对应的SetExecuteStreamId、GetExecuteStreamId、GetExecuteStream方法。
- 该功能只支持设置ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1,开启全局中间tensor内存分配。
父主题: 图内多流并行