mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
d8c8ba2440
Pull Request resolved: https://github.com/pytorch/pytorch/pull/136964 Approved by: https://github.com/justinchuby, https://github.com/albanD
2329 lines
95 KiB
Python
2329 lines
95 KiB
Python
# Owner(s): ["module: codegen"]
|
|
# ruff: noqa: F841
|
|
|
|
import unittest
|
|
from contextlib import nullcontext
|
|
|
|
import torch
|
|
from torch._dispatch.python import (
|
|
enable_crossref_functionalize,
|
|
enable_python_dispatcher,
|
|
)
|
|
from torch._subclasses.functional_tensor import (
|
|
dispatch_functionalize,
|
|
FunctionalTensor,
|
|
FunctionalTensorMode,
|
|
)
|
|
from torch.fx.experimental.proxy_tensor import make_fx
|
|
from torch.fx.passes.reinplace import reinplace
|
|
from torch.multiprocessing.reductions import StorageWeakRef
|
|
from torch.testing._internal.common_utils import (
|
|
IS_WINDOWS,
|
|
run_tests,
|
|
skipIfTorchDynamo,
|
|
TEST_WITH_TORCHDYNAMO,
|
|
TestCase,
|
|
xfail_inherited_tests,
|
|
)
|
|
from torch.testing._internal.logging_tensor import capture_logs, LoggingTensor
|
|
from torch.utils import _pytree as pytree
|
|
from torch.utils._pytree import tree_map_only
|
|
|
|
|
|
def are_aliased(x, y):
|
|
x_storage = StorageWeakRef(x.storage())
|
|
y_storage = StorageWeakRef(y.storage())
|
|
return x_storage == y_storage
|
|
|
|
|
|
# We can unify testing and use functionalize() here instead
|
|
# if/when functorch moves into core.
|
|
# This is basically a crappy version of `functionalize()`.
|
|
def _functionalize(
|
|
f, *, reapply_views: bool, crossref: bool, skip_input_mutations: bool = False
|
|
):
|
|
def to_fun(t: torch.Tensor):
|
|
func_t = torch._to_functional_tensor(t)
|
|
func_t.requires_grad = t.requires_grad
|
|
return func_t
|
|
|
|
def wrapped(*inputs):
|
|
ctx = nullcontext()
|
|
if crossref:
|
|
ctx = enable_crossref_functionalize()
|
|
with ctx:
|
|
inputs_functional = tree_map_only(torch.Tensor, to_fun, inputs)
|
|
torch._enable_functionalization(reapply_views=reapply_views)
|
|
try:
|
|
out = f(*inputs_functional)
|
|
finally:
|
|
torch._disable_functionalization()
|
|
flat_inputs = pytree.tree_leaves(inputs)
|
|
flat_inputs_functional = pytree.tree_leaves(inputs_functional)
|
|
|
|
for inpt, input_functional in zip(flat_inputs, flat_inputs_functional):
|
|
torch._sync(input_functional)
|
|
inpt_new = torch._from_functional_tensor(input_functional)
|
|
if inpt_new is not inpt and not skip_input_mutations:
|
|
# Existing deficiency in functionalize():
|
|
# we don't correctly mutate input metadata (yet?)
|
|
if inpt_new.shape == inpt.shape:
|
|
inpt.copy_(inpt_new)
|
|
tree_map_only(torch.Tensor, torch._sync, out)
|
|
out_unwrapped = tree_map_only(
|
|
torch.Tensor, torch._from_functional_tensor, out
|
|
)
|
|
return out_unwrapped
|
|
|
|
return wrapped
|
|
|
|
|
|
@unittest.skipIf(
|
|
TEST_WITH_TORCHDYNAMO, "https://github.com/pytorch/pytorch/issues/81457"
|
|
)
|
|
class TestFunctionalization(TestCase):
|
|
crossref = False
|
|
|
|
def get_logs(self, func, *inpts, reapply_views=False, run_reinplace=False):
|
|
inpts_clone = tree_map_only(torch.Tensor, torch.clone, inpts)
|
|
traced_f = make_fx(
|
|
_functionalize(func, reapply_views=reapply_views, crossref=self.crossref)
|
|
)(*inpts)
|
|
if run_reinplace:
|
|
traced_f = reinplace(traced_f, *inpts_clone)
|
|
return traced_f.code
|
|
|
|
def assert_functionalization(
|
|
self, func, *inpts, reapply_views=False, mutated_input_metadata=False
|
|
):
|
|
clones1 = tree_map_only(torch.Tensor, torch.clone, inpts)
|
|
clones2 = tree_map_only(torch.Tensor, torch.clone, inpts)
|
|
clones3 = tree_map_only(torch.Tensor, torch.clone, inpts)
|
|
|
|
# Compare outputs (and mutated inputs), with and without functionalization.
|
|
out_ref = func(*inpts)
|
|
out_functional = _functionalize(
|
|
func, reapply_views=reapply_views, crossref=self.crossref
|
|
)(*clones1)
|
|
|
|
# The reinplacing pass is only valid to run with reapply_views=True.
|
|
functional_func = make_fx(
|
|
_functionalize(func, reapply_views=True, crossref=self.crossref)
|
|
)(*clones2)
|
|
reinplace_func = reinplace(functional_func, *clones2)
|
|
|
|
# NOTE: for now, need to pass in fresh inputs here, because make_fx
|
|
# will directly mutate the inputs that you trace with.
|
|
# Once this is fixed we can clean this up.
|
|
out_reinplace = reinplace_func(*clones3)
|
|
|
|
# functionalize() deficiency: input metadata mutations aren't propagated properly,
|
|
# so we just need to skip checks here for the tests that exercise that.
|
|
if not mutated_input_metadata:
|
|
flat_inpts = pytree.tree_leaves(inpts)
|
|
flat_clones1 = pytree.tree_leaves(clones1)
|
|
flat_clones3 = pytree.tree_leaves(clones3)
|
|
for inpt, input_clone, input_clone3 in zip(
|
|
flat_inpts, flat_clones1, flat_clones3
|
|
):
|
|
self.assertEqual(
|
|
inpt, input_clone
|
|
) # input mutations should still occur
|
|
self.assertEqual(inpt, input_clone3)
|
|
|
|
# Handle tests with multi-tensor outputs
|
|
if isinstance(out_ref, tuple):
|
|
out_refs, out_functionals, out_reinplaces = (
|
|
list(out_ref),
|
|
list(out_functional),
|
|
list(out_reinplace),
|
|
)
|
|
else:
|
|
out_refs, out_functionals, out_reinplaces = (
|
|
[out_ref],
|
|
[out_functional],
|
|
[out_reinplace],
|
|
)
|
|
|
|
for out_ref_, out_functional_, out_reinplace_ in zip(
|
|
out_refs, out_functionals, out_reinplaces
|
|
):
|
|
self.assertEqual(out_ref_, out_functional_)
|
|
self.assertEqual(out_ref_, out_reinplace_)
|
|
|
|
def test_save_for_backwards_segfault(self):
|
|
inp = torch._to_functional_tensor(
|
|
LoggingTensor(torch.randn(2, 2))
|
|
).requires_grad_(True)
|
|
inp.exp()
|
|
|
|
def test_multiple_views_of_same_base(self):
|
|
def f(x):
|
|
y = x.view(-1)
|
|
z = x.view(-1)
|
|
x.add_(1)
|
|
# y should have been updated.
|
|
y2 = y + 1
|
|
# z should have been updated too.
|
|
z2 = z + 1
|
|
return z2
|
|
|
|
self.assert_functionalization(f, torch.ones(4))
|
|
|
|
def test_freeze(self):
|
|
def f(x):
|
|
y = x.clone()
|
|
z = y[0]
|
|
torch._freeze_functional_tensor(y)
|
|
x.add_(1)
|
|
self.assertRaises(RuntimeError, lambda: y.add_(1))
|
|
self.assertRaises(RuntimeError, lambda: z.add_(1))
|
|
return z
|
|
|
|
_functionalize(f, reapply_views=True, crossref=self.crossref)(torch.ones(3, 3))
|
|
|
|
def test_copy_stride_mismatch(self):
|
|
def f(x):
|
|
y = torch.empty_strided((2, 2), (5, 1))
|
|
y.copy_(x)
|
|
return y
|
|
|
|
r = _functionalize(f, reapply_views=True, crossref=self.crossref)(
|
|
torch.ones(2, 2)
|
|
)
|
|
self.assertEqual(r.stride(), (5, 1))
|
|
|
|
def test_set_(self):
|
|
def f(x):
|
|
y = torch.ones(2)
|
|
y.set_(x.storage())
|
|
return y
|
|
|
|
# We should probaby get the crossref test to work,
|
|
# but fixing it for Storage() objects is annoying.
|
|
r = _functionalize(f, reapply_views=True, crossref=False)(torch.ones(2))
|
|
self.assertEqual(str(r.device), "cpu")
|
|
|
|
def test_advanced_indexing(self):
|
|
def f():
|
|
x = torch.zeros(3, 3)
|
|
idx = torch.tensor([0])
|
|
val = torch.ones(3, 1)
|
|
x[:, idx] = val
|
|
return x
|
|
|
|
self.assert_functionalization(f)
|
|
|
|
def test_view_clone_view_inplace(self):
|
|
def f(input):
|
|
shape = [1, 1024, 128, 128]
|
|
input_reshaped = input.view(shape)
|
|
out = input_reshaped.clone()
|
|
r = out.view(input.shape)
|
|
r.relu_()
|
|
return r
|
|
|
|
def g(x):
|
|
loss = f(x).sum()
|
|
import torch.fx.traceback as fx_traceback
|
|
from torch._functorch.aot_autograd import (
|
|
setup_stacktrace_preservation_hooks,
|
|
)
|
|
|
|
setup_stacktrace_preservation_hooks([loss.grad_fn])
|
|
with fx_traceback.preserve_node_meta():
|
|
loss.backward()
|
|
return x.grad
|
|
|
|
with torch.autograd.detect_anomaly(check_nan=False):
|
|
logs = self.get_logs(g, torch.ones(16, 64, 128, 128, requires_grad=True))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [1, 1024, 128, 128]); arg0_1 = None
|
|
clone = torch.ops.aten.clone.default(view_copy); view_copy = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(clone, [16, 64, 128, 128])
|
|
relu = torch.ops.aten.relu.default(view_copy_1); view_copy_1 = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(relu, [1, 1024, 128, 128]); relu = None
|
|
view_copy_3 = torch.ops.aten.view_copy.default(view_copy_2, [16, 64, 128, 128]); view_copy_2 = None
|
|
view_copy_4 = torch.ops.aten.view_copy.default(clone, [16, 64, 128, 128]); clone = view_copy_4 = None
|
|
sum_1 = torch.ops.aten.sum.default(view_copy_3)
|
|
ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format); sum_1 = None
|
|
expand_copy = torch.ops.aten.expand_copy.default(ones_like, [16, 64, 128, 128]); ones_like = None
|
|
view_copy_5 = torch.ops.aten.view_copy.default(expand_copy, [1, 1024, 128, 128]); expand_copy = None
|
|
new_empty_strided = torch.ops.aten.new_empty_strided.default(view_copy_5, [1, 1024, 128, 128], [16777216, 16384, 128, 1])
|
|
copy = torch.ops.aten.copy.default(new_empty_strided, view_copy_5); new_empty_strided = view_copy_5 = None
|
|
view_copy_6 = torch.ops.aten.view_copy.default(copy, [16, 64, 128, 128]); view_copy_6 = None
|
|
view_copy_7 = torch.ops.aten.view_copy.default(copy, [16, 64, 128, 128])
|
|
clone_1 = torch.ops.aten.clone.default(view_copy_7, memory_format = torch.contiguous_format)
|
|
threshold_backward = torch.ops.aten.threshold_backward.default(clone_1, view_copy_3, 0); clone_1 = view_copy_3 = None
|
|
copy_1 = torch.ops.aten.copy.default(view_copy_7, threshold_backward); view_copy_7 = threshold_backward = None
|
|
view_copy_8 = torch.ops.aten.view_copy.default(copy_1, [1, 1024, 128, 128]); copy_1 = None
|
|
view_copy_9 = torch.ops.aten.view_copy.default(view_copy_8, [16, 64, 128, 128]); view_copy_9 = None
|
|
view_copy_10 = torch.ops.aten.view_copy.default(copy, [16, 64, 128, 128]); copy = None
|
|
detach_copy = torch.ops.aten.detach_copy.default(view_copy_10); view_copy_10 = detach_copy = None
|
|
view_copy_11 = torch.ops.aten.view_copy.default(view_copy_8, [16, 64, 128, 128]); view_copy_8 = None
|
|
detach_copy_1 = torch.ops.aten.detach_copy.default(view_copy_11); view_copy_11 = None
|
|
return detach_copy_1
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_simple(self):
|
|
def f(x):
|
|
# simple test: 1 view op, 1 inplace op
|
|
tmp = torch.ones(4, 2)
|
|
y = x.view(4, 2)
|
|
y.add_(tmp)
|
|
z = x * x
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [4, 2])
|
|
add = torch.ops.aten.add.Tensor(view_copy, ones); view_copy = ones = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(add, [4, 2]); add = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(view_copy_1, [4, 2])
|
|
mul = torch.ops.aten.mul.Tensor(view_copy_1, view_copy_1); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, view_copy_1); arg0_1 = view_copy_1 = copy_ = None
|
|
return view_copy_2
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view = torch.ops.aten.view.default(arg0_1, [4, 2])
|
|
add = torch.ops.aten.add.Tensor(view, ones); view = ones = None
|
|
view_1 = torch.ops.aten.view.default(add, [4, 2]); add = None
|
|
view_2 = torch.ops.aten.view.default(view_1, [4, 2])
|
|
mul = torch.ops.aten.mul.Tensor(view_1, view_1); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, view_1); arg0_1 = view_1 = copy_ = None
|
|
return view_2
|
|
""",
|
|
)
|
|
|
|
def test_simple_out(self):
|
|
def f(x):
|
|
tmp = torch.ones(4, 2)
|
|
y = x.view(4, 2)
|
|
# the out= tensor will get resized, since it has size=0 to start.
|
|
z = torch.empty(())
|
|
torch.add(y, tmp, out=z)
|
|
w = z * z
|
|
return w
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [4, 2]); arg0_1 = None
|
|
empty = torch.ops.aten.empty.memory_format([], device = device(type='cpu'), pin_memory = False); empty = None
|
|
add = torch.ops.aten.add.Tensor(view_copy, ones); view_copy = ones = None
|
|
mul = torch.ops.aten.mul.Tensor(add, add); add = None
|
|
return mul
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view = torch.ops.aten.view.default(arg0_1, [4, 2]); arg0_1 = None
|
|
empty = torch.ops.aten.empty.memory_format([], device = device(type='cpu'), pin_memory = False); empty = None
|
|
add = torch.ops.aten.add.Tensor(view, ones); view = ones = None
|
|
mul = torch.ops.aten.mul.Tensor(add, add); add = None
|
|
return mul
|
|
""",
|
|
)
|
|
|
|
def test_multi_out(self):
|
|
def f(x):
|
|
# aminmax.out returns a tuple of tensors.
|
|
# functionalization should properly handle the tuple.
|
|
out_min = torch.empty(4)
|
|
out_max = torch.empty(4)
|
|
torch.aminmax(x, dim=0, out=(out_max, out_min))
|
|
return out_max
|
|
|
|
self.assert_functionalization(f, torch.arange(8, dtype=torch.float32))
|
|
logs = self.get_logs(f, torch.arange(8, dtype=torch.float32))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
empty = torch.ops.aten.empty.memory_format([4], device = device(type='cpu'), pin_memory = False); empty = None
|
|
empty_1 = torch.ops.aten.empty.memory_format([4], device = device(type='cpu'), pin_memory = False); empty_1 = None
|
|
aminmax = torch.ops.aten.aminmax.default(arg0_1, dim = 0); arg0_1 = None
|
|
getitem = aminmax[0]
|
|
getitem_1 = aminmax[1]; aminmax = getitem_1 = None
|
|
return getitem
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f,
|
|
torch.arange(8, dtype=torch.float32),
|
|
reapply_views=True,
|
|
run_reinplace=True,
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
empty = torch.ops.aten.empty.memory_format([4], device = device(type='cpu'), pin_memory = False); empty = None
|
|
empty_1 = torch.ops.aten.empty.memory_format([4], device = device(type='cpu'), pin_memory = False); empty_1 = None
|
|
aminmax = torch.ops.aten.aminmax.default(arg0_1, dim = 0); arg0_1 = None
|
|
getitem = aminmax[0]
|
|
getitem_1 = aminmax[1]; aminmax = getitem_1 = None
|
|
return getitem
|
|
""",
|
|
)
|
|
|
|
def test_tensor_ctr(self):
|
|
def f(x):
|
|
y = torch.tensor((1, 2, 3))
|
|
z = y.view(-1)
|
|
z.add_(1)
|
|
return y
|
|
|
|
inpt = torch.arange(3, dtype=torch.float32)
|
|
self.assert_functionalization(f, inpt)
|
|
|
|
logs = self.get_logs(f, inpt)
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
_tensor_constant0 = self._tensor_constant0
|
|
lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0); _tensor_constant0 = None
|
|
view_copy = torch.ops.aten.view_copy.default(lift_fresh_copy, [-1]); lift_fresh_copy = None
|
|
add = torch.ops.aten.add.Tensor(view_copy, 1); view_copy = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(add, [3]); add = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(view_copy_1, [-1]); view_copy_2 = None
|
|
return view_copy_1
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(f, inpt, reapply_views=True, run_reinplace=True)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
_tensor_constant0 = self._tensor_constant0
|
|
lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0); _tensor_constant0 = None
|
|
view = torch.ops.aten.view.default(lift_fresh_copy, [-1]); lift_fresh_copy = None
|
|
add = torch.ops.aten.add_.Tensor(view, 1); add = None
|
|
view_1 = torch.ops.aten.view.default(view, [3]); view = None
|
|
view_2 = torch.ops.aten.view.default(view_1, [-1]); view_2 = None
|
|
return view_1
|
|
""",
|
|
)
|
|
|
|
def test_advanced_indexing_correct_strides(self):
|
|
def f(a):
|
|
# This test requires that *_scatter ops are able to return
|
|
# non-contiguous tensors.
|
|
b = a.clone()[:, 1]
|
|
c = torch.ones_like(b, dtype=torch.bool)
|
|
d = b.masked_fill_(c, 0)
|
|
return d
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2), reapply_views=True)
|
|
|
|
def test_tensor_list_mixed_functional_nonfunctional(self):
|
|
nonfunctional_tensor = torch.ones(2, dtype=torch.long)
|
|
|
|
def f(x):
|
|
# simple test: 1 view op, 1 inplace op
|
|
functional_tensor = torch.ones(2, dtype=torch.long)
|
|
out = x[functional_tensor, nonfunctional_tensor]
|
|
return out
|
|
|
|
out = f(torch.ones(2, 2))
|
|
out_functional = _functionalize(f, reapply_views=True, crossref=self.crossref)(
|
|
torch.ones(2, 2)
|
|
)
|
|
self.assertEqual(out, out_functional)
|
|
|
|
def test_inplace_on_non_view(self):
|
|
def f(x):
|
|
# test for the case where we functionalize an inplace op on the other tensor - not a view.
|
|
# This is worth checking because the tensor will have an empty ViewMeta stack, which needs to be special cased.
|
|
tmp = torch.ones(4, 2)
|
|
y = x.view(4, 2)
|
|
x.add_(tmp)
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [4, 2]); view_copy = None
|
|
add = torch.ops.aten.add.Tensor(arg0_1, ones); ones = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, add); arg0_1 = copy_ = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(add, [4, 2]); add = None
|
|
return view_copy_1
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view = torch.ops.aten.view.default(arg0_1, [4, 2]); view = None
|
|
add = torch.ops.aten.add.Tensor(arg0_1, ones); ones = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, add); arg0_1 = copy_ = None
|
|
view_1 = torch.ops.aten.view.default(add, [4, 2]); add = None
|
|
return view_1
|
|
""",
|
|
)
|
|
|
|
# Some ops that are mutable are neither inplace nor out= ops.
|
|
# They also need special handling.
|
|
def test_mutable_op_not_inplace_or_other(self):
|
|
def f(x):
|
|
return torch._fused_moving_avg_obs_fq_helper(
|
|
x, x, x, x, x, x, x, 1.0, 0, 1, 0
|
|
)
|
|
|
|
logs = self.get_logs(f, torch.ones(1))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
_fused_moving_avg_obs_fq_helper_functional = torch.ops.aten._fused_moving_avg_obs_fq_helper_functional.default(arg0_1, arg0_1, arg0_1, arg0_1, arg0_1, arg0_1, arg0_1, 1.0, 0, 1, 0)
|
|
getitem = _fused_moving_avg_obs_fq_helper_functional[0]
|
|
getitem_1 = _fused_moving_avg_obs_fq_helper_functional[1]
|
|
getitem_2 = _fused_moving_avg_obs_fq_helper_functional[2]; getitem_2 = None
|
|
getitem_3 = _fused_moving_avg_obs_fq_helper_functional[3]; getitem_3 = None
|
|
getitem_4 = _fused_moving_avg_obs_fq_helper_functional[4]; getitem_4 = None
|
|
getitem_5 = _fused_moving_avg_obs_fq_helper_functional[5]; _fused_moving_avg_obs_fq_helper_functional = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, getitem_5); arg0_1 = getitem_5 = copy_ = None
|
|
return (getitem, getitem_1)
|
|
""", # noqa: B950
|
|
)
|
|
|
|
def test_as_strided(self):
|
|
def f(x):
|
|
y = x.as_strided((2,), (2,), 1)
|
|
y.add_(1)
|
|
return x
|
|
|
|
self.assert_functionalization(f, torch.ones(9))
|
|
logs = self.get_logs(f, torch.ones(9))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
as_strided_copy = torch.ops.aten.as_strided_copy.default(arg0_1, [2], [2], 1)
|
|
add = torch.ops.aten.add.Tensor(as_strided_copy, 1); as_strided_copy = None
|
|
as_strided_scatter = torch.ops.aten.as_strided_scatter.default(arg0_1, add, [2], [2], 1); add = None
|
|
as_strided_copy_1 = torch.ops.aten.as_strided_copy.default(as_strided_scatter, [2], [2], 1); as_strided_copy_1 = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, as_strided_scatter); arg0_1 = copy_ = None
|
|
return as_strided_scatter
|
|
""",
|
|
)
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
as_strided = torch.ops.aten.as_strided.default(arg0_1, [2], [2], 1)
|
|
add = torch.ops.aten.add.Tensor(as_strided, 1); as_strided = None
|
|
as_strided_scatter = torch.ops.aten.as_strided_scatter.default(arg0_1, add, [2], [2], 1); add = None
|
|
as_strided_1 = torch.ops.aten.as_strided.default(as_strided_scatter, [2], [2], 1); as_strided_1 = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, as_strided_scatter); arg0_1 = copy_ = None
|
|
return as_strided_scatter
|
|
""",
|
|
)
|
|
|
|
def test_tensor_list_composite(self):
|
|
def f(x):
|
|
# Test an op with TensorList input
|
|
y = torch.block_diag(x, x)
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2))
|
|
logs = self.get_logs(f, torch.ones(2, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
block_diag = torch.ops.aten.block_diag.default([arg0_1, arg0_1]); arg0_1 = None
|
|
return block_diag
|
|
""",
|
|
)
|
|
|
|
def test_cat(self):
|
|
def f(x):
|
|
out = torch.empty(0)
|
|
torch.cat((x,), out=out)
|
|
return out
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2))
|
|
logs = self.get_logs(f, torch.ones(2, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
empty = torch.ops.aten.empty.memory_format([0], device = device(type='cpu'), pin_memory = False); empty = None
|
|
cat = torch.ops.aten.cat.default([arg0_1]); arg0_1 = None
|
|
return cat
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
empty = torch.ops.aten.empty.memory_format([0], device = device(type='cpu'), pin_memory = False); empty = None
|
|
cat = torch.ops.aten.cat.default([arg0_1]); arg0_1 = None
|
|
return cat
|
|
""",
|
|
)
|
|
|
|
def test_diagonal(self):
|
|
def f(x):
|
|
# test: view ops that take a subset of the original tensor (select/diagonal)
|
|
tmp = torch.ones(2)
|
|
y = x.clone().diagonal()
|
|
y.add_(tmp)
|
|
z = x * x
|
|
return z
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2))
|
|
logs = self.get_logs(f, torch.ones(2, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
clone = torch.ops.aten.clone.default(arg0_1)
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(clone)
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy, ones); diagonal_copy = ones = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(clone, add); clone = add = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter); diagonal_scatter = diagonal_copy_1 = None
|
|
mul = torch.ops.aten.mul.Tensor(arg0_1, arg0_1); arg0_1 = None
|
|
return mul
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
clone = torch.ops.aten.clone.default(arg0_1)
|
|
diagonal = torch.ops.aten.diagonal.default(clone)
|
|
add = torch.ops.aten.add_.Tensor(diagonal, ones); diagonal = ones = add = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(clone); clone = diagonal_1 = None
|
|
mul = torch.ops.aten.mul.Tensor(arg0_1, arg0_1); arg0_1 = None
|
|
return mul
|
|
""",
|
|
)
|
|
|
|
def test_diagonal_mutated_input(self):
|
|
def f(x):
|
|
# simple test: there are pending updates afterwards, which the test syncs manually
|
|
tmp = torch.ones(2)
|
|
y = x.diagonal()
|
|
y.add_(tmp)
|
|
return x
|
|
|
|
x = torch.ones(2, 2)
|
|
self.assert_functionalization(f, x)
|
|
logs = self.get_logs(f, torch.ones(2, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(arg0_1)
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy, ones); diagonal_copy = ones = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(arg0_1, add); add = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter); diagonal_copy_1 = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, diagonal_scatter); arg0_1 = copy_ = None
|
|
return diagonal_scatter
|
|
""",
|
|
)
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal = torch.ops.aten.diagonal.default(arg0_1)
|
|
add = torch.ops.aten.add.Tensor(diagonal, ones); diagonal = ones = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(arg0_1, add); add = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(diagonal_scatter); diagonal_1 = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, diagonal_scatter); arg0_1 = copy_ = None
|
|
return diagonal_scatter
|
|
""",
|
|
)
|
|
|
|
def test_channels_last_contiguous(self):
|
|
def f(x):
|
|
return x.contiguous(memory_format=torch.channels_last)
|
|
tmp = torch.ones(2)
|
|
y = x.diagonal()
|
|
y.add_(tmp)
|
|
return x
|
|
|
|
x = torch.randn(4, 8, 8, 3).permute(0, 3, 1, 2)
|
|
self.assert_functionalization(f, x)
|
|
logs = self.get_logs(f, x).strip()
|
|
# There should be no clone in the graph
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
def forward(self, arg0_1):
|
|
return arg0_1""",
|
|
)
|
|
|
|
def test_split(self):
|
|
def f(x):
|
|
# test: view ops that return multiple tensors (split)
|
|
tmp = torch.ones(2)
|
|
y1, y2 = x.split(2)
|
|
y3 = y2.diagonal()
|
|
y3.add_(tmp)
|
|
z = x * x
|
|
return y3
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
split_copy = torch.ops.aten.split_copy.Tensor(arg0_1, 2)
|
|
getitem = split_copy[0]; getitem = None
|
|
getitem_1 = split_copy[1]; split_copy = None
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(getitem_1); getitem_1 = None
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy, ones); diagonal_copy = ones = None
|
|
split_copy_1 = torch.ops.aten.split_copy.Tensor(arg0_1, 2)
|
|
getitem_2 = split_copy_1[0]; getitem_2 = None
|
|
getitem_3 = split_copy_1[1]; split_copy_1 = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(getitem_3, add); getitem_3 = add = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(arg0_1, diagonal_scatter, 0, 2, 4); diagonal_scatter = None
|
|
split_copy_2 = torch.ops.aten.split_copy.Tensor(slice_scatter, 2)
|
|
getitem_4 = split_copy_2[0]; getitem_4 = None
|
|
getitem_5 = split_copy_2[1]; split_copy_2 = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(getitem_5); getitem_5 = None
|
|
mul = torch.ops.aten.mul.Tensor(slice_scatter, slice_scatter); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, slice_scatter); arg0_1 = slice_scatter = copy_ = None
|
|
return diagonal_copy_1
|
|
""",
|
|
) # noqa: B950
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
split = torch.ops.aten.split.Tensor(arg0_1, 2)
|
|
getitem = split[0]; getitem = None
|
|
getitem_1 = split[1]; split = None
|
|
diagonal = torch.ops.aten.diagonal.default(getitem_1); getitem_1 = None
|
|
add = torch.ops.aten.add.Tensor(diagonal, ones); diagonal = ones = None
|
|
split_1 = torch.ops.aten.split.Tensor(arg0_1, 2)
|
|
getitem_2 = split_1[0]; getitem_2 = None
|
|
getitem_3 = split_1[1]; split_1 = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(getitem_3, add); getitem_3 = add = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(arg0_1, diagonal_scatter, 0, 2, 4); diagonal_scatter = None
|
|
split_2 = torch.ops.aten.split.Tensor(slice_scatter, 2)
|
|
getitem_4 = split_2[0]; getitem_4 = None
|
|
getitem_5 = split_2[1]; split_2 = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(getitem_5); getitem_5 = None
|
|
mul = torch.ops.aten.mul.Tensor(slice_scatter, slice_scatter); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, slice_scatter); arg0_1 = slice_scatter = copy_ = None
|
|
return diagonal_1
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_split_with_sizes(self):
|
|
def f(x):
|
|
# test: view ops that return multiple tensors (split_with_sizes)
|
|
tmp = torch.ones(2)
|
|
y1, y2 = x.split_with_sizes([2, 2])
|
|
y3 = y1.diagonal()
|
|
y3.add_(tmp)
|
|
z = x * x
|
|
return y3
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
split_with_sizes_copy = torch.ops.aten.split_with_sizes_copy.default(arg0_1, [2, 2])
|
|
getitem = split_with_sizes_copy[0]
|
|
getitem_1 = split_with_sizes_copy[1]; split_with_sizes_copy = getitem_1 = None
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(getitem); getitem = None
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy, ones); diagonal_copy = ones = None
|
|
split_with_sizes_copy_1 = torch.ops.aten.split_with_sizes_copy.default(arg0_1, [2, 2])
|
|
getitem_2 = split_with_sizes_copy_1[0]
|
|
getitem_3 = split_with_sizes_copy_1[1]; split_with_sizes_copy_1 = getitem_3 = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(getitem_2, add); getitem_2 = add = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(arg0_1, diagonal_scatter, 0, 0, 2); diagonal_scatter = None
|
|
split_with_sizes_copy_2 = torch.ops.aten.split_with_sizes_copy.default(slice_scatter, [2, 2])
|
|
getitem_4 = split_with_sizes_copy_2[0]
|
|
getitem_5 = split_with_sizes_copy_2[1]; split_with_sizes_copy_2 = getitem_5 = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(getitem_4); getitem_4 = None
|
|
mul = torch.ops.aten.mul.Tensor(slice_scatter, slice_scatter); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, slice_scatter); arg0_1 = slice_scatter = copy_ = None
|
|
return diagonal_copy_1
|
|
""",
|
|
) # noqa: B950
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2], device = device(type='cpu'), pin_memory = False)
|
|
split_with_sizes = torch.ops.aten.split_with_sizes.default(arg0_1, [2, 2])
|
|
getitem = split_with_sizes[0]
|
|
getitem_1 = split_with_sizes[1]; split_with_sizes = getitem_1 = None
|
|
diagonal = torch.ops.aten.diagonal.default(getitem); getitem = None
|
|
add = torch.ops.aten.add.Tensor(diagonal, ones); diagonal = ones = None
|
|
split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(arg0_1, [2, 2])
|
|
getitem_2 = split_with_sizes_1[0]
|
|
getitem_3 = split_with_sizes_1[1]; split_with_sizes_1 = getitem_3 = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(getitem_2, add); getitem_2 = add = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(arg0_1, diagonal_scatter, 0, 0, 2); diagonal_scatter = None
|
|
split_with_sizes_2 = torch.ops.aten.split_with_sizes.default(slice_scatter, [2, 2])
|
|
getitem_4 = split_with_sizes_2[0]
|
|
getitem_5 = split_with_sizes_2[1]; split_with_sizes_2 = getitem_5 = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(getitem_4); getitem_4 = None
|
|
mul = torch.ops.aten.mul.Tensor(slice_scatter, slice_scatter); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, slice_scatter); arg0_1 = slice_scatter = copy_ = None
|
|
return diagonal_1
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_slice(self):
|
|
def f(x):
|
|
tmp = torch.ones(4)
|
|
x.transpose_(1, 0)
|
|
y = x[0:2]
|
|
y.add_(tmp)
|
|
return x
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2), mutated_input_metadata=True)
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4], device = device(type='cpu'), pin_memory = False)
|
|
transpose_copy = torch.ops.aten.transpose_copy.int(arg0_1, 1, 0)
|
|
slice_copy = torch.ops.aten.slice_copy.Tensor(transpose_copy, 0, 0, 2); transpose_copy = None
|
|
add = torch.ops.aten.add.Tensor(slice_copy, ones); slice_copy = ones = None
|
|
transpose_copy_1 = torch.ops.aten.transpose_copy.int(arg0_1, 1, 0); arg0_1 = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(transpose_copy_1, add, 0, 0, 2); transpose_copy_1 = add = None
|
|
transpose_copy_2 = torch.ops.aten.transpose_copy.int(slice_scatter, 1, 0); slice_scatter = None
|
|
transpose_copy_3 = torch.ops.aten.transpose_copy.int(transpose_copy_2, 1, 0)
|
|
slice_copy_1 = torch.ops.aten.slice_copy.Tensor(transpose_copy_3, 0, 0, 2); transpose_copy_3 = slice_copy_1 = None
|
|
transpose_copy_4 = torch.ops.aten.transpose_copy.int(transpose_copy_2, 1, 0); transpose_copy_2 = None
|
|
return transpose_copy_4
|
|
""",
|
|
) # noqa: B950
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4], device = device(type='cpu'), pin_memory = False)
|
|
transpose = torch.ops.aten.transpose.int(arg0_1, 1, 0)
|
|
slice_1 = torch.ops.aten.slice.Tensor(transpose, 0, 0, 2); transpose = None
|
|
add = torch.ops.aten.add.Tensor(slice_1, ones); slice_1 = ones = None
|
|
transpose_1 = torch.ops.aten.transpose.int(arg0_1, 1, 0); arg0_1 = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(transpose_1, add, 0, 0, 2); transpose_1 = add = None
|
|
transpose_2 = torch.ops.aten.transpose.int(slice_scatter, 1, 0); slice_scatter = None
|
|
transpose_3 = torch.ops.aten.transpose.int(transpose_2, 1, 0)
|
|
slice_2 = torch.ops.aten.slice.Tensor(transpose_3, 0, 0, 2); transpose_3 = slice_2 = None
|
|
transpose_4 = torch.ops.aten.transpose.int(transpose_2, 1, 0); transpose_2 = None
|
|
return transpose_4
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_view_inplace(self):
|
|
def f(x):
|
|
# test: view + inplace op (transpose_)
|
|
tmp = torch.ones(4)
|
|
x.transpose_(1, 0)
|
|
y = x[0]
|
|
y.add_(tmp)
|
|
return x
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2), mutated_input_metadata=True)
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4], device = device(type='cpu'), pin_memory = False)
|
|
transpose_copy = torch.ops.aten.transpose_copy.int(arg0_1, 1, 0)
|
|
select_copy = torch.ops.aten.select_copy.int(transpose_copy, 0, 0); transpose_copy = None
|
|
add = torch.ops.aten.add.Tensor(select_copy, ones); select_copy = ones = None
|
|
transpose_copy_1 = torch.ops.aten.transpose_copy.int(arg0_1, 1, 0); arg0_1 = None
|
|
select_scatter = torch.ops.aten.select_scatter.default(transpose_copy_1, add, 0, 0); transpose_copy_1 = add = None
|
|
transpose_copy_2 = torch.ops.aten.transpose_copy.int(select_scatter, 1, 0); select_scatter = None
|
|
transpose_copy_3 = torch.ops.aten.transpose_copy.int(transpose_copy_2, 1, 0)
|
|
select_copy_1 = torch.ops.aten.select_copy.int(transpose_copy_3, 0, 0); transpose_copy_3 = select_copy_1 = None
|
|
transpose_copy_4 = torch.ops.aten.transpose_copy.int(transpose_copy_2, 1, 0); transpose_copy_2 = None
|
|
return transpose_copy_4
|
|
""",
|
|
) # noqa: B950
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4], device = device(type='cpu'), pin_memory = False)
|
|
transpose = torch.ops.aten.transpose.int(arg0_1, 1, 0)
|
|
select = torch.ops.aten.select.int(transpose, 0, 0); transpose = None
|
|
add = torch.ops.aten.add.Tensor(select, ones); select = ones = None
|
|
transpose_1 = torch.ops.aten.transpose.int(arg0_1, 1, 0); arg0_1 = None
|
|
select_scatter = torch.ops.aten.select_scatter.default(transpose_1, add, 0, 0); transpose_1 = add = None
|
|
transpose_2 = torch.ops.aten.transpose.int(select_scatter, 1, 0); select_scatter = None
|
|
transpose_3 = torch.ops.aten.transpose.int(transpose_2, 1, 0)
|
|
select_1 = torch.ops.aten.select.int(transpose_3, 0, 0); transpose_3 = select_1 = None
|
|
transpose_4 = torch.ops.aten.transpose.int(transpose_2, 1, 0); transpose_2 = None
|
|
return transpose_4
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_unbind(self):
|
|
def f(x):
|
|
# test: view + inplace op (transpose_)
|
|
tmp = torch.ones(4)
|
|
x.transpose_(1, 0)
|
|
y, _ = x.unbind(0)
|
|
y.add_(tmp)
|
|
return x
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2), mutated_input_metadata=True)
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4], device = device(type='cpu'), pin_memory = False)
|
|
transpose_copy = torch.ops.aten.transpose_copy.int(arg0_1, 1, 0)
|
|
unbind_copy = torch.ops.aten.unbind_copy.int(transpose_copy); transpose_copy = None
|
|
getitem = unbind_copy[0]
|
|
getitem_1 = unbind_copy[1]; unbind_copy = getitem_1 = None
|
|
add = torch.ops.aten.add.Tensor(getitem, ones); getitem = ones = None
|
|
transpose_copy_1 = torch.ops.aten.transpose_copy.int(arg0_1, 1, 0); arg0_1 = None
|
|
select_scatter = torch.ops.aten.select_scatter.default(transpose_copy_1, add, 0, 0); transpose_copy_1 = add = None
|
|
transpose_copy_2 = torch.ops.aten.transpose_copy.int(select_scatter, 1, 0); select_scatter = None
|
|
transpose_copy_3 = torch.ops.aten.transpose_copy.int(transpose_copy_2, 1, 0)
|
|
unbind_copy_1 = torch.ops.aten.unbind_copy.int(transpose_copy_3); transpose_copy_3 = None
|
|
getitem_2 = unbind_copy_1[0]; getitem_2 = None
|
|
getitem_3 = unbind_copy_1[1]; unbind_copy_1 = getitem_3 = None
|
|
transpose_copy_4 = torch.ops.aten.transpose_copy.int(transpose_copy_2, 1, 0); transpose_copy_2 = None
|
|
return transpose_copy_4
|
|
""",
|
|
) # noqa: B950
|
|
|
|
# NB: even with reapply_views=True, we expect to see scatter op
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=False
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4], device = device(type='cpu'), pin_memory = False)
|
|
transpose = torch.ops.aten.transpose.int(arg0_1, 1, 0)
|
|
unbind = torch.ops.aten.unbind.int(transpose); transpose = None
|
|
getitem = unbind[0]
|
|
getitem_1 = unbind[1]; unbind = getitem_1 = None
|
|
add = torch.ops.aten.add.Tensor(getitem, ones); getitem = ones = None
|
|
transpose_1 = torch.ops.aten.transpose.int(arg0_1, 1, 0); arg0_1 = None
|
|
select_scatter = torch.ops.aten.select_scatter.default(transpose_1, add, 0, 0); transpose_1 = add = None
|
|
transpose_2 = torch.ops.aten.transpose.int(select_scatter, 1, 0); select_scatter = None
|
|
transpose_3 = torch.ops.aten.transpose.int(transpose_2, 1, 0)
|
|
unbind_1 = torch.ops.aten.unbind.int(transpose_3); transpose_3 = None
|
|
getitem_2 = unbind_1[0]; getitem_2 = None
|
|
getitem_3 = unbind_1[1]; unbind_1 = getitem_3 = None
|
|
transpose_4 = torch.ops.aten.transpose.int(transpose_2, 1, 0); transpose_2 = None
|
|
return transpose_4
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_optional_tensor_list(self):
|
|
def f(x):
|
|
# test: an operator that takes in a List[Optional[Tensor]] argument
|
|
# (index_put)
|
|
y = x.view(8)
|
|
indices = torch.arange(4)
|
|
values = torch.arange(4, dtype=y.dtype)
|
|
y.index_put_((indices,), values, accumulate=False)
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [8])
|
|
arange = torch.ops.aten.arange.default(4, device = device(type='cpu'), pin_memory = False)
|
|
arange_1 = torch.ops.aten.arange.default(4, dtype = torch.float32, device = device(type='cpu'), pin_memory = False)
|
|
index_put = torch.ops.aten.index_put.default(view_copy, [arange], arange_1); view_copy = arange = arange_1 = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(index_put, [4, 2]); index_put = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(view_copy_1, [8])
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, view_copy_1); arg0_1 = view_copy_1 = copy_ = None
|
|
return view_copy_2
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_scalars(self):
|
|
def f(x):
|
|
# test: the pass can handle scalar inputs properly
|
|
tmp = torch.ones(4, 2)
|
|
y = x.view(4, 2)
|
|
y.add_(1)
|
|
z = 2 * y
|
|
z.div_(1)
|
|
return z
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False); ones = None
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [4, 2])
|
|
add = torch.ops.aten.add.Tensor(view_copy, 1); view_copy = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(add, [4, 2]); add = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(view_copy_1, [4, 2])
|
|
mul = torch.ops.aten.mul.Tensor(view_copy_2, 2); view_copy_2 = None
|
|
div = torch.ops.aten.div.Tensor(mul, 1); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, view_copy_1); arg0_1 = view_copy_1 = copy_ = None
|
|
return div
|
|
""",
|
|
)
|
|
|
|
@skipIfTorchDynamo("Test does not work with TorchDynamo")
|
|
def test_metadata_change(self):
|
|
def f(x):
|
|
# ops like ge_() are allowed to change the dtype of the input.
|
|
# functionalization should pick up on that.
|
|
y = x.clone()
|
|
out = y.ge_(0)
|
|
return out
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
clone = torch.ops.aten.clone.default(arg0_1); arg0_1 = None
|
|
ge = torch.ops.aten.ge.Scalar(clone, 0); clone = None
|
|
_to_copy = torch.ops.aten._to_copy.default(ge, dtype = torch.float32, layout = torch.strided); ge = None
|
|
return _to_copy
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
clone = torch.ops.aten.clone.default(arg0_1); arg0_1 = None
|
|
ge = torch.ops.aten.ge.Scalar(clone, 0); clone = None
|
|
_to_copy = torch.ops.aten._to_copy.default(ge, dtype = torch.float32, layout = torch.strided); ge = None
|
|
return _to_copy
|
|
""",
|
|
) # noqa: B950
|
|
|
|
@skipIfTorchDynamo("Test does not work with TorchDynamo")
|
|
def test_metadata_change_out_op(self):
|
|
def f(t, y):
|
|
out_1 = torch.ones(1)
|
|
return torch.add(t, y, out=out_1)
|
|
|
|
inpt1, inpt2 = torch.tensor([1]), torch.tensor([1])
|
|
inpt1_func, inpt2_func = (
|
|
torch._to_functional_tensor(inpt1),
|
|
torch._to_functional_tensor(inpt2),
|
|
)
|
|
|
|
out_ref = f(inpt1, inpt2)
|
|
torch._enable_functionalization(reapply_views=True)
|
|
try:
|
|
out_functional = f(inpt1_func, inpt2_func)
|
|
finally:
|
|
torch._disable_functionalization()
|
|
self.assertEqual(out_ref, torch._from_functional_tensor(out_functional))
|
|
|
|
def test_only_one_view(self):
|
|
def f(x):
|
|
# This tests that we don't have any unnecessary views in the trace.
|
|
# If the input wasn't mutated, we don't need to regenerate it,
|
|
# so there should be a total of 1 op in the output trace.
|
|
return x.view(4, 2)
|
|
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [4, 2]); arg0_1 = None
|
|
return view_copy
|
|
""",
|
|
)
|
|
|
|
def test_everything(self):
|
|
def f(x):
|
|
# test: everything
|
|
tmp = torch.ones(2, 2)
|
|
x2 = x + x
|
|
y = x2.view(8)
|
|
z0 = y.reshape(2, 4)
|
|
z1 = z0.transpose(1, 0)
|
|
z1.unsqueeze_(0)
|
|
z1.squeeze_()
|
|
z2, z3 = z1.split(2)
|
|
z2.add_(tmp)
|
|
z4 = z0[0] + z2.reshape(4)
|
|
return z2
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2))
|
|
logs = self.get_logs(f, torch.ones(4, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
add = torch.ops.aten.add.Tensor(arg0_1, arg0_1); arg0_1 = None
|
|
view_copy = torch.ops.aten.view_copy.default(add, [8])
|
|
view_copy_1 = torch.ops.aten.view_copy.default(view_copy, [2, 4]); view_copy = None
|
|
transpose_copy = torch.ops.aten.transpose_copy.int(view_copy_1, 1, 0)
|
|
unsqueeze_copy = torch.ops.aten.unsqueeze_copy.default(transpose_copy, 0); transpose_copy = None
|
|
squeeze_copy = torch.ops.aten.squeeze_copy.default(unsqueeze_copy); unsqueeze_copy = None
|
|
split_copy = torch.ops.aten.split_copy.Tensor(squeeze_copy, 2); squeeze_copy = None
|
|
getitem = split_copy[0]
|
|
getitem_1 = split_copy[1]; split_copy = getitem_1 = None
|
|
add_1 = torch.ops.aten.add.Tensor(getitem, ones); getitem = ones = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(add, [8]); add = None
|
|
view_copy_3 = torch.ops.aten.view_copy.default(view_copy_2, [2, 4]); view_copy_2 = None
|
|
transpose_copy_1 = torch.ops.aten.transpose_copy.int(view_copy_3, 1, 0); view_copy_3 = None
|
|
unsqueeze_copy_1 = torch.ops.aten.unsqueeze_copy.default(transpose_copy_1, 0); transpose_copy_1 = None
|
|
squeeze_copy_1 = torch.ops.aten.squeeze_copy.default(unsqueeze_copy_1); unsqueeze_copy_1 = None
|
|
slice_scatter = torch.ops.aten.slice_scatter.default(squeeze_copy_1, add_1, 0, 0, 2); squeeze_copy_1 = add_1 = None
|
|
unsqueeze_copy_2 = torch.ops.aten.unsqueeze_copy.default(slice_scatter, 0); slice_scatter = None
|
|
squeeze_copy_2 = torch.ops.aten.squeeze_copy.dim(unsqueeze_copy_2, 0); unsqueeze_copy_2 = None
|
|
transpose_copy_2 = torch.ops.aten.transpose_copy.int(squeeze_copy_2, 1, 0); squeeze_copy_2 = None
|
|
view_copy_4 = torch.ops.aten.view_copy.default(transpose_copy_2, [8]); transpose_copy_2 = None
|
|
view_copy_5 = torch.ops.aten.view_copy.default(view_copy_4, [4, 2]); view_copy_4 = None
|
|
view_copy_6 = torch.ops.aten.view_copy.default(view_copy_5, [8])
|
|
view_copy_7 = torch.ops.aten.view_copy.default(view_copy_6, [2, 4]); view_copy_6 = None
|
|
transpose_copy_3 = torch.ops.aten.transpose_copy.int(view_copy_7, 1, 0); view_copy_7 = None
|
|
unsqueeze_copy_3 = torch.ops.aten.unsqueeze_copy.default(transpose_copy_3, 0); transpose_copy_3 = None
|
|
squeeze_copy_3 = torch.ops.aten.squeeze_copy.default(unsqueeze_copy_3); unsqueeze_copy_3 = None
|
|
split_copy_1 = torch.ops.aten.split_copy.Tensor(squeeze_copy_3, 2); squeeze_copy_3 = None
|
|
getitem_2 = split_copy_1[0]
|
|
getitem_3 = split_copy_1[1]; split_copy_1 = getitem_3 = None
|
|
select_copy = torch.ops.aten.select_copy.int(view_copy_1, 0, 0); view_copy_1 = select_copy = None
|
|
view_copy_8 = torch.ops.aten.view_copy.default(getitem_2, [4]); view_copy_8 = None
|
|
view_copy_9 = torch.ops.aten.view_copy.default(view_copy_5, [8])
|
|
view_copy_10 = torch.ops.aten.view_copy.default(view_copy_9, [2, 4]); view_copy_9 = None
|
|
select_copy_1 = torch.ops.aten.select_copy.int(view_copy_10, 0, 0); view_copy_10 = None
|
|
view_copy_11 = torch.ops.aten.view_copy.default(view_copy_5, [8]); view_copy_5 = None
|
|
view_copy_12 = torch.ops.aten.view_copy.default(view_copy_11, [2, 4]); view_copy_11 = None
|
|
transpose_copy_4 = torch.ops.aten.transpose_copy.int(view_copy_12, 1, 0); view_copy_12 = None
|
|
unsqueeze_copy_4 = torch.ops.aten.unsqueeze_copy.default(transpose_copy_4, 0); transpose_copy_4 = None
|
|
squeeze_copy_4 = torch.ops.aten.squeeze_copy.default(unsqueeze_copy_4); unsqueeze_copy_4 = None
|
|
split_copy_2 = torch.ops.aten.split_copy.Tensor(squeeze_copy_4, 2); squeeze_copy_4 = None
|
|
getitem_4 = split_copy_2[0]
|
|
getitem_5 = split_copy_2[1]; split_copy_2 = getitem_5 = None
|
|
view_copy_13 = torch.ops.aten.view_copy.default(getitem_4, [4]); getitem_4 = None
|
|
add_2 = torch.ops.aten.add.Tensor(select_copy_1, view_copy_13); select_copy_1 = view_copy_13 = add_2 = None
|
|
return getitem_2
|
|
""",
|
|
) # noqa: B950
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(4, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
add = torch.ops.aten.add.Tensor(arg0_1, arg0_1); arg0_1 = None
|
|
view = torch.ops.aten.view.default(add, [8])
|
|
view_1 = torch.ops.aten.view.default(view, [2, 4]); view = None
|
|
transpose = torch.ops.aten.transpose.int(view_1, 1, 0)
|
|
unsqueeze = torch.ops.aten.unsqueeze.default(transpose, 0); transpose = None
|
|
squeeze = torch.ops.aten.squeeze.default(unsqueeze); unsqueeze = None
|
|
split = torch.ops.aten.split.Tensor(squeeze, 2); squeeze = None
|
|
getitem = split[0]
|
|
getitem_1 = split[1]; split = getitem_1 = None
|
|
add_1 = torch.ops.aten.add_.Tensor(getitem, ones); getitem = ones = add_1 = None
|
|
view_2 = torch.ops.aten.view.default(add, [8]); add = None
|
|
view_3 = torch.ops.aten.view.default(view_2, [2, 4]); view_2 = None
|
|
transpose_1 = torch.ops.aten.transpose.int(view_3, 1, 0); view_3 = None
|
|
unsqueeze_1 = torch.ops.aten.unsqueeze.default(transpose_1, 0); transpose_1 = None
|
|
squeeze_1 = torch.ops.aten.squeeze.default(unsqueeze_1); unsqueeze_1 = None
|
|
unsqueeze_2 = torch.ops.aten.unsqueeze.default(squeeze_1, 0); squeeze_1 = None
|
|
squeeze_2 = torch.ops.aten.squeeze.dim(unsqueeze_2, 0); unsqueeze_2 = None
|
|
transpose_2 = torch.ops.aten.transpose.int(squeeze_2, 1, 0); squeeze_2 = None
|
|
view_4 = torch.ops.aten.view.default(transpose_2, [8]); transpose_2 = None
|
|
view_5 = torch.ops.aten.view.default(view_4, [4, 2]); view_4 = None
|
|
view_6 = torch.ops.aten.view.default(view_5, [8])
|
|
view_7 = torch.ops.aten.view.default(view_6, [2, 4]); view_6 = None
|
|
transpose_3 = torch.ops.aten.transpose.int(view_7, 1, 0); view_7 = None
|
|
unsqueeze_3 = torch.ops.aten.unsqueeze.default(transpose_3, 0); transpose_3 = None
|
|
squeeze_3 = torch.ops.aten.squeeze.default(unsqueeze_3); unsqueeze_3 = None
|
|
split_1 = torch.ops.aten.split.Tensor(squeeze_3, 2); squeeze_3 = None
|
|
getitem_2 = split_1[0]
|
|
getitem_3 = split_1[1]; split_1 = getitem_3 = None
|
|
select = torch.ops.aten.select.int(view_1, 0, 0); view_1 = select = None
|
|
clone = torch.ops.aten.clone.default(getitem_2, memory_format = torch.contiguous_format)
|
|
_unsafe_view = torch.ops.aten._unsafe_view.default(clone, [4]); clone = None
|
|
view_8 = torch.ops.aten.view.default(view_5, [8]); view_5 = None
|
|
view_9 = torch.ops.aten.view.default(view_8, [2, 4]); view_8 = None
|
|
select_1 = torch.ops.aten.select.int(view_9, 0, 0); view_9 = None
|
|
add_2 = torch.ops.aten.add.Tensor(select_1, _unsafe_view); select_1 = _unsafe_view = add_2 = None
|
|
return getitem_2
|
|
""",
|
|
)
|
|
|
|
def test_reapply_views_simple(self):
|
|
def f(x):
|
|
tmp = torch.ones(4, 2)
|
|
y = x.view(4, 2)
|
|
y.add_(tmp)
|
|
z = x * x
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(4, 2), reapply_views=True)
|
|
logs = self.get_logs(f, torch.ones(4, 2), reapply_views=True)
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
ones = torch.ops.aten.ones.default([4, 2], device = device(type='cpu'), pin_memory = False)
|
|
view = torch.ops.aten.view.default(arg0_1, [4, 2])
|
|
add = torch.ops.aten.add.Tensor(view, ones); view = ones = None
|
|
view_1 = torch.ops.aten.view.default(add, [4, 2]); add = None
|
|
view_2 = torch.ops.aten.view.default(view_1, [4, 2])
|
|
mul = torch.ops.aten.mul.Tensor(view_1, view_1); mul = None
|
|
copy_ = torch.ops.aten.copy_.default(arg0_1, view_1); arg0_1 = view_1 = copy_ = None
|
|
return view_2
|
|
""",
|
|
)
|
|
|
|
def test_aliases_maintained_after_pass_when_reapplying_views(self):
|
|
def f(x):
|
|
tmp = torch.ones(4, 2)
|
|
y = x.view(4, 2)
|
|
z = x.view(4, 2)
|
|
y.add_(tmp)
|
|
return y, z
|
|
|
|
input_functional = torch._to_functional_tensor(torch.ones(4, 2))
|
|
torch._enable_functionalization(reapply_views=True)
|
|
try:
|
|
y, z = f(input_functional)
|
|
torch._sync(y)
|
|
torch._sync(z)
|
|
finally:
|
|
torch._disable_functionalization()
|
|
|
|
# y and z are aliases inside of the function, and that aliasing relationship should be maintained.
|
|
_y = torch._from_functional_tensor(y)
|
|
_z = torch._from_functional_tensor(z)
|
|
self.assertTrue(are_aliased(_y, _z))
|
|
|
|
# copy_() gets its own test, because it used to be special cased in functionalization.
|
|
# However, now it works pretty similar to other functional ops
|
|
def test_copy_(self):
|
|
def f(x):
|
|
tmp = torch.zeros(2, 2)
|
|
tmp_slice = tmp.diagonal()
|
|
y = tmp_slice.copy_(x)
|
|
z = y.add_(x)
|
|
return z
|
|
|
|
# Test 1: copy_() with same dtype and shape
|
|
# to() is a composite op that noops when the dtype/shape match, so nothing gets logged.
|
|
# self.assert_functionalization(f, torch.ones(2))
|
|
logs = self.get_logs(f, torch.ones(2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(zeros)
|
|
copy = torch.ops.aten.copy.default(diagonal_copy, arg0_1); diagonal_copy = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(zeros, copy); zeros = copy = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter)
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy_1, arg0_1); diagonal_copy_1 = arg0_1 = None
|
|
diagonal_scatter_1 = torch.ops.aten.diagonal_scatter.default(diagonal_scatter, add); diagonal_scatter = add = None
|
|
diagonal_copy_2 = torch.ops.aten.diagonal_copy.default(diagonal_scatter_1); diagonal_scatter_1 = None
|
|
return diagonal_copy_2
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal = torch.ops.aten.diagonal.default(zeros)
|
|
copy = torch.ops.aten.copy_.default(diagonal, arg0_1); diagonal = copy = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(zeros)
|
|
add = torch.ops.aten.add_.Tensor(diagonal_1, arg0_1); diagonal_1 = arg0_1 = add = None
|
|
diagonal_2 = torch.ops.aten.diagonal.default(zeros); zeros = None
|
|
return diagonal_2
|
|
""",
|
|
)
|
|
|
|
# Test 2: copy_() with same dtype, different shape
|
|
self.assert_functionalization(f, torch.ones(1))
|
|
logs = self.get_logs(f, torch.ones(1))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(zeros)
|
|
copy = torch.ops.aten.copy.default(diagonal_copy, arg0_1); diagonal_copy = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(zeros, copy); zeros = copy = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter)
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy_1, arg0_1); diagonal_copy_1 = arg0_1 = None
|
|
diagonal_scatter_1 = torch.ops.aten.diagonal_scatter.default(diagonal_scatter, add); diagonal_scatter = add = None
|
|
diagonal_copy_2 = torch.ops.aten.diagonal_copy.default(diagonal_scatter_1); diagonal_scatter_1 = None
|
|
return diagonal_copy_2
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(1), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal = torch.ops.aten.diagonal.default(zeros)
|
|
copy = torch.ops.aten.copy_.default(diagonal, arg0_1); diagonal = copy = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(zeros)
|
|
add = torch.ops.aten.add_.Tensor(diagonal_1, arg0_1); diagonal_1 = arg0_1 = add = None
|
|
diagonal_2 = torch.ops.aten.diagonal.default(zeros); zeros = None
|
|
return diagonal_2
|
|
""",
|
|
)
|
|
|
|
# Test 3: copy_() with different dtype, same shape
|
|
self.assert_functionalization(f, torch.ones(2, dtype=torch.long))
|
|
logs = self.get_logs(f, torch.ones(2, dtype=torch.long))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(zeros)
|
|
copy = torch.ops.aten.copy.default(diagonal_copy, arg0_1); diagonal_copy = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(zeros, copy); zeros = copy = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter)
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy_1, arg0_1); diagonal_copy_1 = arg0_1 = None
|
|
diagonal_scatter_1 = torch.ops.aten.diagonal_scatter.default(diagonal_scatter, add); diagonal_scatter = add = None
|
|
diagonal_copy_2 = torch.ops.aten.diagonal_copy.default(diagonal_scatter_1); diagonal_scatter_1 = None
|
|
return diagonal_copy_2
|
|
""",
|
|
) # noqa: B950
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, dtype=torch.long), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal = torch.ops.aten.diagonal.default(zeros)
|
|
copy = torch.ops.aten.copy_.default(diagonal, arg0_1); diagonal = copy = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(zeros)
|
|
add = torch.ops.aten.add_.Tensor(diagonal_1, arg0_1); diagonal_1 = arg0_1 = add = None
|
|
diagonal_2 = torch.ops.aten.diagonal.default(zeros); zeros = None
|
|
return diagonal_2
|
|
""",
|
|
) # noqa: B950
|
|
|
|
# Test 4: copy_() with different dtype, different shape
|
|
self.assert_functionalization(f, torch.ones(1, dtype=torch.long))
|
|
logs = self.get_logs(f, torch.ones(1, dtype=torch.long))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(zeros)
|
|
copy = torch.ops.aten.copy.default(diagonal_copy, arg0_1); diagonal_copy = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(zeros, copy); zeros = copy = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter)
|
|
add = torch.ops.aten.add.Tensor(diagonal_copy_1, arg0_1); diagonal_copy_1 = arg0_1 = None
|
|
diagonal_scatter_1 = torch.ops.aten.diagonal_scatter.default(diagonal_scatter, add); diagonal_scatter = add = None
|
|
diagonal_copy_2 = torch.ops.aten.diagonal_copy.default(diagonal_scatter_1); diagonal_scatter_1 = None
|
|
return diagonal_copy_2
|
|
""",
|
|
) # noqa: B950
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(1, dtype=torch.long), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([2, 2], device = device(type='cpu'), pin_memory = False)
|
|
diagonal = torch.ops.aten.diagonal.default(zeros)
|
|
copy = torch.ops.aten.copy_.default(diagonal, arg0_1); diagonal = copy = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(zeros)
|
|
add = torch.ops.aten.add_.Tensor(diagonal_1, arg0_1); diagonal_1 = arg0_1 = add = None
|
|
diagonal_2 = torch.ops.aten.diagonal.default(zeros); zeros = None
|
|
return diagonal_2
|
|
""",
|
|
) # noqa: B950
|
|
|
|
def test_expand_symint(self):
|
|
# Once some existing SymInt bugs are ironed out, we should update
|
|
# this test to plumb FakeSymbolicTensors through it
|
|
def f(x):
|
|
return x.expand(x.size(0), x.size(1))
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2))
|
|
logs = self.get_logs(f, torch.ones(2, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
expand_copy = torch.ops.aten.expand_copy.default(arg0_1, [2, 2]); arg0_1 = None
|
|
return expand_copy
|
|
""",
|
|
)
|
|
|
|
def test_fill_(self):
|
|
def f(x):
|
|
y = x + x
|
|
z = y.diagonal()
|
|
z.fill_(0)
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2))
|
|
logs = self.get_logs(f, torch.ones(2, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
add = torch.ops.aten.add.Tensor(arg0_1, arg0_1); arg0_1 = None
|
|
diagonal_copy = torch.ops.aten.diagonal_copy.default(add)
|
|
fill = torch.ops.aten.fill.Scalar(diagonal_copy, 0); diagonal_copy = None
|
|
diagonal_scatter = torch.ops.aten.diagonal_scatter.default(add, fill); add = fill = None
|
|
diagonal_copy_1 = torch.ops.aten.diagonal_copy.default(diagonal_scatter); diagonal_copy_1 = None
|
|
return diagonal_scatter
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
add = torch.ops.aten.add.Tensor(arg0_1, arg0_1); arg0_1 = None
|
|
diagonal = torch.ops.aten.diagonal.default(add)
|
|
fill = torch.ops.aten.fill_.Scalar(diagonal, 0); diagonal = fill = None
|
|
diagonal_1 = torch.ops.aten.diagonal.default(add); diagonal_1 = None
|
|
return add
|
|
""",
|
|
)
|
|
|
|
def test_resize_smaller(self):
|
|
def f(w):
|
|
# Resizing to a smaller size doesn't affect storage
|
|
x = w + 1
|
|
y = x.view(4, 4)
|
|
y.resize_(3, 3)
|
|
y2 = y.view(-1)
|
|
y2.add_(1)
|
|
z = y + 1
|
|
return z
|
|
|
|
self.assert_functionalization(f, torch.ones(8, 2))
|
|
logs = self.get_logs(f, torch.ones(8, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
add = torch.ops.aten.add.Tensor(arg0_1, 1); arg0_1 = None
|
|
view_copy = torch.ops.aten.view_copy.default(add, [4, 4])
|
|
resize = torch.ops.aten.resize.default(view_copy, [3, 3]); resize = None
|
|
as_strided_copy = torch.ops.aten.as_strided_copy.default(view_copy, [3, 3], [3, 1]); view_copy = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(as_strided_copy, [-1]); as_strided_copy = None
|
|
add_1 = torch.ops.aten.add.Tensor(view_copy_1, 1); view_copy_1 = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(add, [4, 4]); add = None
|
|
as_strided_copy_1 = torch.ops.aten.as_strided_copy.default(view_copy_2, [3, 3], [3, 1]); as_strided_copy_1 = None
|
|
view_copy_3 = torch.ops.aten.view_copy.default(add_1, [3, 3]); add_1 = None
|
|
as_strided_scatter = torch.ops.aten.as_strided_scatter.default(view_copy_2, view_copy_3, [3, 3], [3, 1]); view_copy_2 = view_copy_3 = None
|
|
view_copy_4 = torch.ops.aten.view_copy.default(as_strided_scatter, [8, 2]); as_strided_scatter = None
|
|
view_copy_5 = torch.ops.aten.view_copy.default(view_copy_4, [4, 4])
|
|
as_strided_copy_2 = torch.ops.aten.as_strided_copy.default(view_copy_5, [3, 3], [3, 1]); view_copy_5 = None
|
|
view_copy_6 = torch.ops.aten.view_copy.default(as_strided_copy_2, [-1]); as_strided_copy_2 = view_copy_6 = None
|
|
view_copy_7 = torch.ops.aten.view_copy.default(view_copy_4, [4, 4]); view_copy_4 = None
|
|
as_strided_copy_3 = torch.ops.aten.as_strided_copy.default(view_copy_7, [3, 3], [3, 1]); view_copy_7 = None
|
|
add_2 = torch.ops.aten.add.Tensor(as_strided_copy_3, 1); as_strided_copy_3 = None
|
|
return add_2
|
|
""", # noqa: B950
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(8, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
add = torch.ops.aten.add.Tensor(arg0_1, 1); arg0_1 = None
|
|
view = torch.ops.aten.view.default(add, [4, 4])
|
|
resize = torch.ops.aten.resize.default(view, [3, 3]); resize = None
|
|
as_strided = torch.ops.aten.as_strided.default(view, [3, 3], [3, 1]); view = None
|
|
view_1 = torch.ops.aten.view.default(as_strided, [-1]); as_strided = None
|
|
add_1 = torch.ops.aten.add_.Tensor(view_1, 1); add_1 = None
|
|
view_2 = torch.ops.aten.view.default(add, [4, 4]); add = None
|
|
as_strided_1 = torch.ops.aten.as_strided.default(view_2, [3, 3], [3, 1]); as_strided_1 = None
|
|
view_3 = torch.ops.aten.view.default(view_1, [3, 3]); view_1 = view_3 = None
|
|
view_4 = torch.ops.aten.view.default(view_2, [8, 2]); view_2 = None
|
|
view_5 = torch.ops.aten.view.default(view_4, [4, 4])
|
|
as_strided_2 = torch.ops.aten.as_strided.default(view_5, [3, 3], [3, 1]); view_5 = None
|
|
view_6 = torch.ops.aten.view.default(as_strided_2, [-1]); as_strided_2 = view_6 = None
|
|
view_7 = torch.ops.aten.view.default(view_4, [4, 4]); view_4 = None
|
|
as_strided_3 = torch.ops.aten.as_strided.default(view_7, [3, 3], [3, 1]); view_7 = None
|
|
add_2 = torch.ops.aten.add_.Tensor(as_strided_3, 1); add_2 = None
|
|
return as_strided_3
|
|
""",
|
|
)
|
|
|
|
def test_resize_same_size_diff_rank(self):
|
|
def f(x):
|
|
y = x.clone()
|
|
y.resize_(25, 5)
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(5, 5, 5))
|
|
|
|
def test_resize_larger_valid(self):
|
|
def f(x):
|
|
y = x + 1
|
|
# resizing a tensor to a larger size is only currently allowed
|
|
# if the tensor-to-resize is not a view / has no outstanding views.
|
|
# See Note [resize_() in functionalization pass]
|
|
y.resize_(5, 5)
|
|
y2 = y.view(25)
|
|
# Do a mutation to ensure that aliases of the output of resize_()
|
|
# propagate mutations correctly.
|
|
# I'm using fill_ specifically because I want to guarantee that
|
|
# none of the output has uninitialized memory at the end
|
|
# (since these tests compare the data output against a reference impl)
|
|
y2.fill_(1)
|
|
out = y + 1
|
|
return y, out
|
|
|
|
self.assert_functionalization(f, torch.ones(8, 2))
|
|
logs = self.get_logs(f, torch.ones(8, 2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
add = torch.ops.aten.add.Tensor(arg0_1, 1); arg0_1 = None
|
|
resize = torch.ops.aten.resize.default(add, [5, 5]); add = None
|
|
view_copy = torch.ops.aten.view_copy.default(resize, [25]); resize = None
|
|
fill = torch.ops.aten.fill.Scalar(view_copy, 1); view_copy = None
|
|
view_copy_1 = torch.ops.aten.view_copy.default(fill, [5, 5]); fill = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(view_copy_1, [25]); view_copy_2 = None
|
|
add_1 = torch.ops.aten.add.Tensor(view_copy_1, 1)
|
|
return (view_copy_1, add_1)
|
|
""",
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(8, 2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
add = torch.ops.aten.add.Tensor(arg0_1, 1); arg0_1 = None
|
|
resize = torch.ops.aten.resize_.default(add, [5, 5]); resize = None
|
|
view = torch.ops.aten.view.default(add, [25]); add = None
|
|
fill = torch.ops.aten.fill_.Scalar(view, 1); fill = None
|
|
view_1 = torch.ops.aten.view.default(view, [5, 5]); view = None
|
|
view_2 = torch.ops.aten.view.default(view_1, [25]); view_2 = None
|
|
add_1 = torch.ops.aten.add.Tensor(view_1, 1)
|
|
return (view_1, add_1)
|
|
""",
|
|
)
|
|
|
|
def test_resize_larger_invalid(self):
|
|
def f(x):
|
|
y = x + 1
|
|
z = y.view(4, 4)
|
|
# resizing a tensor to a larger size is only currently allowed
|
|
# if the tensor-to-resize is not a view / has no outstanding views.
|
|
# See Note [resize_() in functionalization pass]
|
|
# This should fail
|
|
z.resize_(5, 5)
|
|
z2 = z.view(25)
|
|
z2.fill_(1)
|
|
out = z + 1
|
|
return y, out
|
|
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"Attempted to resize a view tensor to a larger size. This is not allowed in the functionalization pass",
|
|
):
|
|
self.assert_functionalization(f, torch.ones(8, 2))
|
|
|
|
def test_nested_functions_propagate_updates(self):
|
|
def g(x):
|
|
# Create a view of x
|
|
y = x[0]
|
|
y.add_(1)
|
|
# The view, y, gets deallocated at the end of this function
|
|
|
|
def f(x):
|
|
# Calling g(x) should mutate x
|
|
g(x)
|
|
# We expect x to be synced here, even though the alias created in g() has been deallocated!
|
|
y = x + x
|
|
return y
|
|
|
|
self.assert_functionalization(f, torch.ones(2, 2))
|
|
|
|
def test_mixed_wrappers_valid(self):
|
|
def f(x, y):
|
|
z = x + y
|
|
z.add_(1)
|
|
return z
|
|
|
|
x1_not_functional = LoggingTensor(torch.ones(4))
|
|
x2_functional = torch._to_functional_tensor(LoggingTensor(torch.ones(4)))
|
|
|
|
with capture_logs() as logs:
|
|
y = f(x1_not_functional, x2_functional)
|
|
|
|
# Make sure that functionalization ran the "+" kernel
|
|
# with a functional + non-functional tensor, and wrapped the output appropriately.
|
|
self.assertExpectedInline(
|
|
"\n".join(logs),
|
|
"""\
|
|
$2: f32[4] = torch._ops.aten.add.Tensor($0, $1)
|
|
$3: f32[4] = torch._ops.aten.add.Tensor($2, 1)""",
|
|
)
|
|
|
|
def test_mixed_wrappers_invalid(self):
|
|
x1_not_functional = torch.ones(4)
|
|
x2_functional = torch._to_functional_tensor(torch.ones(4))
|
|
|
|
# When dealing with mixed functional + non functional tensors,
|
|
# normal_tensor.add_(functional_tensor) is not valid
|
|
# because normal_tensor would need to be "promoted" to a functional tensor.
|
|
with self.assertRaises(RuntimeError):
|
|
x1_not_functional.add_(x2_functional)
|
|
|
|
def test_index_mutation_on_non_input(self):
|
|
def f(x):
|
|
tmp = torch.zeros(10)
|
|
tmp[5].fill_(1)
|
|
return tmp
|
|
|
|
self.assert_functionalization(f, torch.ones(2))
|
|
logs = self.get_logs(f, torch.ones(2))
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([10], device = device(type='cpu'), pin_memory = False)
|
|
select_copy = torch.ops.aten.select_copy.int(zeros, 0, 5)
|
|
fill = torch.ops.aten.fill.Scalar(select_copy, 1); select_copy = None
|
|
select_scatter = torch.ops.aten.select_scatter.default(zeros, fill, 0, 5); zeros = fill = None
|
|
select_copy_1 = torch.ops.aten.select_copy.int(select_scatter, 0, 5); select_copy_1 = None
|
|
return select_scatter
|
|
""",
|
|
) # noqa: B950
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f, torch.ones(2), reapply_views=True, run_reinplace=True
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1):
|
|
zeros = torch.ops.aten.zeros.default([10], device = device(type='cpu'), pin_memory = False)
|
|
select = torch.ops.aten.select.int(zeros, 0, 5)
|
|
fill = torch.ops.aten.fill_.Scalar(select, 1); select = fill = None
|
|
select_1 = torch.ops.aten.select.int(zeros, 0, 5); select_1 = None
|
|
return zeros
|
|
""",
|
|
)
|
|
|
|
def test_instance_norm(self):
|
|
size = 100
|
|
|
|
def f(x, running_mean, running_var):
|
|
with enable_python_dispatcher():
|
|
return torch.instance_norm(
|
|
x,
|
|
None,
|
|
None,
|
|
running_mean,
|
|
running_var,
|
|
use_input_stats=True,
|
|
momentum=0.1,
|
|
eps=1e-5,
|
|
cudnn_enabled=False,
|
|
)
|
|
|
|
self.assert_functionalization(
|
|
f, torch.randn(20, size, 35, 45), torch.zeros(size), torch.ones(size)
|
|
)
|
|
# On Windows, for instance_norm, the alias_copy's are reordered to come right before they need to be used
|
|
# whereas on other platforms, the alias_copy's are before the view_copy's.
|
|
# e.g., the alias_copy after the getitem_4 assignment would be moved to be right before the copy assignment.
|
|
if not IS_WINDOWS:
|
|
logs = self.get_logs(
|
|
f, torch.randn(20, size, 35, 45), torch.zeros(size), torch.ones(size)
|
|
)
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1, arg1_1, arg2_1):
|
|
repeat = torch.ops.aten.repeat.default(arg1_1, [20])
|
|
repeat_1 = torch.ops.aten.repeat.default(arg2_1, [20])
|
|
view_copy = torch.ops.aten.view_copy.default(arg0_1, [1, 2000, 35, 45]); arg0_1 = None
|
|
empty = torch.ops.aten.empty.memory_format([0], dtype = torch.uint8, layout = torch.strided, device = device(type='cpu')); empty = None
|
|
_native_batch_norm_legit_functional = torch.ops.aten._native_batch_norm_legit_functional.default(view_copy, None, None, repeat, repeat_1, True, 0.1, 1e-05); view_copy = repeat = repeat_1 = None
|
|
getitem = _native_batch_norm_legit_functional[0]
|
|
getitem_1 = _native_batch_norm_legit_functional[1]; getitem_1 = None
|
|
getitem_2 = _native_batch_norm_legit_functional[2]; getitem_2 = None
|
|
getitem_3 = _native_batch_norm_legit_functional[3]
|
|
getitem_4 = _native_batch_norm_legit_functional[4]; _native_batch_norm_legit_functional = None
|
|
alias_copy = torch.ops.aten.alias_copy.default(arg1_1)
|
|
view_copy_1 = torch.ops.aten.view_copy.default(getitem_3, [20, 100]); view_copy_1 = None
|
|
view_copy_2 = torch.ops.aten.view_copy.default(getitem_3, [20, 100]); getitem_3 = None
|
|
mean = torch.ops.aten.mean.dim(view_copy_2, [0]); view_copy_2 = None
|
|
copy = torch.ops.aten.copy.default(alias_copy, mean); alias_copy = mean = None
|
|
alias_copy_1 = torch.ops.aten.alias_copy.default(copy); copy = None
|
|
alias_copy_2 = torch.ops.aten.alias_copy.default(alias_copy_1); alias_copy_2 = None
|
|
alias_copy_3 = torch.ops.aten.alias_copy.default(arg2_1)
|
|
view_copy_3 = torch.ops.aten.view_copy.default(getitem_4, [20, 100]); view_copy_3 = None
|
|
view_copy_4 = torch.ops.aten.view_copy.default(getitem_4, [20, 100]); getitem_4 = None
|
|
mean_1 = torch.ops.aten.mean.dim(view_copy_4, [0]); view_copy_4 = None
|
|
copy_1 = torch.ops.aten.copy.default(alias_copy_3, mean_1); alias_copy_3 = mean_1 = None
|
|
alias_copy_4 = torch.ops.aten.alias_copy.default(copy_1); copy_1 = None
|
|
alias_copy_5 = torch.ops.aten.alias_copy.default(alias_copy_4); alias_copy_5 = None
|
|
view_copy_5 = torch.ops.aten.view_copy.default(getitem, [20, 100, 35, 45]); getitem = None
|
|
copy_ = torch.ops.aten.copy_.default(arg1_1, alias_copy_1); arg1_1 = alias_copy_1 = copy_ = None
|
|
copy__1 = torch.ops.aten.copy_.default(arg2_1, alias_copy_4); arg2_1 = alias_copy_4 = copy__1 = None
|
|
return view_copy_5
|
|
""", # noqa: B950
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f,
|
|
torch.randn(20, size, 35, 45),
|
|
torch.zeros(size),
|
|
torch.ones(size),
|
|
reapply_views=True,
|
|
run_reinplace=True,
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1, arg1_1, arg2_1):
|
|
repeat = torch.ops.aten.repeat.default(arg1_1, [20])
|
|
repeat_1 = torch.ops.aten.repeat.default(arg2_1, [20])
|
|
view = torch.ops.aten.view.default(arg0_1, [1, 2000, 35, 45]); arg0_1 = None
|
|
empty = torch.ops.aten.empty.memory_format([0], dtype = torch.uint8, layout = torch.strided, device = device(type='cpu')); empty = None
|
|
_native_batch_norm_legit_functional = torch.ops.aten._native_batch_norm_legit_functional.default(view, None, None, repeat, repeat_1, True, 0.1, 1e-05); view = repeat = repeat_1 = None
|
|
getitem = _native_batch_norm_legit_functional[0]
|
|
getitem_1 = _native_batch_norm_legit_functional[1]; getitem_1 = None
|
|
getitem_2 = _native_batch_norm_legit_functional[2]; getitem_2 = None
|
|
getitem_3 = _native_batch_norm_legit_functional[3]
|
|
getitem_4 = _native_batch_norm_legit_functional[4]; _native_batch_norm_legit_functional = None
|
|
alias = torch.ops.aten.alias.default(arg1_1)
|
|
view_1 = torch.ops.aten.view.default(getitem_3, [20, 100]); view_1 = None
|
|
view_2 = torch.ops.aten.view.default(getitem_3, [20, 100]); getitem_3 = None
|
|
mean = torch.ops.aten.mean.dim(view_2, [0]); view_2 = None
|
|
copy = torch.ops.aten.copy.default(alias, mean); alias = mean = None
|
|
alias_1 = torch.ops.aten.alias.default(copy); copy = None
|
|
alias_2 = torch.ops.aten.alias.default(alias_1); alias_2 = None
|
|
alias_3 = torch.ops.aten.alias.default(arg2_1)
|
|
view_3 = torch.ops.aten.view.default(getitem_4, [20, 100]); view_3 = None
|
|
view_4 = torch.ops.aten.view.default(getitem_4, [20, 100]); getitem_4 = None
|
|
mean_1 = torch.ops.aten.mean.dim(view_4, [0]); view_4 = None
|
|
copy_1 = torch.ops.aten.copy.default(alias_3, mean_1); alias_3 = mean_1 = None
|
|
alias_4 = torch.ops.aten.alias.default(copy_1); copy_1 = None
|
|
alias_5 = torch.ops.aten.alias.default(alias_4); alias_5 = None
|
|
view_5 = torch.ops.aten.view.default(getitem, [20, 100, 35, 45]); getitem = None
|
|
copy_ = torch.ops.aten.copy_.default(arg1_1, alias_1); arg1_1 = alias_1 = copy_ = None
|
|
copy__1 = torch.ops.aten.copy_.default(arg2_1, alias_4); arg2_1 = alias_4 = copy__1 = None
|
|
return view_5
|
|
""", # noqa: B950
|
|
)
|
|
|
|
def test_mutation_overlapping_mem(self):
|
|
def fn(x):
|
|
# x: (1, 5)
|
|
t1 = torch.add(x, x)
|
|
t2 = t1.unfold(1, 3, 2)
|
|
t3 = t2.abs_()
|
|
return t3
|
|
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"encountered a tensor being mutated that has internal overlap",
|
|
):
|
|
x = torch.ones(1, 5)
|
|
out = _functionalize(fn, reapply_views=True, crossref=False)(x)
|
|
|
|
def test_batch_norm(self):
|
|
def f(x, running_mean, running_var):
|
|
with enable_python_dispatcher():
|
|
return torch.batch_norm(
|
|
x, None, None, running_mean, running_var, True, 0.1, 1e-5, False
|
|
)
|
|
|
|
self.assert_functionalization(
|
|
f, torch.randn(20, 100, 35, 45), torch.zeros(100), torch.ones(100)
|
|
)
|
|
logs = self.get_logs(
|
|
f, torch.randn(20, 100, 35, 45), torch.zeros(100), torch.ones(100)
|
|
)
|
|
self.assertExpectedInline(
|
|
logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1, arg1_1, arg2_1):
|
|
empty = torch.ops.aten.empty.memory_format([0], dtype = torch.uint8, layout = torch.strided, device = device(type='cpu')); empty = None
|
|
_native_batch_norm_legit_functional = torch.ops.aten._native_batch_norm_legit_functional.default(arg0_1, None, None, arg1_1, arg2_1, True, 0.1, 1e-05); arg0_1 = None
|
|
getitem = _native_batch_norm_legit_functional[0]
|
|
getitem_1 = _native_batch_norm_legit_functional[1]; getitem_1 = None
|
|
getitem_2 = _native_batch_norm_legit_functional[2]; getitem_2 = None
|
|
getitem_3 = _native_batch_norm_legit_functional[3]
|
|
getitem_4 = _native_batch_norm_legit_functional[4]; _native_batch_norm_legit_functional = None
|
|
copy_ = torch.ops.aten.copy_.default(arg1_1, getitem_3); arg1_1 = getitem_3 = copy_ = None
|
|
copy__1 = torch.ops.aten.copy_.default(arg2_1, getitem_4); arg2_1 = getitem_4 = copy__1 = None
|
|
return getitem
|
|
""", # noqa: B950
|
|
)
|
|
|
|
reinplaced_logs = self.get_logs(
|
|
f,
|
|
torch.randn(20, 100, 35, 45),
|
|
torch.zeros(100),
|
|
torch.ones(100),
|
|
reapply_views=True,
|
|
run_reinplace=True,
|
|
)
|
|
self.assertExpectedInline(
|
|
reinplaced_logs,
|
|
"""\
|
|
|
|
|
|
|
|
def forward(self, arg0_1, arg1_1, arg2_1):
|
|
empty = torch.ops.aten.empty.memory_format([0], dtype = torch.uint8, layout = torch.strided, device = device(type='cpu')); empty = None
|
|
_native_batch_norm_legit_functional = torch.ops.aten._native_batch_norm_legit_functional.default(arg0_1, None, None, arg1_1, arg2_1, True, 0.1, 1e-05); arg0_1 = None
|
|
getitem = _native_batch_norm_legit_functional[0]
|
|
getitem_1 = _native_batch_norm_legit_functional[1]; getitem_1 = None
|
|
getitem_2 = _native_batch_norm_legit_functional[2]; getitem_2 = None
|
|
getitem_3 = _native_batch_norm_legit_functional[3]
|
|
getitem_4 = _native_batch_norm_legit_functional[4]; _native_batch_norm_legit_functional = None
|
|
copy_ = torch.ops.aten.copy_.default(arg1_1, getitem_3); arg1_1 = getitem_3 = copy_ = None
|
|
copy__1 = torch.ops.aten.copy_.default(arg2_1, getitem_4); arg2_1 = getitem_4 = copy__1 = None
|
|
return getitem
|
|
""", # noqa: B950
|
|
)
|
|
|
|
# This tests our python shims around C++ Functionalization: FunctionalTensor and FunctionalTensorMode
|
|
def test_python_functionalization(self):
|
|
def f(x):
|
|
x_view = x.view(-1)
|
|
x.mul_(2)
|
|
return x_view + 1
|
|
|
|
def f_functionalized(x):
|
|
# Note [Disabling Functionalize TLS Above Python Functionalization]
|
|
# This UX is pretty annoying (although python functionalization's main customer is AOTAutograd,
|
|
# and is not really advertised as a user API).
|
|
# We need to explicitly disable functionalization when using python FunctionalTensor and FunctionalTensorMode.
|
|
# Why? FunctionalTensor is a wrapper tensor that holds an inner FunctionalTensorWrapper.
|
|
# Since the inner tensor has `DispatchKey.Functionalize` in its keyset, then by default,
|
|
# our FunctionalTensor will inherit the same keyset.
|
|
# We don't have an easy way of directly mutating a tensor's keyset from python,
|
|
# so globally disabling functionalization here is easier.
|
|
maybe_disable = torch._C._ExcludeDispatchKeyGuard(
|
|
torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
|
|
)
|
|
with maybe_disable, FunctionalTensorMode():
|
|
x_wrapped = FunctionalTensor.to_functional(x)
|
|
out_wrapped = f(x_wrapped)
|
|
out_unwrapped = out_wrapped.elem
|
|
torch._sync(out_unwrapped)
|
|
return torch._from_functional_tensor(out_unwrapped)
|
|
|
|
# Make a non-leaf
|
|
x = torch.randn(2, requires_grad=True) + 1
|
|
fx_g = make_fx(f_functionalized)(x)
|
|
# NB: view_1 below is expected (though unused) due to view replay. AOTAutograd runs a
|
|
# DCE pass that will remove nodes like this later on.
|
|
self.assertExpectedInline(
|
|
fx_g.code.strip(),
|
|
"""\
|
|
def forward(self, x_1):
|
|
view = torch.ops.aten.view.default(x_1, [-1]); view = None
|
|
mul = torch.ops.aten.mul.Tensor(x_1, 2); x_1 = None
|
|
view_1 = torch.ops.aten.view.default(mul, [-1]); view_1 = None
|
|
view_2 = torch.ops.aten.view.default(mul, [-1]); mul = None
|
|
add = torch.ops.aten.add.Tensor(view_2, 1); view_2 = None
|
|
return add""",
|
|
)
|
|
|
|
def test_python_functionalization_zero_tensor(self):
|
|
def f(x):
|
|
y = torch.ops.aten._efficientzerotensor([4])
|
|
out = x + y
|
|
out.mul_(2)
|
|
return out
|
|
|
|
x = torch.randn(4)
|
|
out_ref = f(x)
|
|
out_test = dispatch_functionalize(f)(x)
|
|
out_test_cpp = _functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)(x)
|
|
self.assertEqual(out_ref, out_test)
|
|
self.assertEqual(out_ref, out_test_cpp)
|
|
fx_g = make_fx(dispatch_functionalize(f))(x)
|
|
fx_g_cpp = make_fx(
|
|
_functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)
|
|
)(x)
|
|
self.assertEqual(fx_g_cpp.code.strip(), fx_g.code.strip())
|
|
|
|
def test_python_functionalization_is_conj(self):
|
|
def f(x):
|
|
out = x.conj()
|
|
return out, out.is_conj()
|
|
|
|
x = torch.randn(4, dtype=torch.complex64)
|
|
out_ref = f(x)
|
|
out_test = dispatch_functionalize(f)(x)
|
|
out_test_cpp = _functionalize(f, reapply_views=True, crossref=False)(x)
|
|
self.assertEqual(out_ref[0], out_test[0])
|
|
self.assertEqual(out_ref[1], out_test[1])
|
|
self.assertEqual(out_ref[0], out_test_cpp[0])
|
|
self.assertEqual(out_ref[1], out_test_cpp[1])
|
|
|
|
def test_python_functionalization_is_neg(self):
|
|
def f(x):
|
|
out = x.neg()
|
|
return out, out.is_neg()
|
|
|
|
x = torch.randn(4, dtype=torch.complex64)
|
|
out_ref = f(x)
|
|
out_test = dispatch_functionalize(f)(x)
|
|
out_test_cpp = _functionalize(f, reapply_views=True, crossref=False)(x)
|
|
self.assertEqual(out_ref[0], out_test[0])
|
|
self.assertEqual(out_ref[1], out_test[1])
|
|
self.assertEqual(out_ref[0], out_test_cpp[0])
|
|
self.assertEqual(out_ref[1], out_test_cpp[1])
|
|
|
|
def test_python_functionalization_conj(self):
|
|
def f(x):
|
|
y = x.clone().conj()
|
|
y.mul_(2)
|
|
return torch.view_as_real(y.resolve_conj())
|
|
|
|
x = torch.randn(4, dtype=torch.complex64)
|
|
out_ref = f(x)
|
|
out_test = dispatch_functionalize(f)(x)
|
|
out_test_cpp = _functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)(x)
|
|
self.assertEqual(out_ref, out_test)
|
|
self.assertEqual(out_test, out_test_cpp)
|
|
fx_g = make_fx(dispatch_functionalize(f))(x)
|
|
fx_g_cpp = make_fx(
|
|
_functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)
|
|
)(x)
|
|
self.assertExpectedInline(
|
|
fx_g.code.strip(),
|
|
"""\
|
|
def forward(self, arg0_1):
|
|
clone = torch.ops.aten.clone.default(arg0_1); arg0_1 = None
|
|
_conj = torch.ops.aten._conj.default(clone); clone = None
|
|
clone_1 = torch.ops.aten.clone.default(_conj)
|
|
mul = torch.ops.aten.mul.Tensor(clone_1, 2); clone_1 = None
|
|
clone_2 = torch.ops.aten.clone.default(_conj); _conj = None
|
|
copy = torch.ops.aten.copy.default(clone_2, mul); clone_2 = mul = None
|
|
_conj_1 = torch.ops.aten._conj.default(copy); copy = None
|
|
_conj_2 = torch.ops.aten._conj.default(_conj_1); _conj_1 = None
|
|
clone_3 = torch.ops.aten.clone.default(_conj_2); _conj_2 = None
|
|
view_as_real = torch.ops.aten.view_as_real.default(clone_3); clone_3 = None
|
|
return view_as_real""",
|
|
)
|
|
self.assertEqual(fx_g_cpp.code.strip(), fx_g.code.strip())
|
|
|
|
def test_python_functionalization_neg(self):
|
|
def f(x):
|
|
y = x._neg_view()
|
|
z = y.resolve_neg()
|
|
return z + 1
|
|
|
|
x = torch.randn(4)
|
|
out_ref = f(x)
|
|
out_test = dispatch_functionalize(f)(x)
|
|
out_test_cpp = _functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)(x)
|
|
self.assertEqual(out_ref, out_test)
|
|
self.assertEqual(out_ref, out_test_cpp)
|
|
fx_g = make_fx(dispatch_functionalize(f))(x)
|
|
fx_g_cpp = make_fx(
|
|
_functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)
|
|
)(x)
|
|
self.assertExpectedInline(
|
|
fx_g.code.strip(),
|
|
"""\
|
|
def forward(self, arg0_1):
|
|
_neg_view = torch.ops.aten._neg_view.default(arg0_1); arg0_1 = None
|
|
clone = torch.ops.aten.clone.default(_neg_view); _neg_view = None
|
|
add = torch.ops.aten.add.Tensor(clone, 1); clone = None
|
|
return add""",
|
|
)
|
|
self.assertEqual(fx_g_cpp.code.strip(), fx_g.code.strip())
|
|
|
|
def test_python_functionalization_lift_fresh_storage(self):
|
|
unlifted = torch.tensor([0.0])
|
|
|
|
maybe_disable = torch._C._ExcludeDispatchKeyGuard(
|
|
torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
|
|
)
|
|
with maybe_disable, FunctionalTensorMode():
|
|
lifted = torch.ops.aten.lift_fresh.default(unlifted)
|
|
|
|
self.assertNotEqual(unlifted.untyped_storage(), lifted.untyped_storage())
|
|
|
|
def test_python_functionalization_lift_fresh(self):
|
|
def f(x):
|
|
tmp = torch.tensor([0.0])
|
|
return tmp + x
|
|
|
|
x = torch.randn(4)
|
|
out_ref = f(x)
|
|
out_test = dispatch_functionalize(f)(x)
|
|
out_test_cpp = _functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)(x)
|
|
self.assertEqual(out_ref, out_test)
|
|
self.assertEqual(out_ref, out_test_cpp)
|
|
fx_g = make_fx(dispatch_functionalize(f))(x)
|
|
fx_g_cpp = make_fx(
|
|
_functionalize(
|
|
f, reapply_views=True, crossref=False, skip_input_mutations=True
|
|
)
|
|
)(x)
|
|
self.assertExpectedInline(
|
|
fx_g.code.strip(),
|
|
"""\
|
|
def forward(self, arg0_1):
|
|
_tensor_constant0 = self._tensor_constant0
|
|
lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0); _tensor_constant0 = None
|
|
add = torch.ops.aten.add.Tensor(lift_fresh_copy, arg0_1); lift_fresh_copy = arg0_1 = None
|
|
return add""",
|
|
)
|
|
self.assertEqual(fx_g_cpp.code.strip(), fx_g.code.strip())
|
|
|
|
|
|
@xfail_inherited_tests(
|
|
[
|
|
"test_as_strided",
|
|
"test_copy_",
|
|
"test_diagonal",
|
|
"test_diagonal_mutated_input",
|
|
"test_everything",
|
|
"test_fill_",
|
|
"test_slice",
|
|
"test_split",
|
|
"test_split_with_sizes",
|
|
"test_unbind",
|
|
"test_view_clone_view_inplace",
|
|
"test_view_inplace",
|
|
]
|
|
)
|
|
@unittest.skipIf(
|
|
TEST_WITH_TORCHDYNAMO, "dynamo-ing code with proxy + fake doesnt work well"
|
|
)
|
|
class TestCrossRefFunctionalization(TestFunctionalization):
|
|
crossref = True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|