def test_basic(self):
def model(x, y):
return (x + y) * y
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
x = torch.randn(3, device="cuda", requires_grad=True)
y = torch.randn(3, device="cuda")
loss = model(x, y).sum()
loss.backward()
def test_htod(self):
def model(x, y):
a = x + y
return a * 3
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
x = torch.randn(3, device="cuda", requires_grad=True)
y = torch.randn((), device="cpu")
loss = model(x, y).sum()
loss.backward()
def test_factory(self):
def model(y):
x = torch.zeros(3, device="cuda:0")
x.add_(3)
return x * y
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
with self.subTest(i):
y = torch.randn(3, device="cuda:0", requires_grad=True)
loss = model(y).sum()
loss.backward()
def test_mutate_constant(self):
def model(x, y):
c = torch.tensor(1)
c.add_(2)
return x * y * 0 + c
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
with self.subTest(i):
x = torch.randn(1, device="cuda", requires_grad=True)
y = torch.randn(1, device="cuda")
loss = model(x, y).sum()
self.assertEqual(loss, torch.tensor(3.0, device="cuda"))
loss.backward()
def test_mutate_input(self):
def model(x, y):
y.add_(3)
return x * y
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
with self.subTest(i):
x = torch.randn(3, device="cuda", requires_grad=True)
y = torch.randn(3, device="cuda")
y_orig = y.clone()
loss = model(x, y).sum()
self.assertEqual(y, y_orig + 3)
loss.backward()
def test_dead_fill(self):
def model(x):
x = x.clone()
y = x[0:0]
x.fill_(2)
y.fill_(3)
return x, y
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
with self.subTest(i):
x = torch.empty(20, device="cuda:0")
rx, ry = model(x)
self.assertEqual(rx,
torch.full((20, ), 2.0, device="cuda:0"))
self.assertEqual(ry, torch.empty(0, device="cuda:0"))
def test_mutated_metadata(self):
# more tortured example at
# https://github.com/pytorch/pytorch/issues/81385
def model(x):
x = x.clone()
x.resize_(20)
x.fill_(2)
return x
with torchdynamo.optimize(aot_autograd_cudagraphs):
for i in range(N_ITERS):
with self.subTest(i):
x = torch.empty(0, device="cuda:0")
rx = model(x)
self.assertEqual(rx,
torch.full((20, ), 2.0, device="cuda:0"))
