Repository: pytorch
Issue: Make broadcast_coalesced to a op for processgroup
Labels: oncall: distributed

### 🚀 The feature, motivation and pitch

Today, broadcast_coalesced functions similarly to braoadcast, like a set communication operator. It's defined here:https://github.com/pytorch/pytorch/blob/4d8906885e672d088849303cc8e821cff81ba6d9/torch/csrc/distributed/c10d/comm.hpp#L12-L16

Because of the custom device part, we may need to have our own special logic for broadcast_coalesced, so one way to make broadcast_coalesced more general is to make it a collection communication operator so that users can customize their own broadcast_coalesced.



### Alternatives

broadcast_coalesced is a generic method for cpu and cuda, so an easy way to do this is to register the current implementation as the default operation. broadcast_coalesced internally uses different backends depending on the device. 

If users need to implement special broadcast_coalesced for cpu/cuda/custom_device, they can register their implementations to the specific device.

Here is a sample:

First, broadcast_coalesced operator needs to be declared and register for all device, and then it will be called inside ProcessGroup

``` C++
// torch/csrc/distributed/c10d/Ops.cpp
  m.def(
      "broadcast_coalesced_(Tensor[] tensors, __torch__.torch.classes.c10d.ProcessGroup process_group, int buffer_size, int rank) -> ()");

void broadcast_coalesced_comm_(
    at::TensorList tensors,
    const c10::intrusive_ptr<ProcessGroup>& process_group,
    int64_t buffer_size,
    int64_t rank) {
  // use current implementation as a default logic
  broadcast_coalesced(process_group, tensors, buffer_size, rank);
  return;
}

TORCH_LIBRARY_IMPL(c10d, CompositeExplicitAutograd, m) {
  m.impl("broadcast_coalesced_", broadcast_coalesced_comm_);
}
```


```C++
// torch/csrc/distributed/c10d/ProcessGroup.hpp
  virtual void broadcast_coalesced(
      std::vector<at::Tensor>& tensors,
      const BroadcastCoalescedOptions& opts = BroadcastCoalescedOptions()) {
    static auto op =
        c10::Dispatcher::singleton()
            .findSchemaOrThrow("c10d::broadcast_coalesced_", "")
            .typed<void(
                    at::TensorList,
                    const
                    c10::intrusive_ptr<::c10d::ProcessGroup>&,
                    int64_t,
                    int64_t)>();
    return op.call(
        tensors,
        c10::intrusive_ptr<ProcessGroup>::unsafe_reclaim_from_nonowning(this),
        opts.bufferSize,
        opts.rank);
  }
```

Of course, we need to bind it to the python
```C++
// torch/csrc/distributed/c10d/init.cpp
          .def(
              "broadcast_coalesced",
              &::c10d::ProcessGroup::broadcast_coalesced,
              py::arg("tensors"),
              py::arg("opts") = ::c10d::BroadcastCoalescedOptions(),
              py::call_guard<py::gil_scoped_release>())

```
```python
// torch/distributed/distributed_c10d.py
def _broadcast_coalesced(group, tensors, buffer_size, src=0):
    # maybe we need to retain the original calling function for ProcessGroupNCCL which is not ProcessGroup
    # rename original _broadcast_coalesced to  _broadcast_coalesced_
    if not issubclass(type(group), ProcessGroup):
        dist._broadcast_coalesced_(group, tensors, buffer_size, src)
        return
    opts = BroadcastCoalescedOptions()
    opts.bufferSize = buffer_size
    opts.rank = src

    return group.broadcast_coalesced(tensors, opts)

```



### Additional context

_No response_

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu