diff --git a/exir/pass_base.py b/exir/pass_base.py index 6071aae2be8..e28448b73b3 100644 --- a/exir/pass_base.py +++ b/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] diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 16ed5af4180..57a87e406d6 100644 --- a/exir/tests/test_pass_infra.py +++ b/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]: