[dynamo] fix graph break, improve hygeine - enforce using ConstantVariable for torch.device,torch.dtype (#112416)

jon-chuang · pytorchmergebot · commit 9bfebf754f90 · 2023-11-01T00:19:52.000Z
Fixes https://github.com/pytorch/pytorch/pull/112332/files#r1375690808 Simplify code paths, fix graph break ``` torch._dynamo.exc.InternalTorchDynamoError: TorchVariable() has no type ``` Pull Request resolved: #112416 Approved by: https://github.com/lezcano
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
@@ -7541,22 +7541,53 @@ def ff2(x):
         self.assertEqual(counter.op_count, 6)
 
     def test_torch_device_python_type(self):
+        for device, device_type, index in [
+            ("cpu", "cpu", None),
+            ("cuda:0", "cuda", 0),
+        ]:
+            if device == "cuda:0" and not TEST_CUDA:
+                continue
+
+            def fn(target):
+                target_device = target.device
+                a = torch.zeros(2, 3, device=target_device)
+                # Constant assert at trace time
+                assert isinstance(target_device, torch.device)
+                assert target_device.type == device_type
+                assert target_device.index == index
+                b = torch.zeros(2, 3, device=target_device)
+                c = torch.zeros(2, 3, device=target_device)
+                return a + b + c
+
+            from torch._dynamo.variables import ConstantVariable
+
+            device = torch.device(device)
+            expected_variable = ConstantVariable(device)
+            self.assertEqual(expected_variable.python_type(), type(device))
+
+            opt_func = torch._dynamo.optimize("eager", nopython=True)(fn)
+            a = torch.tensor([2, 3], device=device)
+            res = opt_func(a)
+            self.assertIsInstance(res, torch.Tensor)
+
+    def test_torch_dtype_python_type(self):
         def fn(target):
-            target_device = target.device
-            self.assertIsInstance(target_device, torch.device)
-            a = torch.zeros(2, 3, device=target_device)
-            b = torch.zeros(2, 3, device=target_device)
-            c = torch.zeros(2, 3, device=target_device)
+            target_dtype = target.dtype
+            a = torch.zeros(2, 3, dtype=target_dtype)
+            # Constant assert at trace time
+            assert isinstance(target_dtype, torch.dtype)
+            b = torch.zeros(2, 3, dtype=target_dtype)
+            c = torch.zeros(2, 3, dtype=target_dtype)
             return a + b + c
 
-        from torch._dynamo.variables import TorchVariable
+        from torch._dynamo.variables import ConstantVariable
 
-        device = torch.device("cpu")
-        expected_variable = TorchVariable(device)
-        self.assertEqual(expected_variable.python_type(), type(device))
+        dtype = torch.float16
+        expected_variable = ConstantVariable(dtype)
+        self.assertEqual(expected_variable.python_type(), type(dtype))
 
-        opt_func = torch._dynamo.optimize("inductor")(fn)
-        a = torch.tensor([2, 3], device=device)
+        opt_func = torch._dynamo.optimize("eager", nopython=True)(fn)
+        a = torch.tensor([2, 3], dtype=dtype)
         res = opt_func(a)
         self.assertIsInstance(res, torch.Tensor)
 
diff --git a/torch/_dynamo/variables/tensor.py b/torch/_dynamo/variables/tensor.py
@@ -233,9 +233,9 @@ def var_getattr(self, tx, name):
         if name == "ndim" and self.ndim is not None:
             result = ConstantVariable.create(self.ndim, **options)
         elif name == "dtype" and self.dtype is not None:
-            result = TorchVariable(self.dtype, **options)
+            result = ConstantVariable.create(self.dtype, **options)
         elif name == "device" and self.device is not None:
-            result = TorchVariable(self.device, **options)
+            result = ConstantVariable.create(self.device, **options)
         elif name == "layout" and self.layout is not None:
             result = TorchVariable(self.layout, **options)
         elif name == "is_cuda" and self.device is not None:
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
@@ -250,6 +250,9 @@ def __init__(self, value, **kwargs):
 
         self.value = value
 
+        assert not isinstance(
+            value, (torch.dtype, torch.device)
+        ), "should use ConstantVariable"
         # the remainder of this is just optional debug checks
         try:
             self_should_be_none = getattr(self.value, "__self__", None)
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -1662,7 +1662,16 @@ def check_skip_condition(node, parent, is_output):
 
 def make_fallback(op, layout_constraint=None, warn=True):
     assert op not in decompositions, f"both a fallback and a decomp for same op: {op}"
-    if get_decompositions([op]) and warn and bool(os.getenv("CI")):
+    if (
+        warn
+        and bool(os.getenv("CI"))
+        and get_decompositions([op])
+        # if fallback_random, we allow not decomposing random
+        and not (
+            config.fallback_random
+            and op in torch._decomp.decompositions_for_rng.extra_random_decomps
+        )
+    ):
         # Note: 'warn' is holdover from when this was a warning, but for ops that previously
         # set warn=False we do not want a CI error.
         # Ignore the 'suppress errors' configs in CI, as this particular warning happens on startup anyway and is not
