[RFC] Add collective_broadcast to the StableHLO specification

[chaserileyroberts]

This RFC proposes adding collective_broadcast as one of the collective communication primitives.
Please provide any feedback you feel is valuable.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[GleasonK]

View this failed invocation of the CLA check for more information.

Also - could you follow this link and sign the CLA when you get a chance.

[chaserileyroberts]

Also - could you follow this link and sign the CLA when you get a chance.

I work at Nvidia so I need to go through the separate channel. I'll take care of it don't worry.

[chaserileyroberts]

Made updates based on the new semantics.

[andydavis1]

Thank you for the RFC. The current set of collective operations are currently used in the SPMD context, where each device is running the same program. With the proposed collective_broadcast, do you have an MPMD use case for this, or were you expecting other devices which were not broadcasting to produce an empty buffer as a result of the broadcast?

[chaserileyroberts]

With the proposed collective_broadcast, do you have an MPMD use case for this, or were you expecting other devices which were not broadcasting to produce an empty buffer as a result of the broadcast

No, the idea is for this to still be targeting exclusively SPMD. On the cuda side, the goal is to have this op lower to the exact same nccl.Broadcast operation on all devices. Devices that are not in any replica_group of the SPMD op return zeros, as per the specification and latest example.

[chaserileyroberts]

To give a very concrete example of how this operation could be used in a realistic SPMD setting, let me show you some code.

Here, we have an example of how we could implement the 2D pgemm algorithm summa in JAX once this spec is in place.

@partial(shard_map, mesh=Mesh(...), 
    in_specs=(P('x', 'y'), P('x', 'y')), out_specs=P('x', 'y'))
def summa_matrix_multiply(a, b):
  for i in range(N):
     abcast = pcollective_broadcast(a, 'y', root=i)
     bbcast = pcollective_broadcast(b, 'x', root=i)
     if i == 0:
        c = abcast @ bbcast
     c += abcast @ bbcast
  return  c

[GleasonK]

RFC Approved. I'll send a follow-up with markdownlint fixes.

[chaserileyroberts]

🥳 🎉

Thanks again Kevin for getting this over the finish line! I'll get to work on getting this implemented in JAX and cuda xla backend.