fix: ensure output/out-var node specs before memory planning

- Add ensure_graph_node_specs() to set meta['spec'] from meta['val'] when
  missing (e.g. after delegation), so memory planning and verifier succeed.
- Call it at start of MemoryPlanningPass.run() for all graph modules.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: NefAI

exir/memory_planning.py

@@ -41,7 +41,7 @@
     InputKind,
 )
 from torch.fx import Node
-from torch.utils._pytree import tree_flatten
+from torch.utils._pytree import tree_flatten, tree_map
 
 REGISTERED_ALGOS: Dict[str, Callable[..., List[int]]] = {}
 
@@ -755,6 +755,26 @@ def get_node_tensor_specs(
         ]
 
 
+def ensure_graph_node_specs(graph_module: torch.fx.GraphModule) -> None:
+    """
+    Set meta["spec"] from meta["val"] for nodes that are missing spec (e.g. output
+    or out-var nodes in delegated graphs that were built after SpecPropPass).
+    """
+    for node in graph_module.graph.nodes:
+        if "spec" in node.meta:
+            continue
+        if "val" not in node.meta:
+            continue
+        val = node.meta["val"]
+
+        def to_spec(x: Any) -> Any:
+            if isinstance(x, torch.Tensor):
+                return TensorSpec.from_tensor(x)
+            return x
+
+        node.meta["spec"] = tree_map(to_spec, val)
+
+
 # Little bit hacky to check if the graph contains
 # XNNPACK delegate
 # Why?

exir/passes/memory_planning_pass.py

@@ -19,6 +19,7 @@
     _is_out_var_node,
     apply_algo,
     collect_specs_from_nodes,
+    ensure_graph_node_specs,
     filter_nodes,
     get_node_tensor_specs,
     MemoryPlanningAlgorithmSuite,
@@ -233,6 +234,9 @@ def run(
         A pass for memory planning. The actual algorithm used will be picked by
         memory_planning_algo
         """
+        for subgm in graph_module.modules():
+            if isinstance(subgm, torch.fx.GraphModule):
+                ensure_graph_node_specs(subgm)
         self._set_alloc_node_spec(graph_module)
         # TODO(shunting) if people have concern of adding a field to GraphModule
         # directly, we should define a GraphModule subclass that we can add our