PR #11182: [NVIDIA GPU] Optimize collective permute using memcpy
PR #11182: [NVIDIA GPU] Optimize collective permute using memcpy

Imported from GitHub PR https://github.com/openxla/xla/pull/11182

When the source and target pair of a collective permute instruction are located within the same node, we will try to use cudaMemCpyPeerToPeer to send data instead of invoking nccl send and recv kernels. Because memcpy will use the copy engine on the GPU instead of launching a kernel, copy engine can fully saturate the bandwidth without occupying any SM.
Since memcpyP2p needs the destination pointer of the other device, we first communicate the pointer values with peer using nccl calls and use the pointer to invoke memcpy.
Some perf comparison:
sending a tensor of this size bf16[512,24576]{1,0}:
Using collective-permute: 445 us
using memcpy: 100us
Copybara import of the project:

--
73a2568520d411eb6f62b6383a3841cb0f8edf22 by TJ <tjx@nvidia.com>:

optimize collective permute using memcpy

--
1dd896f544852fa898e937af8c0b94e075db3ec2 by TJ <tjx@nvidia.com>:

Remove nccl calls and use Async value

--
a2fbc7d29d3d8106753f36e2a477e0889da6cb58 by TJ <tjx@nvidia.com>:

wrap async map and mutex with a class

Merging this change closes #11182

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11182 from Tixxx:tixxx/optimize_cp a2fbc7d29d3d8106753f36e2a477e0889da6cb58
