Supported kwargs in ThunderCompiler (#2759)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Yan Wang <[email protected]>

Author: 3 people

thunder/dynamo/compiler.py

@@ -129,7 +129,7 @@ def __init__(self, **thunder_options):
         )
         self.thunder_options = thunder_options
 
-    def __call__(self, gm: torch.fx.GraphModule, sample_args: list[torch.SymInt, torch.Tensor]):
+    def __call__(self, gm: torch.fx.GraphModule, sample_args: list[torch.SymInt, torch.Tensor], **compile_options):
         from thunder import jit
 
         remove_empty_autocast(gm)
@@ -148,6 +148,7 @@ def __call__(self, gm: torch.fx.GraphModule, sample_args: list[torch.SymInt, tor
             gm,
             partial(jit, **thunder_options),
             thunder_options,
+            **compile_options,
         )
         self.subgraph_infos.append(subgraph_info)
         return split_module

thunder/dynamo/splitter.py

@@ -34,6 +34,7 @@ def _splitter(
     gm: torch.fx.GraphModule,
     thunder_jit: Callable,
     thunder_options: dict[str, Any] | None = None,
+    **compile_options,
 ) -> tuple[torch.fx.GraphModule, SubgraphInfo]:
     """
     This method will split graph into multiple graph modules based on thunder supported operations.
@@ -225,7 +226,7 @@ def is_thunder_supported_partition(node: torch.fx.Node) -> bool:
             fake_mode = torch._guards.detect_fake_mode()
             # Delay Inductor compilation until invocation with real tensors,
             # because we do not know the strides of tensors that Thunder-compiled submodules return.
-            jit_fn = LazyInductorModule(graph_module, fake_mode)
+            jit_fn = LazyInductorModule(graph_module, fake_mode, **compile_options)
 
             # Update the node name from "submod_*" to "inductor_*" for more user-friendly names
             update_node_and_submodule(split_gm, node, node.name.replace("submod", "inductor"), jit_fn)

thunder/dynamo/utils.py

@@ -20,6 +20,8 @@
 from torch._guards import tracing, TracingContext
 from torch._subclasses.fake_tensor import DynamicOutputShapeException
 
+from torch._inductor import list_mode_options
+
 if torch.distributed.is_available():
     from torch.distributed.tensor import DTensor
 else:
@@ -160,11 +162,12 @@ def is_thunder_supported_partition(self, node: torch.fx.Node) -> bool:
 
 
 class LazyInductorModule(torch.nn.Module):
-    def __init__(self, graph_module, fake_mode):
+    def __init__(self, graph_module, fake_mode, **compile_options):
         super().__init__()
         self.graph_module = graph_module
         self.compiled_fn = None
         self.fake_mode = fake_mode
+        self.compile_options = compile_options
 
         # For ease of debugging, we add graph attribute so GraphModule.print_readable will print it
         self.graph = graph_module.graph
@@ -200,7 +203,14 @@ def forward(self, *args):
                 with tracing(TracingContext(fake_mode=self.fake_mode)):
                     try:
                         original_graph = copy.deepcopy(self.graph_module.graph)
-                        self.compiled_fn = torch._inductor.compile(self.graph_module, args)
+                        # Extract and merge options from compile_options
+                        options = self.compile_options.get("options", {}).copy()
+                        mode = self.compile_options.get("mode")
+                        if mode:
+                            mode_options = list_mode_options().get(mode, {})
+                            options.update(mode_options)
+
+                        self.compiled_fn = torch._inductor.compile(self.graph_module, args, options=options)
                     except DynamicOutputShapeException as e:
                         # This exception is meant to be handled by Dynamo, which is responsible for graph break
                         # TODO: Use torch.compile for fallback. Ensure its correctness.

thunder/tests/test_dynamo.py

@@ -180,6 +180,55 @@ def func(x):
     assert mock_inductor.call_count == 2
 
 
+@instantiate(
+    dtypes=NOTHING,
+    executors=[DynamoThunderExecutor],
+    decorators=(
+        pytest.mark.parametrize("mode", ("default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs")),
+        pytest.mark.skipif(
+            condition=IS_WINDOWS,
+            reason="torch.compile Windows support is still WIP - https://github.com/pytorch/pytorch/issues/122094",
+        ),
+    ),
+)
+def test_kwargs_forwarding_to_inductor(executor, device, dtype, mode):
+    """Test that torch.compile kwargs (like mode) are forwarded to inductor for unsupported regions."""
+    from torch._inductor import list_mode_options
+
+    x = torch.randn(2, 2, device=device, dtype=dtype)
+
+    def func(x):
+        # torch.sinc has automatic fallback registered,
+        # so that operation will be given to inductor.
+        return x.sinc()
+
+    cfunc = thunderfx(func, mode=mode)
+
+    # Mock torch._inductor.compile to verify arguments are passed correctly
+    original_compile = torch._inductor.compile
+
+    with patch("torch._inductor.compile", side_effect=original_compile) as mock_compile:
+        cfunc(x)
+
+        # Verify the mock was called (inductor fallback occurred)
+        assert mock_compile.called
+
+        # Get the kwargs from the call
+        _, compile_options = mock_compile.call_args
+
+        # Check if mode was expanded into options
+        expected_options = list_mode_options().get(mode, {})
+
+        # At least empty dict options should be passed
+        assert "options" in compile_options
+
+        if expected_options:
+            # If the mode has options, verify they were passed
+            options = compile_options["options"]
+            for k, v in expected_options.items():
+                assert options.get(k) == v, f"Expected option {k}={v} for mode {mode}, but got {options.get(k)}"
+
+
 @instantiate(
     dtypes=NOTHING,
     executors=[DynamoThunderExecutor],