diff --git a/exir/memory_planning.py b/exir/memory_planning.py index ee04ab9ddac..501265f9364 100644 --- a/exir/memory_planning.py +++ b/exir/memory_planning.py @@ -271,6 +271,14 @@ def verify_graph_input_output(self) -> None: for nd in graph_module.graph.nodes: if nd.op in check_list: if not (specs := get_node_tensor_specs(nd)): + # A node that carries no tensor specs has nothing to + # allocate. An output node can legitimately be in that + # position, e.g. when a delegated subgraph's only results + # are symints, and skipping it silently leaves + # graph_output_allocated unset so the assert below fires on + # a graph that is in fact fully planned. + if nd.op == "output": + graph_output_allocated = self.alloc_graph_output continue if _is_mutable_buffer(nd, self.graph_signature): continue diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index ea9a2362da5..f50b9d1247f 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -2115,3 +2115,27 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: if node.op == "call_function" ) self.assertFalse(has_inplace) + + +class TestTensorlessGraphOutput(unittest.TestCase): + def test_verifier_accepts_output_carrying_no_tensors(self) -> None: + """An output node is allowed to carry no tensors at all. + + A delegated subgraph whose only results are symints is shaped like + this, and there is nothing for the planner to allocate. The verifier + used to skip such a node without recording that it had seen an output, + then assert `graph_output_allocated not set` on a graph that was in + fact fully planned. + """ + graph = Graph() + sym = graph.placeholder("sym") + graph.output((sym,)) + graph_module = GraphModule({}, graph) + + verifier = Verifier( + graph_module, + alloc_graph_input=True, + alloc_graph_output=True, + alloc_mutable_buffers=True, + ) + verifier.verify_graph_input_output()