Only make wait_tensor as a side_effect op (pytorch#132341)

Summary:
pytorch#131023 add all the collective ops to the side effect list. But we should only make wait_tensor as a side_effect op because all collective ops should have a corresponding wait_tensor.

We should switch to use high_order effect token.

Pull Request resolved: pytorch#132341
Approved by: https://github.com/yf225

Author: fegin

test/distributed/test_functional_api.py

@@ -587,6 +587,28 @@ def allreduce(t, pg):
         )
         allreduce(torch.randn(8, device=self.device), pg=dist.group.WORLD)
 
+    @unittest.skipIf(not has_triton(), "Inductor+gpu needs triton and recent GPU arch")
+    @requires_nccl()
+    @with_comms()
+    def test_tracing_with_dce_code(self):
+        if self.world_size > 2:
+            return
+
+        def func(batch, group, rank):
+            ret = ft_c.permute_tensor(batch, [1, 0], group)
+            if hasattr(ret, "wait"):
+                ret = ret.wait()
+            if rank == 0:
+                return ret
+            else:
+                return batch * 5
+
+        compiled_func = torch.compile(func)
+        ret = compiled_func(
+            torch.ones((100,), device="cuda"), self.process_group, self.rank
+        )
+        dist.barrier()
+
 
 class TestNCCLCollectivesWithWorldSize4(TestCollectivesWithNCCL):
     @property

torch/distributed/_functional_collectives.py

@@ -963,66 +963,7 @@ def _reduce_scatter_tensor_coalesced_native_meta(
 
     # mark these ops has side effect so that they won't be removed by DCE
     torch.fx.node.has_side_effect(torch.ops._c10d_functional.wait_tensor.default)
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.all_gather_into_tensor_out.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.all_gather_into_tensor.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default
-    )
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_reduce.default)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_reduce_.default)
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.all_reduce_coalesced.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.all_reduce_coalesced_.default
-    )
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_to_all_single.default)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.broadcast.default)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.broadcast_.default)
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.reduce_scatter_tensor.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.reduce_scatter_tensor_coalesced.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional_autograd.all_to_all_single.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional_autograd.reduce_scatter_tensor.default
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional_autograd.all_gather_into_tensor.default
-    )
-    # also the no-overload version
     torch.fx.node.has_side_effect(torch.ops._c10d_functional.wait_tensor)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_gather_into_tensor_out)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_gather_into_tensor)
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.all_gather_into_tensor_coalesced
-    )
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_reduce)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_reduce_)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_reduce_coalesced)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_reduce_coalesced_)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.all_to_all_single)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.broadcast)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.broadcast_)
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional.reduce_scatter_tensor)
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional.reduce_scatter_tensor_coalesced
-    )
-    torch.fx.node.has_side_effect(torch.ops._c10d_functional_autograd.all_to_all_single)
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional_autograd.reduce_scatter_tensor
-    )
-    torch.fx.node.has_side_effect(
-        torch.ops._c10d_functional_autograd.all_gather_into_tensor
-    )
 
     # Register legacy ops for backward compatibility
     # TODO(yifu): remove these in functional collective beta release