[dynamo][dynamic] Recompilation hint for nn module integer attributes (pytorch#154867)

For program like this

```
class Mod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.c = 0

    def forward(self, x):
        self.c += 1
        return x * self.c
```

You can check the recompile reasons at https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpzv9z6Q/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000

![image](https://github.com/user-attachments/assets/856a95fd-0533-4abc-a213-1f73ae2cb766)

Pull Request resolved: pytorch#154867
Approved by: https://github.com/zou3519

Author: anijain2305

torch/_dynamo/guards.py

@@ -1665,7 +1665,7 @@ def metadata_checker(x):
             metadata_checker, get_verbose_code_parts(global_name, guard)
         )
 
-    def EQUALS_MATCH(self, guard: Guard):
+    def EQUALS_MATCH(self, guard: Guard, recompile_hint: Optional[str] = None):
         ref = self.arg_ref(guard)
         val = self.get(guard.name)
         if np:
@@ -1762,9 +1762,14 @@ def EQUALS_MATCH(self, guard: Guard):
             # is immutable. For a few corner cases like sets and lists, we make a deepcopy to purposefully fail the
             # pointer equality check.
             val = deepcopy(val)
-        self.get_guard_manager(guard).add_equals_match_guard(
-            val, get_verbose_code_parts(code, guard)
-        )
+
+        verbose_code_parts = get_verbose_code_parts(code, guard)
+        if recompile_hint:
+            verbose_code_parts = [
+                f"{part} (HINT: {recompile_hint})" for part in verbose_code_parts
+            ]
+
+        self.get_guard_manager(guard).add_equals_match_guard(val, verbose_code_parts)
         self._set_guard_export_info(guard, code)
         return
 

torch/_dynamo/source.py

@@ -21,6 +21,7 @@
 
 import dataclasses
 import enum
+import functools
 from typing import Any, Optional, TYPE_CHECKING, Union
 
 from torch._guards import ChainedSource, GuardSource, Source
@@ -901,6 +902,7 @@ def is_from_source(source: Source, target: Source):
     return source == target
 
 
+@functools.lru_cache
 def is_from_unspecialized_nn_module_source(source: Source):
     if isinstance(source, UnspecializedNNModuleSource):
         return True
@@ -909,6 +911,16 @@ def is_from_unspecialized_nn_module_source(source: Source):
     return False
 
 
+@functools.lru_cache
+def is_from_unspecialized_builtin_nn_module_source(source: Source):
+    if isinstance(source, UnspecializedBuiltinNNModuleSource):
+        return True
+    if isinstance(source, ChainedSource):
+        return is_from_unspecialized_builtin_nn_module_source(source.base)
+    return False
+
+
+@functools.lru_cache
 def is_from_unspecialized_param_buffer_source(source: Source):
     if isinstance(source, UnspecializedParamBufferSource):
         return True
@@ -917,6 +929,7 @@ def is_from_unspecialized_param_buffer_source(source: Source):
     return False
 
 
+@functools.lru_cache
 def is_from_flatten_script_object_source(source: Source):
     if isinstance(source, FlattenScriptObjectSource):
         return True
@@ -925,6 +938,7 @@ def is_from_flatten_script_object_source(source: Source):
     return False
 
 
+@functools.lru_cache
 def is_from_optimizer_source(source: Source):
     if isinstance(source, OptimizerSource):
         return True
@@ -935,6 +949,7 @@ def is_from_optimizer_source(source: Source):
 
 # TODO: can probably write a generic "test this on everything in the chain"
 # helper
+@functools.lru_cache
 def is_from_defaults(source: Source):
     if isinstance(source, DefaultsSource):
         return True

torch/_dynamo/variables/builder.py

@@ -1807,7 +1807,28 @@ def wrap_literal(self, value):
                 # unspecializing int by default, but still
                 # specialize for the following conditions
                 if is_int_specialization_case(value, self.source):
-                    self.install_guards(GuardBuilder.CONSTANT_MATCH)
+                    recompile_hint = None
+                    if (
+                        self.source.guard_source().is_unspecialized_builtin_nn_module()
+                        or self.source.guard_source().is_unspecialized_nn_module()
+                    ):
+                        # This means that it is an integer from a NN module.
+                        # Dynamo considers nn module int attributes to be static
+                        # (a good heursitic). But a user might want to mark the
+                        # int attribute to be a symint, so track this integer
+                        # for recompilation later.
+                        recompile_hint = (
+                            "torch.compile considers integer attributes of the nn.Module to be static. "
+                            "If you are observing recompilation, you might want to make this integer dynamic "
+                            "using torch._dynamo.config.allow_unspec_int_on_nn_module = True, or convert this "
+                            "integer into a tensor."
+                        )
+
+                    self.install_guards(
+                        functools.partial(
+                            GuardBuilder.EQUALS_MATCH, recompile_hint=recompile_hint
+                        )
+                    )
                     return ConstantVariable.create(value=value, source=self.source)
 
             return self.wrap_symint(value)