Revert "Put Python Dispatcher cache in dict, clear it on new registrations. (#88329)"

pytorchmergebot · pytorchmergebot · commit 410ce96a23a3 · 2022-11-03T21:57:19.000Z
This reverts commit 86c7cd2. Reverted #88329 on behalf of https://github.com/clee2000 due to test_decomp takes an extra 2 hours in some jobs, windows takes so long it times out
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -243,8 +243,6 @@ def __init__(self, overloadpacket, op, op_dk, schema, tags):
         op.__module__ = overloadpacket.__module__
         self.__qualname__ = self._name
         self.__annotations__ = {}
-        # NB: This name is hard-coded in torch/csrc/autograd/python_variable.cpp
-        self._dispatch_cache = {}
 
     # it's a no-op since OpOverload object is immutable and must be unique for a given op overload.
     def __deepcopy__(self, memo=None):
@@ -291,7 +289,6 @@ def inner(fn):
                 assert mode not in self.python_key_mode_table
                 # TODO(voz): Should we replace setting torch._C.DispatchKey.Python entirely with setting mode keys?
                 self.python_key_mode_table[mode] = fn
-                self._dispatch_cache.clear()
                 return fn
 
             assert isinstance(dispatch_key_or_mode, torch._C.DispatchKey)
@@ -304,19 +301,23 @@ def inner(fn):
                     f"Trying to override a python impl for {dispatch_key_or_mode} on operator {self._name}"
                 )
             self.py_kernels[dispatch_key_or_mode] = fn
-            self._dispatch_cache.clear()
             return fn
 
         return inner
 
     # This implements the pre-computation logic for the Python dispatcher.
-    def _get_dispatch(self, key):
-        # This is only called upon a cache miss
-        assert key not in self._dispatch_cache
+    def __getattr__(self, attr):
+        if len(attr) == 0 or not attr[0].isupper():
+            raise AttributeError()
+
+        try:
+            key = torch._C._dispatch_key_parse(attr)
+        except Exception as e:
+            raise AttributeError()
 
         if key == torch._C.DispatchKey.Python:
             if not self.python_key_mode_table:
-                self._dispatch_cache[key] = key
+                setattr(self, attr, key)
                 return key
 
             def handler(*args, **kwargs):
@@ -335,12 +336,12 @@ def handler(*args, **kwargs):
                 # TODO(voz): The idea behind this is that we do not yet support dispatch by key + mode, only key.
                 return self.python_key_mode_table[curr_mode](*args, **kwargs)
 
-            self._dispatch_cache[key] = handler
+            setattr(self, attr, handler)
             return handler
 
         key = resolve_key(self, key)
         r = self.py_kernels.get(key, key)
-        self._dispatch_cache[key] = r
+        setattr(self, attr, r)
         return r
 
     def name(self):
diff --git a/torch/csrc/autograd/python_variable.cpp b/torch/csrc/autograd/python_variable.cpp
@@ -2321,20 +2321,11 @@ void ConcretePyInterpreterVTable::python_dispatcher(
     torch::jit::Stack* stack) const {
   py::gil_scoped_acquire g;
   py::handle torch_api_function_overload = getTorchApiFunction(op);
-  // TODO: if necessary, can optimize to cache the cache lookup
-  // TODO: if necessary, can optimize OpOverload to have slots
-  auto cache = py::dict(torch_api_function_overload.attr("_dispatch_cache"));
-  if (cache.ptr() == nullptr) {
-    throw python_error();
-  }
 
   c10::DispatchKey k = ks.highestPriorityTypeId();
-  // TODO: allow this to be non-owning
-  auto handler = py::reinterpret_borrow<py::object>(
-      PyDict_GetItem(cache.ptr(), py::cast(k).ptr()));
+  auto handler = torch_api_function_overload.attr(toString(k));
   if (handler.ptr() == nullptr) {
-    // Slow path
-    handler = torch_api_function_overload.attr("_get_dispatch")(k);
+    throw python_error();
   }
   if (py::isinstance<c10::DispatchKey>(handler)) {
     // NB: not redispatch, as that will permanently remove the python
