Enable AllGather Triton Backend

[mfrancepillois]

This PR enables AllGather triton backend:

• Enhances the common collective-emitter to handle the AllGather op (that returns a tuple)
• Adds 2 tritons kernel implementations: one default implementation and one using swizzled for loading data
• Adds e2e tests.
  (This support needed the triton-xla atomics operations to be implemented. That's why it is based on top of the branch ci_maxime_allreduce_triton_rocm_elementwise_rocm)

[claude]

Review Summary

This PR extends the collective emitter infrastructure (originally built for AllReduce) to support AllGather via the Triton backend. It adds two kernel implementations (default and swizzled), tuple unpacking for AllGatherStart's (input, output) shape, a new KernelArguments::Create overload, and comprehensive e2e tests. The feature is gated behind xla_gpu_unsupported_use_all_gather_triton_backend (default: false).

Key issues found:

• Bug — dtype check bypassed for AllGatherStart: all_gather.shape().element_type() returns TUPLE for AllGatherStart, so the F8/S4 unsupported-type guard never fires. Should use the operand's element type instead.
• Bug — inconsistent kMaxBlocksPerGrid constants: Launch dimensions use 32, but signal buffer shape uses 24. Should be a single shared constant.
• Correctness — GetTupleElement identity indexing map: The new GTE case in ComputeOutputToInputIndexing maps output shape to the tuple operand, which is semantically incorrect for general GTE operations. Needs a guard or scoping.
• Dead code with latent bugs: EmitAllGatherSwizzled hardcodes gather dim=0 and has a potential division-by-zero. Both will break if the swizzled path is enabled.
• Logging noise: ~24 new LOG(INFO) calls on common code paths (including non-collective Triton fusions). Should be VLOG(n).

Details in inline comments.

Automated review by Claude

[i-chaochen]

wondering is this branch is based on upstream or xla-0.9.1?

[mfrancepillois]

wondering is this branch is based on upstream or xla-0.9.1?

This branch is based on ci_maxime_allreduce_triton_rocm_elementwise_rocm because we need triton_xla atomic operations to be implemented. But ci_maxime_allreduce_triton_rocm_elementwise_rocm is based on rocm-jaxlib-v0.9.1.

[mfrancepillois]

Currently, the swizzled kernel is not called but I'm keeping it until the performance evaluation is complete.