Replace exir cond with torch cond

Summary: This diff removes exir.control_flow.cond and replace its existing usage with torch.cond.

Reviewed By: angelayi

Differential Revision: D47924374

fbshipit-source-id: d296479ee1f708cb423a27feb55e258632238902

Author: ydwu4

exir/control_flow.py

@@ -125,67 +125,6 @@ def _make_submodule(
     return gm
 
 
-def cond(
-    pred: bool,
-    true_fn: Callable[..., Tuple[torch.Tensor]],
-    false_fn: Callable[..., Tuple[torch.Tensor]],
-    inputs: pytree.PyTree,
-) -> Union[List[torch.Tensor], Value]:
-    """
-    A higher order function returning result based on passed predicate
-    value and conditionally execute one of true_fn and false_fn.
-
-    Detects whether a tracer is present in the context, and if so will
-    trace_through both true_fn and false_fn with local inputs provided
-    by tracing_context dictionary from the current tracer. When
-    returning, wraps two traced graphs into a cond() call and construct
-    a call_function node in the tracer's graph.
-
-    Checks and enforces that the returning value(s) from both
-    branches has the same Tensor type. For now enforces that both
-    branches have the same number of tensor inputs.
-    """
-    flattened_inputs, _ = pytree.tree_flatten(inputs)
-
-    if not all([isinstance(i, torch.Tensor) for i in flattened_inputs]):
-        raise ExportError(
-            ExportErrorType.INVALID_INPUT_TYPE,
-            f"control_flow.cond() expects all inputs values to be tensors, actual inputs: {inputs}",
-        )
-
-    with using_tracer(None):
-        outputs = true_fn(*inputs) if pred else false_fn(*inputs)
-
-    flattened_outputs, _ = pytree.tree_flatten(outputs)
-
-    if not all([isinstance(r, torch.Tensor) for r in flattened_outputs]):
-        raise ExportError(
-            ExportErrorType.INVALID_OUTPUT_TYPE,
-            f"control_flow.cond() only supports tensors as output, actual output: {outputs}",
-        )
-
-    tracer = DispatchTracer.get()
-
-    if tracer is None:
-        return outputs
-
-    # Once global tracer is present, we need to assume all tensors are
-    # PythonTensor wrapped with FunctionalTensorWrapper.
-
-    gm_true = _make_submodule(true_fn, example_returns=flattened_outputs)
-    gm_false = _make_submodule(false_fn, example_returns=flattened_outputs)
-    proxies = tuple([unwrap_proxy(i) for i in flattened_inputs])
-
-    proxy = tracer.create_proxy(
-        "call_function",
-        cond,
-        (unwrap_proxy(pred), gm_true, gm_false, proxies),
-        {},
-    )
-
-    return tree_return(outputs, proxy, update_with_proxy)
-
-
 def while_loop(
     cond_fn: Callable[..., torch.Tensor],
     body_fn: Callable[..., Tuple[torch.Tensor]],

exir/passes/__init__.py

@@ -236,7 +236,7 @@ def callWithLoggerEnabled(self, graph_module: torch.fx.GraphModule) -> None:
 # pyre-ignore
 to_out_var_skiplist: Set[Callable[[Any], Any]] = {
     _operator.getitem,
-    control_flow.cond,
+    torch.ops.higher_order.cond,
     control_flow.while_loop,
     # memory.alloc will be added after the to_out_variant pass so usually
     # we won't see it in the input graph to the to_out_variant pass, unless
@@ -321,7 +321,7 @@ def get_submodule(node: torch.fx.Node) -> torch.fx.GraphModule:
                 continue
 
             target = node.target
-            if target == control_flow.cond or target == torch.ops.higher_order.cond:
+            if target == torch.ops.higher_order.cond:
                 self.call(get_submodule(node.args[1]))
                 self.call(get_submodule(node.args[2]))
                 continue

test/end2end/test_end2end.py

@@ -56,6 +56,7 @@
 from executorch.exir.tests.dynamic_shape_models import BatchNormModel
 
 from executorch.exir.tests.transformer import Transformer
+from functorch.experimental.control_flow import cond
 
 kernel_mode = None  # either aten mode or lean mode
 try:
@@ -245,15 +246,13 @@ def addloop(x, n):
                 out = out + x
             return out
 
-        @control_flow.tracing_context(inputs=(torch.randn(1), torch.randn(10)))
         def true_branch(c, x):
             return addloop(x, 3)
 
-        @control_flow.tracing_context(inputs=(torch.randn(1), torch.randn(10)))
         def false_branch(c, x):
             return addloop(x, 4)
 
-        y = control_flow.cond(c, true_branch, false_branch, (c, x))
+        y = cond(c, true_branch, false_branch, (c, x))
         return y * y
 
     def get_random_inputs(self):
@@ -273,18 +272,14 @@ def addloop(x, n):
                 out = out + x
             return out
 
-        @control_flow.tracing_context(inputs=(torch.randn(1), torch.randn(10)))
         def true_branch(c, x):
             return addloop(x, 3)
 
-        @control_flow.tracing_context(inputs=(torch.randn(1), torch.randn(10)))
         def false_branch(c, x):
             return addloop(x, 4)
 
-        # pyre-fixme[6]: Incompatible parameter type
-        y = control_flow.cond(c, true_branch, false_branch, (c, x))
+        y = cond(c, true_branch, false_branch, (c, x))
 
-        # pyre-fixme[58]: Unsupported operand type for binary operator '*'
         return y * y
 
     def get_random_inputs(self):
@@ -319,7 +314,7 @@ def true_branch(cnt):
             def false_branch(cnt):
                 return torch.zeros([1], dtype=torch.long)
 
-            accum = accum + control_flow.cond(
+            accum = accum + cond(
                 torch.BoolTensor([True]), true_branch, false_branch, (cnt,)
             )
             # 'cnt - 1' does not work yet since the runtime does not expect
@@ -372,9 +367,9 @@ def loop_body(accum, cnt):
         def false_branch(accum, cnt):
             return accum, cnt
 
-        return control_flow.cond(
-            torch.BoolTensor([True]), true_branch, false_branch, (accum, cnt)
-        )[0]
+        return cond(torch.BoolTensor([True]), true_branch, false_branch, (accum, cnt))[
+            0
+        ]
 
     def get_random_inputs(self):
         return (torch.zeros([1]).to(dtype=torch.long), torch.randint(10, 100, [1]))