[while_loop] disable closure capturing and manually set the inputs. (#122244)

ydwu4 · pytorchmergebot · commit 61ff41f0ca86 · 2024-03-20T20:14:35.000Z
For while_loop operator, it's important to keep the output ordering consistent with input ordering. Previously, we're using set_graph_inputs="automatic", which doesn't respect such ordering. This PR changes it to "manual" and respects the original user inputs' ordering. We disable closures for body and cond fn as they require some additional designs. This PR is just to prevent the bleeding. Repro: ```python import torch from torch._higher_order_ops.while_loop import while_loop from torch._functorch.aot_autograd import aot_export_module class Nested(torch.nn.Module): def forward(self, ci, cj, a, b): def cond_fn(i1, j1, x1, y1): return i1 > 0 def body_fn(i1, j1, x1, y1): def cond_fn_nested(i2, j2, x2, y2): return j2 > 0 def body_fn_nested(i2, j2, x2, y2): return i2.clone(), j2 - 1, x2 + 3.14, y2 - 2.71 i1, j1, x1, y1 = while_loop( cond_fn_nested, body_fn_nested, [i1, j1, x1, y1] ) return i1 - 1, j1.clone(), x1 * 2, y1 / 2 return while_loop(cond_fn, body_fn, (ci, cj, a, b)) nested = Nested() torch.compile(nested, backend="eager", fullgraph=True)(torch.tensor(2), torch.tensor(2), torch.randn(2, 2), torch.randn(2, 2)) ``` Test plan: add new test. Pull Request resolved: #122244 Approved by: https://github.com/aakhundov
diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
@@ -94,15 +94,36 @@ def outer_body_fn(out_iter, it, y):
 
         return while_loop(outer_cond_fn, outer_body_fn, (out_iter, it, y))
 
+    class Nested(torch.nn.Module):
+        def forward(self, ci, cj, a, b):
+
+            def cond_fn(i1, j1, x1, y1):
+                return i1 > 0
+
+            def body_fn(i1, j1, x1, y1):
+
+                def cond_fn_nested(i2, j2, x2, y2):
+                    return j2 > 0
+
+                def body_fn_nested(i2, j2, x2, y2):
+                    return i2.clone(), j2 - 1, x2 + 3.14, y2 - 2.71
+
+                i1, j1, x1, y1 = while_loop(
+                    cond_fn_nested, body_fn_nested, [i1, j1, x1, y1]
+                )
+                return i1 - 1, j1.clone(), x1 * 2, y1 / 2
+            return while_loop(cond_fn, body_fn, (ci, cj, a, b))
+
+    nested2 = Nested()
 
     x = torch.zeros(1)
     y = torch.zeros(1)
     z = torch.zeros(1)
     return {"simple": (simple, (x,)),
             "nested": (nested, (x, y, z)),
+            "nested2": (nested2, (torch.tensor(2), torch.tensor(2), torch.ones(2, 2), torch.ones(2, 2))),
             "simple_with_mutation": (simple_with_mutation, (x,))}
 
-
 def collect_meta_for_filtered_nodes(gm: torch.fx.GraphModule, node_names, meta_field_name):
     ret = []
     for mod in gm.modules():
@@ -376,7 +397,7 @@ def forward(self, arg0_1, arg1_1, arg2_1):
 def forward(self, arg0_1, arg1_1, arg2_1):
     while_loop_cond_graph_0 = self.while_loop_cond_graph_0
     while_loop_body_graph_0 = self.while_loop_body_graph_0
-    while_loop = torch.ops.higher_order.while_loop(while_loop_cond_graph_0, while_loop_body_graph_0, (arg1_1, arg0_1, arg2_1));  while_loop_cond_graph_0 = while_loop_body_graph_0 = arg1_1 = arg0_1 = arg2_1 = None
+    while_loop = torch.ops.higher_order.while_loop(while_loop_cond_graph_0, while_loop_body_graph_0, (arg0_1, arg1_1, arg2_1));  while_loop_cond_graph_0 = while_loop_body_graph_0 = arg0_1 = arg1_1 = arg2_1 = None
     getitem = while_loop[0]
     getitem_1 = while_loop[1]
     getitem_2 = while_loop[2];  while_loop = None
@@ -491,6 +512,69 @@ def test_while_loop_functionalize(self, func_type):
             with mode:
                 self._check_tracing(fn, inp)
 
+    def test_while_loop_nested2_traced(self):
+        fn, inp = _while_loop_tests()["nested2"]
+        graphs = self._check_tracing(fn, inp)
+        gm = graphs["symbolic"]
+        outer_body = gm.while_loop_body_graph_0
+        outer_cond = gm.while_loop_cond_graph_0
+        inner_body = outer_body.while_loop_body_graph_0
+        inner_cond = outer_body.while_loop_cond_graph_0
+        self.assertExpectedInline(gm.code.strip("\n"), """\
+def forward(self, arg0_1, arg1_1, arg2_1, arg3_1):
+    while_loop_cond_graph_0 = self.while_loop_cond_graph_0
+    while_loop_body_graph_0 = self.while_loop_body_graph_0
+    while_loop = torch.ops.higher_order.while_loop(while_loop_cond_graph_0, while_loop_body_graph_0, (arg0_1, arg1_1, arg2_1, arg3_1));  while_loop_cond_graph_0 = while_loop_body_graph_0 = arg0_1 = arg1_1 = arg2_1 = arg3_1 = None
+    getitem = while_loop[0]
+    getitem_1 = while_loop[1]
+    getitem_2 = while_loop[2]
+    getitem_3 = while_loop[3];  while_loop = None
+    return (getitem, getitem_1, getitem_2, getitem_3)
+    """)  # noqa: B950
+        self.assertExpectedInline(outer_body.code.strip("\n"), """\
+def forward(self, arg0_1, arg1_1, arg2_1, arg3_1):
+    while_loop_cond_graph_0 = self.while_loop_cond_graph_0
+    while_loop_body_graph_0 = self.while_loop_body_graph_0
+    while_loop = torch.ops.higher_order.while_loop(while_loop_cond_graph_0, while_loop_body_graph_0, (arg0_1, arg1_1, arg2_1, arg3_1));  while_loop_cond_graph_0 = while_loop_body_graph_0 = arg0_1 = arg1_1 = arg2_1 = arg3_1 = None
+    getitem = while_loop[0]
+    getitem_1 = while_loop[1]
+    getitem_2 = while_loop[2]
+    getitem_3 = while_loop[3];  while_loop = None
+    sub = torch.ops.aten.sub.Tensor(getitem, 1);  getitem = None
+    clone = torch.ops.aten.clone.default(getitem_1);  getitem_1 = None
+    mul = torch.ops.aten.mul.Tensor(getitem_2, 2);  getitem_2 = None
+    div = torch.ops.aten.div.Tensor(getitem_3, 2);  getitem_3 = None
+    return (sub, clone, mul, div)
+    """)  # noqa: B950
+        self.assertExpectedInline(outer_body.code.strip("\n"), """\
+def forward(self, arg0_1, arg1_1, arg2_1, arg3_1):
+    while_loop_cond_graph_0 = self.while_loop_cond_graph_0
+    while_loop_body_graph_0 = self.while_loop_body_graph_0
+    while_loop = torch.ops.higher_order.while_loop(while_loop_cond_graph_0, while_loop_body_graph_0, (arg0_1, arg1_1, arg2_1, arg3_1));  while_loop_cond_graph_0 = while_loop_body_graph_0 = arg0_1 = arg1_1 = arg2_1 = arg3_1 = None
+    getitem = while_loop[0]
+    getitem_1 = while_loop[1]
+    getitem_2 = while_loop[2]
+    getitem_3 = while_loop[3];  while_loop = None
+    sub = torch.ops.aten.sub.Tensor(getitem, 1);  getitem = None
+    clone = torch.ops.aten.clone.default(getitem_1);  getitem_1 = None
+    mul = torch.ops.aten.mul.Tensor(getitem_2, 2);  getitem_2 = None
+    div = torch.ops.aten.div.Tensor(getitem_3, 2);  getitem_3 = None
+    return (sub, clone, mul, div)
+    """)  # noqa: B950
+        self.assertExpectedInline(inner_body.code.strip("\n"), """\
+def forward(self, arg0_1, arg1_1, arg2_1, arg3_1):
+    clone = torch.ops.aten.clone.default(arg0_1);  arg0_1 = None
+    sub = torch.ops.aten.sub.Tensor(arg1_1, 1);  arg1_1 = None
+    add = torch.ops.aten.add.Tensor(arg2_1, 3.14);  arg2_1 = None
+    sub_1 = torch.ops.aten.sub.Tensor(arg3_1, 2.71);  arg3_1 = None
+    return (clone, sub, add, sub_1)
+    """)
+        self.assertExpectedInline(inner_cond.code.strip("\n"), """\
+def forward(self, arg0_1, arg1_1, arg2_1, arg3_1):
+    gt = torch.ops.aten.gt.Scalar(arg1_1, 0);  arg1_1 = None
+    return gt
+    """)
+
     def test_cond_nested_traced(self):
         def true_nested(y):
             return y * y
diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
@@ -802,8 +802,20 @@ def _check_supported_callable(fn_var):
             cond_graph,
             cond_lifted_freevars,
         ) = speculate_subgraph(
-            tx, args[0], operands, {}, "while_loop", source_target=self.value
+            tx,
+            args[0],
+            operands,
+            {},
+            "while_loop",
+            source_target=self.value,
+            set_subgraph_inputs="manual",
         )
