[PT2][Optimus] Add unbind_stack_to_cat_pass (#132542)

mengluy0125 · pytorchmergebot · commit 919e38424729 · 2024-08-07T02:26:40.000Z
Summary: We observe the stack mpde can be transformed to cat node to elimiate split nodes, which could further enable the unbind cat optimization, thus we add a more advanced pattern to do the graph transformation Test Plan: # unit test ``` CUDA_VISIBLE_DEVICES=3 OC_CAUSE=1 buck2 test //caffe2/test/inductor:split_cat_fx_passes ``` Buck UI: https://www.internalfb.com/buck2/de6c1cda-3d74-4a30-8980-7b209b6fe5dc Test UI: https://www.internalfb.com/intern/testinfra/testrun/12103424042268125 Network: Up: 485KiB Down: 728KiB (reSessionID-2f2c01c3-79bb-4e37-b5be-fb77ec09b264) Jobs completed: 29. Time elapsed: 5:19.8s. Cache hits: 0%. Commands: 4 (cached: 0, remote: 0, local: 4) Tests finished: Pass 9. Fail 0. Fatal 0. Skip 1. Build failure 0 # benchmark ``` CUDA_VISIBLE_DEVICES=3 OC_CAUSE=1 buck2 run mode/opt //scripts/jackiexu0313/pt2:local_model_with_pt2 -- --test_mode batch-split --model_type "ig_ctr" --flow_id 584880697 ``` P1503698962 before and after graph transformation https://www.internalfb.com/intern/diffing/?paste_number=1504050718 Differential Revision: D60411560 Pull Request resolved: #132542 Approved by: https://github.com/jackiexu1992
diff --git a/test/inductor/test_split_cat_fx_passes.py b/test/inductor/test_split_cat_fx_passes.py
@@ -1128,11 +1128,11 @@ def mutate_cat_node_with_some_getitmes(x):
 
         @torch._inductor.config.patch(
             pre_grad_fusion_options={
-                "optimize_cat_inputs_pass": {},
+                "split_cat_to_slices_pass": {},
             },
             post_grad_fusion_options={},
         )
-        def optimize_cat_inputs(x):
+        def split_cat_to_slices(x):
             x_c = x.clone()
             x_c_2 = x.clone()
             l1_out = torch.split(x, [50, 50, 50, 50, 50, 50, 50, 50, 50, 50], dim=0)
@@ -1229,22 +1229,75 @@ def unbind_cat_to_view(x):
             )
             return torch.cat((cat, item7, item8, item9), dim=1)
 
