Repository: pytorch
Issue: torch.compile reorder_for_compute_comm_overlap sink_waits pass does not work
Labels: oncall: distributed,triaged,oncall: pt2

### 🐛 Describe the bug

The sink_waits pass fails to optimize a trivial example due to a nop node creating a fake use

### Error logs

_No response_

### Minified repro

Running
```
import torch
import torch._inductor.config
torch._inductor.config.reorder_for_compute_comm_overlap = True
torch.distributed.init_process_group(backend="nccl")
from torch.distributed._functional_collectives import all_reduce

torch.set_default_device(f'cuda:{torch.distributed.get_rank()}')
print(torch.distributed.get_rank())

def toy(x, y):
    out1 = all_reduce(x, "sum", '0')
    # out1.wait() is scheduled here
    out2 = y * 2
    # out1.wait() should be scheduled here for optimal performance
    return out1, out2

toy = torch.compile(toy, backend='inductor')
args = (torch.zeros((3,)), torch.ones((4,)))
print(toy(*args))
print('done')
```
The inductor output is
```
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((3, ), (1, ), torch.float32)
        # Source Nodes: [tensor], Original ATen: [_c10d_functional.all_reduce]
        stream0 = get_raw_stream(0)
        triton_poi_fused_all_reduce_0.run(arg0_1, buf0, 3, grid=grid(3), stream=stream0)
        del arg0_1
        # Source Nodes: [tensor], Original ATen: [_c10d_functional.all_reduce]
        buf1 = torch.ops._c10d_functional.all_reduce_.default(buf0, 'sum', '0')
        # Source Nodes: [out1], Original ATen: [_c10d_functional.wait_tensor]
        buf3 = torch.ops._c10d_functional.wait_tensor.default(buf0)
        buf5 = empty_strided_cuda((4, ), (1, ), torch.float32)
        # Source Nodes: [out2], Original ATen: [aten.mul]
        triton_poi_fused_mul_1.run(arg1_1, buf5, 4, grid=grid(4), stream=stream0)
        del arg1_1
```
This is suboptimal because the wait can be moved to after the multiply to overlap comm with compute. (In fact, if I add a second collective operation, on out2, before the return, then `reorder_compute_for_overlap` correctly optimzies the waits.)

The issue is that the inductor graph looks like

```
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
    %buf0 : [num_users=3] = call_function[target=torch._inductor.debug.compute: fused_all_reduce](args = (%arg0_1,), kwargs = {device: cuda:0})
    %buf1 : [num_users=2] = call_function[target=torch._inductor.debug.extern: fused_all_reduce](args = (%buf0,), kwargs = {device: cuda:0})
    %buf2 : [num_users=3] = call_function[target=torch._inductor.debug.nop: fused_all_reduce](args = (%buf0, %buf1), kwargs = {device: cuda:0})
    %buf3 : [num_users=1] = call_function[target=torch._inductor.debug.extern: fused_wait_tensor](args = (%buf2,), kwargs = {device: cuda:0})
    %buf4 : [num_users=1] = call_function[target=torch._inductor.debug.nop: fused_wait_tensor](args = (%buf1, %buf2, %buf3), kwargs = {device: cuda:0})
    %buf5 : [num_users=1] = call_function[target=torch._inductor.debug.compute: fused_mul](args = (%arg1_1,), kwargs = {device: cuda:0})
    return (buf0, buf2, buf4, buf5)
```

where `buf4` is a nop which depends on `buf3`, the wait, and so when `sink_waits` gets to the nop and runs https://github.com/pytorch/pytorch/blob/f600faf2480ddd6600ad88fbfc5dd28da132d61d/torch/_inductor/comms.py#L27-L30

it schedules the wait early because it thinks the nop depends on the wait.

### Versions

Nightly

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @ezyang @msaroufim @bdhirsh @anijain2305

@yf225