+        if len(cond_lifted_freevars) > 0:
+            unimplemented(
+                f"while_loop's cond_fn doesn't support capturing free variables yet."
+                f" All used inputs must be passed in as arguments explicitly."
+                f" Proxies for the lifted vars:{cond_lifted_freevars}."
+            )
         cond_nn_modules = dict(tx.output.nn_modules)
         if not isinstance(cond_r, TensorVariable):
             unimplemented(
@@ -831,28 +843,17 @@ def _check_supported_callable(fn_var):
             {},
             "while_loop",
             source_target=self.value,
+            set_subgraph_inputs="manual",
             should_flatten_outputs=True,
         )
         body_nn_modules = dict(tx.output.nn_modules)
 
-        (
-            cond_graph,
-            body_graph,
-            cond_shared,
-            body_shared,
-            cond_unique,
-            body_unique,
-        ) = _merge_graph_inputs(
-            cond_graph,
-            cond_lifted_freevars,
-            "cond_fn",
-            body_graph,
-            body_lifted_freevars,
-            "body_fn",
-        )
-        # We pick cond_shared but it shouldn't matter
-        merged_input = tuple(cond_shared + cond_unique + body_unique)
-
+        if len(body_lifted_freevars) > 0:
+            unimplemented(
+                f"while_loop's body_fn doesn't support capturing free variables yet."
+                f" All used inputs must be passed in as arguments explicitly."
+                f" Proxies for the lifted vars:{body_lifted_freevars}."
+            )
         cond_name = add_subgraph(
             tx,
             self.source,
@@ -872,7 +873,7 @@ def _check_supported_callable(fn_var):
         p_args = (
             cond_node,
             body_node,
-            merged_input,
+            tuple([operand.as_proxy() for operand in operands]),
         )
 
         return _call_function_and_unflatten_output(