+        @torch._inductor.config.patch(
+            pre_grad_fusion_options={
+                "split_stack_to_cats_pass": {},
+            },
+            post_grad_fusion_options={},
+        )
+        def split_stack_to_cats(x):
+            x_c = x.view(10, 50, 500)
+            l1_out = torch.unbind(x_c, dim=0)
+            item0 = l1_out[0]
+            item1 = l1_out[1]
+            item2 = l1_out[2]
+            item3 = l1_out[3]
+            item4 = l1_out[4]
+            item5 = l1_out[5]
+            split1 = torch.split(item0, [250, 250], dim=1)
+            split2 = torch.split(item1, [250, 250], dim=1)
+            split3 = torch.split(item2, [250, 250], dim=1)
+            split4 = torch.split(item3, [250, 250], dim=1)
+            split5 = torch.split(item4, [250, 250], dim=1)
+            split6 = torch.split(item5, [250, 250], dim=1)
+            getitem0, getitem1 = split1[0], split1[1]
+            getitem2, getitem3 = split2[0], split2[1]
+            getitem4, getitem5 = split3[0], split3[1]
+            getitem6, getitem7 = split4[0], split4[1]
+            getitem8, getitem9 = split5[0], split5[1]
+            getitem10, getitem11 = split6[0], split6[1]
+            getitem0_c = getitem0.clone()
+            getitem1_c = getitem1.clone()
+            getitem2_c = getitem2.clone()
+            return torch.stack(
+                (
+                    getitem0,
+                    getitem1,
+                    getitem2,
+                    getitem3,
+                    getitem4,
+                    getitem5,
+                    getitem0_c,
+                    getitem1_c,
+                    getitem6,
+                    getitem7,
+                    getitem8,
+                    getitem9,
+                    getitem10,
+                    getitem11,
+                    getitem2_c,
+                ),
+                dim=1,
+            )
+
         args = [
             torch.randn(500, 500),
         ]
         for (
             fn,
             expected_getitem_cat_merged,
             expected_cat_removed,
-            expected_cat_optimized,
+            expected_split_cat_to_slices,
             exptected_unbind_to_cat_view,
+            expected_split_stack_to_cats,
         ) in [
-            (split_cat_split, 2, 0, 0, 0),
-            (split_cat_split_kwarg, 2, 0, 0, 0),
-            (remove_cat_node_with_all_getitmes, 0, 2, 0, 0),
-            (mutate_cat_node_with_some_getitmes, 0, 1, 0, 0),
-            (optimize_cat_inputs, 0, 0, 1, 0),
-            (unbind_cat_to_view, 0, 0, 0, 1),
+            (split_cat_split, 2, 0, 0, 0, 0),
+            (split_cat_split_kwarg, 2, 0, 0, 0, 0),
+            (remove_cat_node_with_all_getitmes, 0, 2, 0, 0, 0),
+            (mutate_cat_node_with_some_getitmes, 0, 1, 0, 0, 0),
+            (split_cat_to_slices, 0, 0, 1, 0, 0),
+            (unbind_cat_to_view, 0, 0, 0, 1, 0),
+            (split_stack_to_cats, 0, 0, 0, 0, 1),
         ]:
             expected = fn(*args)
             actual = torch.compile(fn)(*args)
@@ -1259,13 +1312,17 @@ def unbind_cat_to_view(x):
                 expected_cat_removed,
             )
             self.assertEqual(
-                counters["inductor"]["optimize_cat_inputs_pass"],
-                expected_cat_optimized,
+                counters["inductor"]["split_cat_to_slices_pass"],
+                expected_split_cat_to_slices,
             )
             self.assertEqual(
                 counters["inductor"]["unbind_cat_to_view_pass"],
                 exptected_unbind_to_cat_view,
             )
+            self.assertEqual(
+                counters["inductor"]["split_stack_to_cats_pass"],
+                expected_split_stack_to_cats,
+            )
             counters.clear()
 
     def test_numpy_compat_normalization(self):
diff --git a/torch/_inductor/fx_passes/group_batch_fusion.py b/torch/_inductor/fx_passes/group_batch_fusion.py
@@ -916,20 +916,24 @@ def fuse(self, graph: torch.fx.GraphModule, subset: List[torch.fx.Node]):
                     args=(stack_inputs,),
                     kwargs={"inplace": subset[0].kwargs.get("inplace", False)},
                 )
+                batch_op.meta["example_value"] = self.op(
+                    stack_inputs.meta["example_value"],
+                    inplace=subset[0].kwargs.get("inplace", False),
+                )
             else:
                 batch_op = graph.call_function(
                     self.op,
                     args=(stack_inputs,),
                 )
-            if is_node_meta_valid(batch_op):
-                batch_op.meta["example_value"] = self.op(batch_op.meta["example_value"])
+                batch_op.meta["example_value"] = self.op(
+                    stack_inputs.meta["example_value"]
+                )
             unbind_op = graph.call_function(
                 torch.unbind, args=(batch_op,), kwargs={"dim": 0}
             )
-            if is_node_meta_valid(unbind_op):
-                unbind_op.meta["example_value"] = torch.unbind(
-                    unbind_op.meta["example_value"], dim=0
-                )
+            unbind_op.meta["example_value"] = torch.unbind(
+                batch_op.meta["example_value"], dim=0
+            )
             for i, node in enumerate(batch_nodes):
                 with graph.inserting_after(unbind_op):
                     getitem = graph.call_function(operator.getitem, args=(unbind_op, i))
diff --git a/torch/_inductor/fx_passes/split_cat.py b/torch/_inductor/fx_passes/split_cat.py
