用户进行异步拷贝需要用到Stream相关接口,MindIE Torch没有C++的Stream相关接口,而是直接复用libtorch的c10::Stream和c10::StreamGuard接口。MindIE Torch的Python相关Stream接口定义请参考API参考(Python)。
具体使用方法如下:
1 2 3 4 5 6 7 8 9 10 | // create stream c10::Stream stream1 = c10::Stream(c10::Stream::DEFAULT, at::Device("npu:0")); c10::Stream stream2 = c10::Stream(c10::Stream::UNSAFE, at::Device("npu:0"), 10); // set stream c10::StreamGuard streamGuard(stream1); // 构造StreamGuard,并设置当前stream为stream1 streamGuard.reset_stream(stream2); // 设置当前stream为stream2 // stream synchronize stream.synchronize(); |
1 2 3 4 5 6 7 8 9 10 11 12 | # create stream stream1 = mindietorch.npu.Stream() stream2 = mindietorch.npu.Stream("npu:0") stream3 = mindietorch.npu.Stream(device = 0, stream_id = 10) # set stream mindietorch.npu.set_stream(stream1) #方式1 with mindietorch.npu.stream(stream2): #方式2 self.assertTrue(mindietorch.npu.current_stream() == stream2) # assertTrue是unittest.TestCase中的断言 # stream synchronize stream.synchronize() |