Pass torch.compile mode/options to all backends (pytorch#99645)

Pull Request resolved: pytorch#99645
Approved by: https://github.com/anijain2305

Author: jansel

test/inductor/test_config.py

@@ -190,6 +190,40 @@ def test_invalid_backend(self):
             lambda: torch.compile(dummy_fn, backend="does_not_exist")(torch.randn(10)),
         )
 
+    def test_non_inductor_backend(self):
+        def assert_options(expected_mode=None, expected_options=None):
+            def backend(gm, _, *, mode=None, options=None):
+                nonlocal call_count
+                self.assertEqual(mode, expected_mode)
+                self.assertEqual(options, expected_options)
+                call_count += 1
+                return gm
+
+            return backend
+
+        inp = torch.randn(8)
+
+        def fn(x):
+            return x + 1
+
+        for mode, options in [
+            (None, None),
+            ("fast-mode", None),
+            (None, {"foo": "bar"}),
+        ]:
+            call_count = 0
+            torch.compile(
+                fn, backend=assert_options(mode, options), mode=mode, options=options
+            )(inp)
+            torch._dynamo.reset()
+            self.assertEqual(call_count, 1)
+
+        # TypeError: eager() got an unexpected keyword argument 'mode'
+        self.assertRaises(
+            torch._dynamo.exc.BackendCompilerFailed,
+            lambda: torch.compile(fn, backend="eager", mode="nope")(inp),
+        )
+
 
 if __name__ == "__main__":
     run_tests()

torch/__init__.py

@@ -1536,6 +1536,35 @@ def reset(self):
                 from torch._inductor.cudagraph_trees import reset_cudagraph_trees
                 reset_cudagraph_trees()
 
+class _TorchCompileWrapper:
+    def __init__(self, backend, mode, options, dynamic):
+        from torch._dynamo.backends.registry import lookup_backend
+
+        if isinstance(backend, str):
+            self.compiler_name = backend
+        elif hasattr(backend, "__name__"):
+            self.compiler_name = backend.__name__
+        else:
+            self.compiler_name = str(backend)
+        self.dynamic = dynamic
+        self.compiler_fn = lookup_backend(backend)
+        self.kwargs = {}
+        # only pass the args if they non-empty
+        if mode and mode != "default":
+            self.kwargs["mode"] = mode
+        if options:
+            self.kwargs["options"] = options
+
+    def __eq__(self, other):
+        return (isinstance(other, _TorchCompileWrapper) and
+                self.compiler_fn == other.compiler_fn and
+                self.kwargs == other.kwargs and
+                self.dynamic == other.dynamic)
+
+    def __call__(self, model_, inputs_):
+        return self.compiler_fn(model_, inputs_, **self.kwargs)
+
+
 def compile(model: Optional[Callable] = None, *,
             fullgraph: builtins.bool = False,
             dynamic: builtins.bool = False,
@@ -1600,6 +1629,8 @@ def fn(model: Callable):
         mode = "default"
     if backend == "inductor":
         backend = _TorchCompileInductorWrapper(mode, options, dynamic)
+    else:
+        backend = _TorchCompileWrapper(backend, mode, options, dynamic)
 
     return torch._dynamo.optimize(backend=backend, nopython=fullgraph, dynamic=dynamic, disable=disable)(model)
 

torch/_dynamo/repro/after_dynamo.py

@@ -111,7 +111,8 @@ def add_paths(exc):
         return compiled_gm
 
     debug_wrapper._torchdynamo_orig_callable = unconfigured_compiler_fn  # type: ignore[attr-defined]
-
+    if hasattr(unconfigured_compiler_fn, "compiler_name"):
+        debug_wrapper.__name__ = unconfigured_compiler_fn.compiler_name
     return debug_wrapper