[POC] Don't decompose matmul

ezyang · ezyang · commit 71d5202643f2 · 2022-12-19T09:42:24.000+08:00
Signed-off-by: Edward Z. Yang &lt;ezyang@fb.com&gt;

[ghstack-poisoned]
diff --git a/a.py b/a.py
@@ -0,0 +1,8 @@
+import torch
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch._dispatch.python import enable_python_dispatcher, patch_py_impls
+
+with enable_python_dispatcher(), patch_py_impls({
+    torch.ops.aten.matmul.default: {torch._C.DispatchKey.AutogradCPU: torch._C.DispatchKey.Autograd}
+}):
+    print(make_fx(torch.matmul)(torch.randn(2, 3), torch.randn(3, 4)))
diff --git a/torch/_dispatch/python.py b/torch/_dispatch/python.py
@@ -4,6 +4,7 @@
 import torch
 import torch.utils._pytree as pytree
 import itertools
+from typing import Callable, Dict
 
 __all__ = ['enable_python_dispatcher', 'no_python_dispatcher']
 
@@ -140,3 +141,26 @@ def enable_crossref_functionalize():
     finally:
         for op in all_known_overloads():
             op._uncache_dispatch(torch._C.DispatchKey.Functionalize)
+
+@contextmanager
+def patch_py_impls(all_patches: Dict[torch._ops.OpOverload, Dict[torch._C.DispatchKey, Callable]]):
+    """
+    Temporarily patch the dispatcher registrations in the Python Dispatcher,
+    undoing them when you exit the context manager.  This is useful for
+    temporarily adding pre-autograd decompositions, among other things.
+    """
+    saved_tables = {}
+    for op, patches in all_patches.items():
+        # TODO: Make this public API on OpOverload instead
+        # of groveling the attribute directly
+        saved_tables[op] = op.py_kernels.copy()
+        for k, fn in patches.items():
+            op.py_impl(k)(fn)
+    try:
+        yield
+    finally:
+        for op in all_patches:
+            # TODO: Make this OpOverload API
+            op.py_kernels.clear()
+            op.py_kernels.update(saved_tables[op])
+            op._dispatch_cache.clear()
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -295,6 +295,7 @@ def namespace(self):
         return self._schema.name.split("::")[0]
 
     def decompose(self, *args, **kwargs):
+        return NotImplemented
         dk = torch._C.DispatchKey.CompositeImplicitAutograd
         if dk in self.py_kernels:
             # NB: This branch is not too necessary anymore, because we can
