XLA CollectiveOpsUtil does not include kReduceScatter

[ptoulme-aws]

When debugging the XLA while loop unroller channel id creation, we found that HloOpcode::kReduceScatter was not included as a collective to generate new channel id for.

While loop unroller location:

xla/xla/service/while_loop_unroller.cc

Line 201 in 79d745f

if (IsCollectiveWithChannelId(body_inst)) {

Collective ops util location:

xla/xla/service/collective_ops_utils.cc

Line 611 in 79d745f

bool IsCollectiveWithChannelId(const HloInstruction* instruction) {

We have added kReduceScatter locally and it fixed our channel id issues.

Questions:

1. Why is kReduceScatter not listed in the list of collective operations to replace channel id with?
2. Can kReduceScatter be added to the list if this is a bug?

[akuegel]

Can you please check whether f584acb fixes this issue?

[akuegel]

I think it is fixed, please let us know in case this is still a problem.