Get channel_id in HLO dump

[DhruvKumar1]

I'm interested in looking at the HLO dumps from XLA when running a JAX program. I was looking through some of the multi-gpu operations like all_gather and all_reduce. I set --xla_dump_to= and am able to see the HLO dumps (before/after optimizations). For the all_reduce I see replica groups and the size of tensors. However, one thing that I am not able to see is the channel_id that is mentioned as an optional argument in the XLA Operation Semantics.

I have seen online some HLO snippets where there is a channel_id populated in the all-reduce-start. Could someone please explain when this field is populated, and a possible simple python example that would induce an HLO with an all_gather with this field populated?

This is the simple code that I am running that induces an all_gather and an all_reduce. (running on 2 GPUs)

import jax
import numpy as np

x = np.array([[1,2,3,4],[5,6,7,8]])

def f(x):
        return jax.lax.all_gather(x, "i")

y = jax.pmap(f, axis_name="i")(x)
out = jax.pmap(lambda x:jax.lax.psum(x,"i"), axis_name="i")(y)

this is the output I am getting for the module relating to the all_reduce

region_0.2 {
  Arg_1.4 = s32[] parameter(1)
  Arg_0.3 = s32[] parameter(0)
  ROOT add.1 = s32[] add(Arg_0.3, Arg_1.4), metadata={op_name="pmap(<lambda>)/jit(main)/add" source_file="/home/training.py" source_line=10}
}

wrapped_copy_computation {
  param_0 = s32[4]{0} parameter(0)
  ROOT copy.3 = s32[4]{0} copy(param_0)
}

ENTRY main.7 {
  Arg_0.1.0 = s32[4]{0} parameter(0)
  wrapped_copy = s32[4]{0} fusion(Arg_0.1.0), kind=kLoop, calls=wrapped_copy_computation
  all-reduce-start = s32[4]{0} all-reduce-start(wrapped_copy), replica_groups={{0,1,2}}, to_apply=region_0.2, metadata={op_name="pmap(<lambda>)/jit(main)/psum[axes=(\'i\',) axis_index_groups=None]" source_file="/home/training.py" source_line=10}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"collective_backend_config":{"is_sync":true,"no_parallel_custom_call":false},"force_earliest_schedule":false}
  ROOT all-reduce-done = s32[4]{0} all-reduce-done(all-reduce-start), metadata={op_name="pmap(<lambda>)/jit(main)/psum[axes=(\'i\',) axis_index_groups=None]" source_file="/home/training.py" source_line=10}
}