Support while in pass base

Differential Revision: D77382557

Pull Request resolved: pytorch#12004

Author: angelayi

exir/pass_base.py

@@ -340,6 +340,11 @@ def call_function(
             elif target == torch.ops.higher_order.cond:
                 pred, true_fn, false_fn, inputs = args
                 return self.callback.call_cond(pred, true_fn, false_fn, inputs, meta)
+            elif target == torch.ops.higher_order.while_loop:
+                cond, body, carried_inputs, additional_inputs = args
+                return self.callback.call_while(
+                    cond, body, carried_inputs, additional_inputs, meta
+                )
             elif target == torch.ops.higher_order.map_impl:
                 f, mapped_args, operands = args  # type: ignore[assignment]
                 return self.callback.call_map(f, mapped_args, operands, meta)
@@ -497,6 +502,31 @@ def call_cond(
             meta,
         )
 
+    def call_while(
+        self,
+        cond_fn: torch.fx.GraphModule,
+        body_fn: torch.fx.GraphModule,
+        carried_inputs: List[Argument],
+        additional_inputs: List[Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        cond_fn = self.call_submodule(cond_fn, (*carried_inputs, *additional_inputs))
+        body_fn = self.call_submodule(body_fn, (*carried_inputs, *additional_inputs))
+        assert cond_fn is not None
+        assert body_fn is not None
+        return self._fx(
+            "call_function",
+            torch.ops.higher_order.while_loop,
+            (
+                cond_fn.graph_module,
+                body_fn.graph_module,
+                carried_inputs,
+                additional_inputs,
+            ),
+            {},
+            meta,
+        )
+
     def call_map(
         self,
         f: torch.fx.GraphModule,

exir/program/test/test_program.py

@@ -294,6 +294,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         for node in ep.graph.nodes:
             self.assertNotEqual(node.op, "get_attr")
 
+    def test_while(self):
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(2, 2)
+                self.dec = torch.nn.Buffer(torch.tensor(1))
+
+            def forward(self, iter, x):
+                def cond_fn(it, x):
+                    return it - self.dec > 0
+
+                def body_fn(it, x):
+                    return it - 1, self.linear(x)
+
+                return torch._higher_order_ops.while_loop(cond_fn, body_fn, (iter, x))
+
+        # Instantiate and export
+        inp = (torch.tensor(3), torch.randn(2, 2))
+        exported = export(M(), inp)
+        to_edge(exported)
+
     def test_constraint_present_after_dce(self):
         import executorch.exir as exir