[FEATURE] Support overlapping pipeline communication and computation

[ZYHowell]

This PR:

• Refactor XLANcclGroup's implementation. We move more things into the XLA side to make code cleaner.
• Support overlapping cross-mesh communication and SPMD computation. We do so by adding GpuEvents during the pipeline execution, then use the event level synchronization. The overlapping is only supported with XLA implementation, because with CuPy implementation, there is a CPU-GPU synchronization when transforming from Xla buffer to Cupy buffers.
• Fix bugs in apply gradient. Previous version cannot handle the following case: a and b are the output of layer i and j, respectively. c = a + b is used in layer k.
• Add overlap friendly pipeline schedule. For more information, please refer to this paper. The broadcast optimization was already merged in earlier prs.
• Add U-Net Transformer model implementation and benchmark script. The model is modified from huggingface diffusers repo, and we insert pipeline markers at correct places.

TODO of next PR:

• Support HloInstruction Rescheduling to overlap backward communication with computation.
  The current blocker is that some code seems to synchronize cpu thread to gpus. We need to find out where they are.
• Either run dynamic slice on other streams or avoid generating shards that make communicated buffers uncontinuous.
  This is because in the slow path, there is a dynamic slice enqueued to the end of computational stream, making communication have to wait until all computations are done, which has no overlap then.