[SWDEV-554558] [dynamo] Use BINARY_SUBSCR for pre-graph bytecode for regular dict ac… (#2669)

…cesses (pytorch#155727)

vLLM profiler sets with_stack=True that shows the dict_getitem on the
profiler, both inflating the numbers and confusing compile users. This
PR keeps BINARY_SUBSCR for regular dicts, while using `dict.__getitem__`
only for dict subclasses.

Using binary_subscr is little bit faster, but not enough to make any
major latency improvements.

Pull Request resolved: pytorch#155727
Approved by: https://github.com/zou3519, https://github.com/StrongerXi,
https://github.com/jansel

(cherry picked from commit a9d5157)

Fixes #ISSUE_NUMBER

Co-authored-by: Animesh Jain <[email protected]>

Author: jataylo

test/dynamo/test_dicts.py

@@ -14,7 +14,6 @@
 from typing import Any, Optional, Tuple
 
 import torch
-import torch._dynamo.config
 import torch._dynamo.test_case
 import torch._dynamo.testing
 import torch._functorch.config

torch/_dynamo/guards.py

@@ -99,6 +99,7 @@
     ConstDictKeySource,
     DefaultsSource,
     DictGetItemSource,
+    DictSubclassGetItemSource,
     FlattenScriptObjectSource,
     FloatTensorSource,
     FSDPNNModuleSource,
@@ -1079,7 +1080,7 @@ def get_guard_manager_from_source(self, source):
                     example_value=example_value,
                     guard_manager_enum=guard_manager_enum,
                 )
-        elif istype(source, DictGetItemSource):
+        elif istype(source, (DictGetItemSource, DictSubclassGetItemSource)):
             assert base_guard_manager  # to make mypy happy
             assert isinstance(base_example_value, (dict, collections.OrderedDict))
             if isinstance(base_guard_manager, DictGuardManager):

torch/_dynamo/source.py

@@ -601,7 +601,44 @@ def __post_init__(self):
     def guard_source(self):
         return self.base.guard_source()
 
-    def reconstruct(self, codegen):
+    def reconstruct(self, codegen: "PyCodegen"):
+        # Load dict
+        codegen(self.base)
+
+        # Load key
+        if isinstance(self.index, Source):
+            codegen(self.index)
+        else:
+            codegen.append_output(codegen.create_load_const(self.index))
+        codegen.append_output(create_instruction("BINARY_SUBSCR"))
+
+    def name(self):
+        if isinstance(self.index, ConstDictKeySource):
+            return f"{self.base.name()}[{self.index.name()}]"
+        else:
+            return f"{self.base.name()}[{self.index!r}]"
+
+
+# Same as DictGetItemSource but used for dict.__getitem__ calls to ensure that
+# torch.compile does not run the overridden __getitem__ method
+@dataclasses.dataclass(frozen=True)
+class DictSubclassGetItemSource(ChainedSource):
+    # Key to access in the dictionary. It can be one of the the following types
+    # 1) ConstDictKeySource
+    # 2) constant - like string, integer
+    index: Any
+
+    def __post_init__(self):
+        from .variables import ConstantVariable
+
+        assert isinstance(
+            self.index, ConstDictKeySource
+        ) or ConstantVariable.is_literal(self.index)
+
+    def guard_source(self):
+        return self.base.guard_source()
+
+    def reconstruct(self, codegen: "PyCodegen"):
         # reconstruct dict.__getitem__(dct, key)
 
         # Load dict.__getitem__

torch/_dynamo/variables/builder.py

@@ -92,6 +92,7 @@
     ConstDictKeySource,
     ConvertIntSource,
     DictGetItemSource,
+    DictSubclassGetItemSource,
     FloatTensorSource,
     GetItemSource,
     GradSource,
@@ -1274,8 +1275,8 @@ def build_key_value(i, k, v):
                 source_key = ConstDictKeySource(self.get_source(), i)
                 key = LazyVariableTracker.create(k, source_key)
 
-                source_value = DictGetItemSource(self.get_source(), source_key)
-                value = LazyVariableTracker.create(v, source_value)
+                source_value = DictSubclassGetItemSource(base, source_key)
+                res_value = LazyVariableTracker.create(v, source_value)
 
                 return key, value