torch代码是否异步执行情况分析

实验

test1 循环没有状态 z ，猜测循环会并行执行，不需要等待上一个循环的结果
test1 循环有状态 z ，猜测循环会顺序执行，需要等待上一个循环的结果

import torch
import time

n = 1000
s = 1024


def test1():
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(n):
        torch.matmul(torch.randn(s, s, device='cuda:0'), torch.randn(s, s, device='cuda:0'))
    torch.cuda.synchronize()
    end = time.time()
    print('Time: ', (end - start) / n * 1000, 'ms')

def test2():
    z = torch.randn(s, s).cuda()
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(n):
        torch.randn(s, s, device='cuda:0')
        z = torch.matmul(z, torch.randn(s, s, device='cuda:0'))
    torch.cuda.synchronize()
    end = time.time()
    print('Time: ', (end - start) / n * 1000, 'ms')

if __name__ == '__main__':
    torch.matmul(torch.randn(100, 100).cuda(), torch.randn(100, 100).cuda())
    test1()
    test2()

测试结果

Time:  1.9771325588226318 ms
Time:  1.9676969051361086 ms

两种情况没太大差异，说明判断错误

结论

torch按顺序执行代码，没有异步