[xla][gpu] Make the p2p-schedule-preparation pass to not rely on certain
[xla][gpu] Make the p2p-schedule-preparation pass to not rely on certain
pattern to find the corresponding Send/Recv for SendDone/RecvDone.

Previously, we relied on certain code patterns to find the Send/Recv for
Send-done/Recv-done for a pipelined loop. Now that we realize copy-insertion
can complicate such code patterns and make such pattern matching fragile.

We modify the p2p-schedule-preparation pass to record Send/Recv in additional
to SendDone/RecvDone for each p2p communication group.
