Remove _deprecated_global_ns from cond

Summary:
Remove _deprecated_global_ns from cond following pytorch/pytorch#104105.

To do that, we need to change how graph_module generates python_code for "cond" target. Otherwise, it will generate target as "torch.ops.cond", which is invalid after the change.

Will import this PR to fix internal tests.

cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78

X-link: pytorch/pytorch#104380

Reviewed By: zou3519

Differential Revision: D47110919

Pulled By: ydwu4

fbshipit-source-id: eed2f7e0aa6bfc0d0a46f0064630abed872e3d75

Author: ydwu4

backends/test/test_backends_nested.py

@@ -115,7 +115,7 @@ def process(gm):
             processed_bytes = ""
             for node in gm.graph.nodes:
                 if node.op == "call_function":
-                    if node.target is torch.ops.cond:
+                    if node.target is torch.ops.higher_order.cond:
                         _, true_gm, _ = _get_submodule(gm, node, 1)
                         _, false_gm, _ = _get_submodule(gm, node, 2)
                         processed_bytes += f"{node.target.__name__}({process(true_gm)},{process(false_gm)});"
@@ -134,7 +134,7 @@ def process(gm):
 
 class CondOperatorSupport(OperatorSupportBase):
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
-        return node.op == "call_function" and node.target is torch.ops.cond
+        return node.op == "call_function" and node.target is torch.ops.higher_order.cond
 
 
 @final
@@ -168,7 +168,10 @@ def partition(
         for partition in partition_list:
             for node in partition.nodes:
                 delegation_tag = f"backend1_tag{partition.id}"
-                if node.op == "call_function" and node.target is torch.ops.cond:
+                if (
+                    node.op == "call_function"
+                    and node.target is torch.ops.higher_order.cond
+                ):
                     # Tag the arguments that take in the submodules to cond
                     # pyre-ignore
                     node.args[1].meta["delegation_tag"] = delegation_tag

docs/website/docs/ir_spec/control_flow.md

@@ -4,7 +4,7 @@ EXIR has a couple of special operators used to help specify control flow within
 some code, similar to jax's control flow operators. Currently these operators
 are only supported for inference.
 
-## torch.ops.cond
+## torch.ops.higher_order.cond
 
 The `cond` function represents an “if” statement in other programming languages.
 It can logically be seen as implemented as follows:
@@ -87,7 +87,7 @@ class GraphModule(torch.nn.Module):
         %gt : Sym(s0 > 4) = call_function[target=operator.gt](args = (%sym_size, 4), kwargs = {})
         %true_graph_0 = get_attr[target=true_graph_0]
         %false_graph_0 = get_attr[target=false_graph_0]
-        %cond : f32[s0, s1] = call_function[target=torch.ops.cond](args = (%gt, %true_graph_0, %false_graph_0, [%arg0]), kwargs = {})
+        %cond : f32[s0, s1] = call_function[target=torch.ops.higher_order.cond](args = (%gt, %true_graph_0, %false_graph_0, [%arg0]), kwargs = {})
         return [cond]
 
     # true_graph_0

exir/graph_module.py

@@ -299,7 +299,7 @@ def get_control_flow_submodules(
 ) -> List[Tuple[str, torch.fx.GraphModule, torch.fx.Node]]:
     """
     Returns a list of submodules used for control flow operations
-    (torch.ops.cond/map) that are in the given toplevel graph (does not look
+    (torch.ops.higher_order.cond/map) that are in the given toplevel graph (does not look
     into submodules). Specifically, the returned value is a list containing a
     tuple of (name of the submodule that's stored in the graph module, the
     submodule itself, and the fx node that uses this submodule).
@@ -309,7 +309,7 @@ def get_control_flow_submodules(
         if node.op != "call_function":
             continue
 
-        if node.target is torch.ops.cond:
+        if node.target is torch.ops.higher_order.cond:
             control_flow_submodules.append(_get_submodule(graph_module, node, 1))
             control_flow_submodules.append(_get_submodule(graph_module, node, 2))
         if node.target is torch.ops.map_impl:

exir/pass_base.py

@@ -53,7 +53,7 @@ def make_inline_interpreter(
 ) -> Type[torch.fx.Interpreter]:
     class InlineInterpreter(parent):
         def call_function(self, target, args, kwargs):
-            if target == torch.ops.cond:
+            if target == torch.ops.higher_order.cond:
                 pred, true, false, params = args
                 return InlineInterpreter(true).run(*params)
             elif target == torch.ops.map:

exir/passes/__init__.py

@@ -314,7 +314,7 @@ def get_submodule(node: torch.fx.Node) -> torch.fx.GraphModule:
                 continue
 
             target = node.target
-            if target == control_flow.cond or target == torch.ops.cond:
+            if target == control_flow.cond or target == torch.ops.higher_order.cond:
                 self.call(get_submodule(node.args[1]))
                 self.call(get_submodule(node.args[2]))
                 continue

exir/tests/control_flow_models.py

@@ -13,7 +13,9 @@ def true_branch(x):
         def false_branch(x):
             return x * x
 
-        return torch.ops.cond(inp.size(0) > 4, true_branch, false_branch, [inp])
+        return torch.ops.higher_order.cond(
+            inp.size(0) > 4, true_branch, false_branch, [inp]
+        )
 
     def get_random_inputs(self):
         return (torch.rand(5),)
@@ -30,7 +32,9 @@ def true_branch(x):
         def false_branch(x):
             return x * x * x
 
-        return torch.ops.cond(inp.size(0) > 4, true_branch, false_branch, [inp])
+        return torch.ops.higher_order.cond(
+            inp.size(0) > 4, true_branch, false_branch, [inp]
+        )
 
     def get_upper_bound_inputs(self):
         return (torch.rand(8),)
@@ -60,7 +64,9 @@ def true_branch(x):
         def false_branch(x):
             return x * 2
 
-        return torch.ops.cond(inp.size(0) > 4, true_branch, false_branch, [inp])
+        return torch.ops.higher_order.cond(
+            inp.size(0) > 4, true_branch, false_branch, [inp]
+        )
 
     def get_random_inputs(self):
         return (torch.eye(5) * 2,)

test/end2end/test_end2end.py

@@ -749,7 +749,7 @@ class DynamicModelE2ETest(unittest.TestCase):
         run_executor=False,
     )
 
-    # basic test for functorch torch.ops.cond
+    # basic test for functorch torch.ops.higher_order.cond
     test_ft_cond_basic = maketest(
         FTCondBasic,
         capture_config=exir.CaptureConfig(