Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions exir/memory_planning.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 24 additions & 0 deletions exir/tests/test_memory_planning.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading