[dynamo] Fix constant propagation in builtins and UserClasses (pytorch#131354)

* Fixes pytorch#118675
* Replaces pytorch#118994

Pull Request resolved: pytorch#131354
Approved by: https://github.com/jansel, https://github.com/anijain2305

Author: rec

test/dynamo/test_higher_order_ops.py

@@ -3818,10 +3818,10 @@ def wrapper_fn(model, params, buffers, inputs):
         if torch._dynamo.config.inline_inbuilt_nn_modules:
             expected = """\
 class GraphModule(torch.nn.Module):
-    def forward(self, L_params_l1_weight_: "f32[1, 1]", L_params_l1_bias_: "f32[1]", L_buffers_buffer_: "f32[1]", L_inputs_: "f32[1, 1]"):
+    def forward(self, L_params_l1_weight_: "f32[1, 1]", L_buffers_buffer_: "f32[1]", L_params_l1_bias_: "f32[1]", L_inputs_: "f32[1, 1]"):
         l_params_l1_weight_ = L_params_l1_weight_
-        l_params_l1_bias_ = L_params_l1_bias_
         l_buffers_buffer_ = L_buffers_buffer_
+        l_params_l1_bias_ = L_params_l1_bias_
         l_inputs_ = L_inputs_
 
         linear: "f32[1, 1]" = torch._C._nn.linear(l_inputs_, l_params_l1_weight_, l_params_l1_bias_);  l_inputs_ = l_params_l1_weight_ = l_params_l1_bias_ = None
@@ -6005,7 +6005,7 @@ def wrapper_fn(x, y):
             return torch.func.vmap(f)(x, y)
 
         actual = wrapper_fn(x, y)
-        expected = torch.compile(wrapper_fn, backend="aot_eager", fullgraph=False)(x, y)
+        expected = torch.compile(wrapper_fn, backend="aot_eager")(x, y)
         self.assertEqual(len(counters["graph_break"]), 0)
         self.assertEqual(actual, expected)
         self.assertEqual(some_list, [1, 1])

test/dynamo/test_misc.py

@@ -101,6 +101,12 @@
 TPFLAGS_MAPPING = 1 << 6
 
 
+# A class defined in the global scope, used in MiscTests.test_const_getattr
+class _B:
+    def __init__(self):
+        pass
+
+
 # Specializes a test to run only if translation validation is set.
 def onlyIfTranslationValidation(fn: typing.Callable) -> typing.Callable:
     @functools.wraps(fn)
@@ -1412,6 +1418,28 @@ def fn(x, s):
         # One recompile per differing input type
         self.assertEqual(cnts.frame_count, 3)
 
+    def test_const_getattr(self):
+        # See https://github.com/pytorch/pytorch/issues/118675
+        def fn(x):
+            y = x[f"{_B.__module__}.{_B.__name__}"]
+            z = x[f"{_B.__class__.__module__}.{_B.__name__}"]
+            u = x[f"{_B.__class__.__module__}.{_B.__class__.__qualname__}"]
+            return y + z + u
+
+        args = (
+            {
+                f"{_B.__module__}._B": torch.randn(10),
+                "builtins._B": torch.randn(10),
+                "builtins.type": torch.randn(10),
+            },
+        )
+
+        cnts = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch._dynamo.optimize(cnts)(fn)
+
+        self.assertEqual(fn(*args), opt_fn(*args))
+        self.assertEqual(cnts.frame_count, 1)
+
     def test_cell_output1(self):
         out = None
 

test/dynamo_expected_failures/TestArrayCreationCopyArgument.test_striding_not_ok

@@ -315,7 +315,7 @@ def uninteresting_files():
 _CLOSURE_VARS: Optional[Dict[str, object]] = None
 
 
-def _get_closure_vars():
+def get_closure_vars():
     global _CLOSURE_VARS
     if _CLOSURE_VARS is None:
         _CLOSURE_VARS = {
@@ -1150,7 +1150,7 @@ def add_python_lambda_leaf_guard_to_root(
         is_epilogue=True,
     ):
         if closure_vars is None:
-            closure_vars = _get_closure_vars()
+            closure_vars = get_closure_vars()
         # Adds a lambda leaf guard to the root guard manager. It wraps the
         # code_parts in a function object which is then passed on to the leaf
         # guard.
@@ -1177,7 +1177,7 @@ def add_python_lambda_leaf_guard_to_root(
     # (like its type) which is what you permanently install into the
     # guard code.
     def get(self, name: str) -> Any:
-        return eval(name, self.scope, _get_closure_vars())
+        return eval(name, self.scope, get_closure_vars())
 
     # Registers the usage of the source name referenced by the
     # string (or stored in the Guard) as being guarded upon.  It's important
@@ -1497,7 +1497,7 @@ def EQUALS_MATCH(self, guard: Guard):
             self._set_guard_export_info(guard, code)
 
             self.get_guard_manager(guard).add_lambda_guard(
-                _get_closure_vars()["__math_isnan"],
+                get_closure_vars()["__math_isnan"],
                 get_verbose_code_parts(code, guard),
             )
             return
@@ -1510,7 +1510,7 @@ def EQUALS_MATCH(self, guard: Guard):
             self._set_guard_export_info(guard, code)
 
             self.get_guard_manager(guard).add_lambda_guard(
-                _get_closure_vars()["__numpy_isnan"],
+                get_closure_vars()["__numpy_isnan"],
                 get_verbose_code_parts(code, guard),
             )
             return
@@ -1786,7 +1786,7 @@ def get_sources(t_id, dim):
         self.add_python_lambda_leaf_guard_to_root(
             code_parts,
             verbose_code_parts,
-            closure_vars={**SYMPY_INTERP, **_get_closure_vars()},
+            closure_vars={**SYMPY_INTERP, **get_closure_vars()},
         )
 
     def TENSOR_MATCH(self, guard: Guard, value=None):
@@ -2375,7 +2375,7 @@ def add_code_part(code_part, guard, log_only=False):
             "___check_global_state": global_state.check,
             "___check_torch_function_mode_stack": torch_function_mode_stack_check_fn,
             **SYMPY_INTERP,
-            **_get_closure_vars(),
+            **get_closure_vars(),
         }
 
         globals_for_guard_fn = {"G": builder.scope["G"]}

torch/_dynamo/guards.py

@@ -6,9 +6,9 @@
 
 from .. import variables
 from ..current_scope_id import current_scope_id
-from ..exc import unimplemented
+from ..exc import unimplemented, Unsupported
 from ..source import AttrSource, Source
-from ..utils import istype
+from ..utils import is_function_or_wrapper, istype
 
 
 if TYPE_CHECKING:
@@ -238,16 +238,24 @@ def make_guard(self, fn):
         raise NotImplementedError
 
     def const_getattr(self, tx: "InstructionTranslator", name: str) -> Any:
-        """getattr(self, name) returning a python constant"""
-        raise NotImplementedError
+        v = self.as_python_constant()
+        try:
+            return getattr(v, name)
+        except AttributeError:
+            raise NotImplementedError from None
 
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
         """getattr(self, name) returning a new variable"""
-        value = self.const_getattr(tx, name)
-        if not variables.ConstantVariable.is_literal(value):
-            raise NotImplementedError
+        from .misc import GetAttrVariable
+
         source = self.source and AttrSource(self.source, name)
-        return variables.ConstantVariable.create(value, source=source)
+        try:
+            value = self.const_getattr(tx, name)
+            if not is_function_or_wrapper(value):
+                return VariableTracker.build(tx, value, source)
+        except (NotImplementedError, Unsupported):
+            pass
+        return GetAttrVariable(self, name, source=source)
 
     def is_proxy(self):
         try:

torch/_dynamo/variables/base.py

@@ -1714,14 +1714,10 @@ def call_getattr(
             if config.replay_record_enabled:
                 tx.exec_recorder.record_module_access(obj.value, name, member)
             return VariableTracker.build(tx, member, source)
-
         elif istype(obj, UserFunctionVariable) and name in ("__name__", "__module__"):
             return ConstantVariable.create(getattr(obj.fn, name))
         else:
-            try:
-                return obj.var_getattr(tx, name)
-            except NotImplementedError:
-                return GetAttrVariable(obj, name, source=source)
+            return obj.var_getattr(tx, name)
 
     def call_setattr(
         self,

torch/_dynamo/variables/builtin.py

@@ -1,7 +1,7 @@
 # mypy: ignore-errors
 import functools
 import inspect
-from typing import Dict, List, TYPE_CHECKING
+from typing import Any, Dict, List, TYPE_CHECKING
 
 import torch
 from torch.fx.experimental._backward_state import BackwardState
@@ -214,6 +214,9 @@ def is_device_mesh(value):
     def as_python_constant(self):
         return self.value
 
+    def const_getattr(self, tx: "InstructionTranslator", name: str) -> Any:
+        raise NotImplementedError
+
     def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
         if name == "ndim":
             return ConstantVariable.create(self.value.ndim)

torch/_dynamo/variables/distributed.py

@@ -164,7 +164,8 @@ def __init__(self, fn, is_constant=False, **kwargs) -> None:
             self.is_constant = False
 
         assert isinstance(
-            fn, (types.FunctionType, torch.jit.ScriptFunction)
+            fn,
+            (types.BuiltinFunctionType, types.FunctionType, torch.jit.ScriptFunction),
         ), f"expected FunctionType found {typestr(fn)} {fn}"
         # TODO(anijain2305) - Replace directly calling UserFunctionVariable with
         # VariableBuilder, which handles the wrapping of _torchdynamo_inline.

torch/_dynamo/variables/functions.py

@@ -82,7 +82,10 @@ def debug_repr_helper(self, prefix, suffix):
         return prefix + ", ".join(i.debug_repr() for i in self.items) + suffix
 
     def as_python_constant(self):
-        return self.python_type()([x.as_python_constant() for x in self.items])
+        try:
+            return self.python_type()([x.as_python_constant() for x in self.items])
+        except RecursionError:
+            unimplemented(f"recursive containment {self}")
 
     def as_proxy(self):
         assert self.python_type() is not SizeVariable

torch/_dynamo/variables/lists.py

@@ -9,7 +9,7 @@
 import sys
 import types
 import warnings
-from typing import Dict, List, Optional, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
 import torch._C
 import torch._numpy as tnp
@@ -1617,6 +1617,9 @@ def python_type(self):
     def as_python_constant(self):
         return self.random
 
+    def const_getattr(self, tx: "InstructionTranslator", name: str) -> Any:
+        raise NotImplementedError
+
     @staticmethod
     def is_supported_random_obj(val):
         if type(val) is not random.Random:

torch/_dynamo/variables/misc.py

@@ -229,7 +229,6 @@ def dynamic_getattr(self, tx: "InstructionTranslator", name):
         # (1) the tensor is a traceable tensor subclass
         # (2) We are getattr'ing an inner tensor from that subclass
         if not self.source and is_traceable_wrapper_subclass(fake_val):
-            fake_val = self.proxy.node.meta["example_value"]
             attrs, ctx = fake_val.__tensor_flatten__()
             proxy = getattr(self.as_proxy(), name)
             example_value = getattr(fake_val, name)
@@ -245,14 +244,19 @@ def dynamic_getattr(self, tx: "InstructionTranslator", name):
                 return VariableTracker.build(tx, example_value)
 
         if not (self.source and self.source.subguards_allowed()):
-            raise NotImplementedError
+            return
+
+        from ..guards import get_closure_vars, GuardBuilder
 
         # For local source, we associate the real value. We use this real value
         # for implementing getattr fallthrough on the variable tracker base class.
-
         # Note - this scope construction is mirrored in guards
         # A subsequent PR will introduce a util.
-        scope = {"L": tx.output.local_scope, "G": tx.output.global_scope}
+        scope = {
+            "L": tx.output.local_scope,
+            "G": tx.output.global_scope,
+            **get_closure_vars(),
+        }
         try:
             # We raise in case we get a typerror bug w/ SuperSource.
             # SuperSource has bugs in it atm, and can produce code like
@@ -261,23 +265,24 @@ def dynamic_getattr(self, tx: "InstructionTranslator", name):
             # Which is incorrect, and violates the invariant that all sources should be eval()-able against the scope.
             _input_associated_real_value = eval(self.source.name(), scope)
         except Exception as exc:
-            raise NotImplementedError from exc
+            msg = f"{exc!r} raised in eval('{self.source.name()}')"
+            raise NotImplementedError(msg) from exc
 
+        real_value = getattr(_input_associated_real_value, name)
         if _input_associated_real_value is None:
-            raise NotImplementedError
+            return
 
         if object_has_getattribute(_input_associated_real_value):
-            raise NotImplementedError
+            return
 
         if get_custom_getattr(_input_associated_real_value):
-            raise NotImplementedError
+            return
 
-        real_value = getattr(_input_associated_real_value, name)
         if callable(real_value):
             # Callables have more nuanced handling, and we should let the existing system delegate here.
             # Raising was past behavior and so should always be sound to fall back.
             # Note - at a certain point we may want to handle
-            raise NotImplementedError
+            return
 
         attr_source = AttrSource(self.source, name)
         install_guard(attr_source.make_guard(GuardBuilder.HASATTR))
@@ -1215,8 +1220,6 @@ def var_getattr(self, tx: "InstructionTranslator", name):
         from ..utils import numpy_attr_wrapper
         from .builder import wrap_fx_proxy
 
-        result = None
-
         example_value = self.as_proxy().node.meta["example_value"]
         example_ndarray = tnp.ndarray(example_value)
 
@@ -1235,7 +1238,7 @@ def insert_into_graph():
                 (self.as_proxy(), name),
                 {},
             )
-            result = NumpyNdarrayVariable.create(tx, proxy)
+            return NumpyNdarrayVariable.create(tx, proxy)
 
         # These are awkward to implement.  The standard playbook for torch._numpy
         # interop is to trace a call into the torch._numpy wrapper which works for
@@ -1264,9 +1267,8 @@ def insert_into_graph():
             unimplemented(f"TODO: add support for ndarray.{name}")
         elif name in ["__version__"]:
             unimplemented("delegate np.__version__ to NumPy")
-        if result is None:
-            raise NotImplementedError
-        return result
+        else:
+            return super().var_getattr(tx, name)
 
     @staticmethod
     def patch_args(name, args, kwargs):

torch/_dynamo/variables/tensor.py

@@ -181,10 +181,7 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracke
         ):
             return super().var_getattr(tx, name)
 
-        try:
-            obj = inspect.getattr_static(self.value, name)
-        except AttributeError:
-            obj = None
+        obj = inspect.getattr_static(self.value, name, None)
 
         if isinstance(obj, staticmethod):
             return VariableTracker.build(tx, obj.__get__(self.value), source)
@@ -206,6 +203,9 @@ def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracke
             ):
                 return VariableTracker.build(tx, obj.__get__(self.value), source)
 
+        if inspect.ismemberdescriptor(obj) or inspect.isdatadescriptor(obj):
+            value = getattr(self.value, name)
+            return VariableTracker.build(tx, value, source)
         if ConstantVariable.is_literal(obj):
             return ConstantVariable.create(obj)
         elif isinstance(obj, enum.Enum):