Fixes in to_executorch for while

Summary: Fixes while loop support in executorch AOT up until emitter (#8769 (comment))

Differential Revision: D77452951

Author: angelayi

exir/passes/__init__.py

@@ -340,7 +340,7 @@ def get_submodule(node: torch.fx.Node) -> torch.fx.GraphModule:
             if target == torch.ops.higher_order.map_impl:
                 self.call(get_submodule(node.args[0]))
                 continue
-            elif target == control_flow.while_loop:
+            elif target == torch.ops.higher_order.while_loop:
                 self.call(get_submodule(node.args[0]))
                 self.call(get_submodule(node.args[1]))
                 continue

exir/passes/spec_prop_pass.py

@@ -19,7 +19,9 @@
 
 # pyre-ignore
 def make_spec(x):
-    if isinstance(x, torch.Tensor):
+    if isinstance(x, ProxyValue):
+        return make_spec(x.node.meta["val"])
+    elif isinstance(x, torch.Tensor):
         return TensorSpec.from_tensor(x)
     elif isinstance(x, (int, bool, float)):
         return x
@@ -109,6 +111,19 @@ def call_cond(self, pred, true_fn, false_fn, inputs, meta):
         meta["spec"] = pytree.tree_map(make_spec, true_out_node.meta["val"])
         return super().call_cond(pred, true_fn, false_fn, inputs, meta)
 
+    def call_while(
+        self,
+        cond_fn: torch.fx.GraphModule,
+        body_fn: torch.fx.GraphModule,
+        carried_inputs: List[ProxyValue],
+        additional_inputs: List[ProxyValue],
+        meta: NodeMetadata,
+    ):
+        meta["spec"] = pytree.tree_map(make_spec, carried_inputs)
+        return super().call_while(
+            cond_fn, body_fn, carried_inputs, additional_inputs, meta
+        )
+
     def call_map(
         self,
         f: torch.fx.GraphModule,

exir/program/test/test_program.py

@@ -313,7 +313,8 @@ def body_fn(it, x):
         # Instantiate and export
         inp = (torch.tensor(3), torch.randn(2, 2))
         exported = export(M(), inp)
-        to_edge(exported)
+        ep = to_edge(exported)
+        # ep.to_executorch()
 
     def test_constraint_present_after_dce(self):
         import executorch.exir as exir