Support collective matmul optimization in mp

[yaochengji]

This PR:

• add the new environment variable ENABLE_COLLECTIVE_MATMUL_IN_MP to turn on the xla config which is required for collective matmul optimization on TPU
• add channel_id and use_global_device_ids to xm.all_gather and xm.reduce_scatter, which is also required to enable collective matmul
• add a utility function to find the best ring order for v5p x 8 and v6 x 8

[tengyifei]

question: do you think it's possible to test that collective matmul is actually triggered? maybe we could look for the "decomposed_reduce_scatter_while" etc signatures in the optimized HLO from the XLA dump?

[yaochengji]

question: do you think it's possible to test that collective matmul is actually triggered? maybe we could look for the "decomposed_reduce_scatter_while" etc signatures in the optimized HLO from the XLA dump?

Thanks for you review, Yifei! I need to wait for tomorrow's libtpu nightly to enable collective matmul.

[tengyifei]

I need to wait for tomorrow's libtpu nightly to enable collective matmul.

Gotcha. I think it's useful to add a test once we roll in a new libtpu. Otherwise it's hard to tell if this worked. LGTM otherwise.