[dynamo][index] Speedup index method for constant data structures (#173612)

anijain2305 · pytorchmergebot · commit 94f18c09b873 · 2026-01-28T17:37:15.000Z
Saves 0.5 seconds on 15.3 seconds of compile time of a model Pull Request resolved: #173612 Approved by: https://github.com/Lucaskabela ghstack dependencies: #173582
diff --git a/torch/_dynamo/variables/lists.py b/torch/_dynamo/variables/lists.py
@@ -44,7 +44,7 @@ class that handles its unique behaviors while integrating with Dynamo's
     range_iterator,
     set_example_value,
 )
-from .base import ValueMutationNew, VariableTracker
+from .base import AsPythonConstantNotImplementedError, ValueMutationNew, VariableTracker
 from .constant import ConstantVariable
 from .functions import UserFunctionVariable
 from .iter import IteratorVariable
@@ -238,11 +238,25 @@ def call_method(
                     f"{len(args)} args and {len(kwargs)} kwargs",
                 )
 
-            return tx.inline_user_function_return(
-                VariableTracker.build(tx, polyfills.index),
-                [self] + list(args),
-                kwargs,
-            )
+            try:
+                # Speedup trace times for constant data structures
+                items = [item.as_python_constant() for item in self.items]
+                const_args = [arg.as_python_constant() for arg in args]
+                const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
+                try:
+                    return ConstantVariable.create(
+                        items.index(*const_args, **const_kwargs)
+                    )
+                except ValueError:
+                    raise_observed_exception(
+                        ValueError, tx, args=[ConstantVariable.create("tuple.index()")]
+                    )
+            except AsPythonConstantNotImplementedError:
+                return tx.inline_user_function_return(
+                    VariableTracker.build(tx, polyfills.index),
+                    [self] + list(args),
+                    kwargs,
+                )
         elif name == "count":
             if len(args) != 1:
                 raise_args_mismatch(
