fix issue with lift_fresh_copy when using export + compile (#108243)

bdhirsh · pytorchmergebot · commit fac7a1f73048 · 2023-09-05T20:02:35.000Z
Fixes #105327. The problem is that `lift_fresh_copy()`'s functionalization implementation currently assumes that the input is always not functional. This is apparently too limiting: when you have "user" code like this (which can potentially come from exporting a model and then running compile on the resulting graph): ``` tensor_constant0 = torch.tensor(2) lift_fresh = torch.ops.aten.lift_fresh_copy.default(tensor_constant0) ``` When we run this through AOTAutograd, the first call (torch.tensor(2)) will **already** be lifted into a functional tensor wrapper - so the `lift_fresh_copy` call doesn't need to do any "lifting" anymore - it just needs to do a clone. Pull Request resolved: #108243 Approved by: https://github.com/albanD ghstack dependencies: #108081, #108235
diff --git a/aten/src/ATen/FunctionalizeFallbackKernel.cpp b/aten/src/ATen/FunctionalizeFallbackKernel.cpp
@@ -181,14 +181,26 @@ static at::Tensor lift_functionalize(const at::Tensor & self) {
 }
 
 static at::Tensor lift_fresh_functionalize(const at::Tensor & self) {
-  TORCH_INTERNAL_ASSERT(!at::functionalization::impl::isFunctionalTensor(self));
+  // See Note [Exporting and compiling a graph with lift_fresh_copy]
+  if (at::functionalization::impl::isFunctionalTensor(self)) {
+    return self.view_as(self);
+  }
+
   at::AutoDispatchSkipFunctionalize guard;
   auto out = at::lift_fresh(self);
   return at::functionalization::impl::to_functional_tensor(out);
 }
 
 static at::Tensor lift_fresh_functionalize_copy(const at::Tensor & self) {
-  TORCH_INTERNAL_ASSERT(!at::functionalization::impl::isFunctionalTensor(self));
+  // Note [Exporting and compiling a graph with lift_fresh_copy]
+  // If out is already a functional tensor, don't wrap it twice.
+  // In theory this could be useful if we want to nest functionalization with itself,
+  // but that isn't really a use case today.
+  // Needed for https://github.com/pytorch/pytorch/issues/105327
+  if (at::functionalization::impl::isFunctionalTensor(self)) {
+    return self.clone();
+  }
+
   at::AutoDispatchSkipFunctionalize guard;
   auto out = at::lift_fresh_copy(self);
   return at::functionalization::impl::to_functional_tensor(out);
diff --git a/test/functorch/test_aotdispatch.py b/test/functorch/test_aotdispatch.py
@@ -2779,6 +2779,27 @@ def forward(self, x, y):
         assert torch.allclose(inputs[0].grad, cloned_inputs[0].grad)
         assert torch.allclose(inputs[1].grad, cloned_inputs[1].grad)
 
+    # https://github.com/pytorch/pytorch/issues/105327
+    def test_lift_fresh_copy_in_graph(self):
+        class MyMod(torch.nn.Module):
+            def forward(self, x):
+                _tensor_constant0 = torch.tensor([1])
+                lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0)
+                y = x.mul(lift_fresh_copy)
+                return (y,)
+
+        mod = MyMod()
+        shape_env = ShapeEnv()
+        fake_mode = FakeTensorMode(shape_env=shape_env)
+        x = torch.ones(4, requires_grad=True)
+        inputs = [x]
+        fake_inputs = [fake_mode.from_tensor(x) for x in inputs]
+        compiled_f = aot_module_simplified(mod, fake_inputs, nop)
+
+        out_ref = mod(x)
+        out_test = compiled_f(x)
+        self.assertEqual(out_ref[0].detach(), out_test[0].detach())
+
     def test_inference_python_dispatcher(self):
         # Extracted from unet
         class MockModule(torch.nn.Module):
