Reuse the graph's fake mode when every input is a lifted constant

exir/pass_base.py

@@ -34,6 +34,7 @@
 from executorch.exir.error import ExportError, ExportErrorType
 from torch import fx
 from torch._dispatch.python import enable_python_dispatcher
+from torch._guards import detect_fake_mode
 from torch._subclasses import FakeTensorMode, UnsupportedFakeTensorException
 from torch._subclasses.fake_tensor import FakeTensor
 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
@@ -863,6 +864,22 @@ def call(self, graph_module: fx.GraphModule) -> PassResult:
                     fake_tensor_mode is None or fake_tensor_mode is i.fake_mode
                 ), "Multiple fake tensor mode detected."
                 fake_tensor_mode = i.fake_mode
+        if fake_tensor_mode is None:
+            # inputs() unwraps a constant-carrying FakeTensor to its real
+            # .constant tensor, so a graph whose placeholders are all lifted
+            # constants yields no FakeTensor above even though the graph itself
+            # is faked. Opening a fresh mode there leaves the retraced nodes in
+            # one mode and the untouched placeholders in another, and the mixed
+            # graph is rejected later by detect_fake_mode(). Recover the
+            # graph's own mode before falling back to a new one.
+            fake_tensor_mode = detect_fake_mode(
+                [
+                    node.meta["val"]
+                    for node in graph_module.graph.nodes
+                    if node.meta.get("val", None) is not None
+                ]
+            )
+
         if fake_tensor_mode is None:
             fake_tensor_mode = FakeTensorMode(allow_non_fake_inputs=True)
             dispatcher_mode = nullcontext()  # type: ignore[assignment]

exir/tests/test_pass_infra.py

@@ -24,7 +24,7 @@
 from executorch.exir.passes import ScalarToTensorPass
 from executorch.exir.passes.pass_registry import PassRegistry
 from executorch.exir.program import to_edge
-from torch._subclasses.fake_tensor import FakeTensor
+from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
 from torch.export import Dim, export, ExportedProgram
 from torch.export.graph_signature import InputKind, InputSpec, TensorArgument
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
@@ -193,6 +193,50 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
                 self.assertIn("val", node.meta)
 
 
+    def test_export_pass_reuses_graph_fake_mode(self) -> None:
+        """A graph whose placeholders are all constant-carrying fake tensors.
+
+        inputs() unwraps such a placeholder to its real .constant tensor, so
+        the fake-mode scan in _ExportPassBase.call finds no FakeTensor at all.
+        Opening a fresh mode there splits the graph in two: the retraced nodes
+        land in the new mode while the untouched placeholders keep the old one,
+        and detect_fake_mode() rejects the result downstream.
+        """
+        mode = FakeTensorMode(allow_non_fake_inputs=True)
+
+        def const_fake(value: float) -> FakeTensor:
+            real = torch.tensor(value)
+            with mode:
+                meta = torch.empty(real.shape, dtype=real.dtype, device="meta")
+            return FakeTensor(mode, meta, torch.device("cpu"), constant=real)
+
+        graph = torch.fx.Graph()
+        lhs = graph.placeholder("c_lhs")
+        rhs = graph.placeholder("c_rhs")
+        lhs.meta["val"] = const_fake(2.0)
+        rhs.meta["val"] = const_fake(3.0)
+        mul = graph.call_function(torch.ops.aten.mul.Tensor, (lhs, rhs))
+        with mode:
+            mul.meta["val"] = torch.ops.aten.mul.Tensor(
+                lhs.meta["val"], rhs.meta["val"]
+            )
+        out = graph.output((mul,))
+        out.meta["val"] = (mul.meta["val"],)
+        gm = torch.fx.GraphModule(torch.nn.Module(), graph)
+
+        result = ExportPass()(gm)
+        self.assertIsNotNone(result)
+
+        fake_modes = set()
+        for node in result.graph_module.graph.nodes:
+            val = node.meta.get("val")
+            for tensor in val if isinstance(val, (list, tuple)) else [val]:
+                if isinstance(tensor, FakeTensor):
+                    fake_modes.add(id(tensor.fake_mode))
+
+        self.assertEqual(len(fake_modes), 1)
+
+
 class TestProxyValueSymbolicCoercions(unittest.TestCase):
     @staticmethod
     def _symbolic_values() -> tuple[torch.SymInt, torch.SymFloat]: