Don't specialize when indexing by SymInt (#99123)

Fixes #99091

Signed-off-by: Edward Z. Yang <ezyang@meta.com>

Pull Request resolved: #99123
Approved by: https://github.com/msaroufim

test/dynamo/test_subgraphs.py

@@ -379,6 +379,19 @@ def fn(a, b):
         # just one graph now rather than 10
         self.assertEqual(cnt_dynamic.frame_count, 1)
 
+    @patch("torch._dynamo.config.dynamic_shapes", True)
+    @patch("torch._dynamo.config.assume_static_by_default", False)
+    def test_dynamic_getitem(self):
+        def fn(a, b):
+            return a[b.size(0) - 1]
+
+        cnt = torch._dynamo.testing.CompileCounter()
+        opt_fn = torch._dynamo.optimize(cnt)(fn)
+        for i in range(3, 12):
+            opt_fn(torch.randn(i), torch.randn(i))
+        # just one graph
+        self.assertEqual(cnt.frame_count, 1)
+
     def test_dynamic_kwarg(self):
         def fn(a, b):
             return a - b * 10

torch/_dynamo/variables/builtin.py

@@ -479,6 +479,16 @@ def call_function(
                     # Work around weird bug in hf_T5
                     fn, args = operator.add, [args[1], args[0]]
 
+                if self.fn is operator.getitem and isinstance(args[1], SymNodeVariable):
+                    # Standard indexing will force specialization due to
+                    # __index__.  Rewrite as a regular torch op which will
+                    # trace fine
+                    fn, args = torch.select, [
+                        args[0],
+                        variables.ConstantVariable(0),
+                        args[1],
+                    ]
+
                 proxy = tx.output.create_proxy(
                     "call_function",
                     fn,