PR #12892: Add collective-permute-valid-iteration-annotator
PR #12892: Add collective-permute-valid-iteration-annotator

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

This patch adds the `collective-permute-valid-iteration-annotator` pass.

If a while loop is known to be a pipelined while loop, has a known trip count and increments with step=1, then this pass annotates the `collective-permute` operations within the while loop with valid iterations for each GPU. This is only done when the source-target pairs of the `collective-permute` operation form a forward or backward cycle.

For example, if the trip count is 10 (iteration 0 to 9), with step=1, and the source-target pairs of a `collective-permute` operation are `{{0,1},{1,2},{2,3},{3,0}}`, then this pass would annotate such operation with `_xla_send_recv_validation="{{0,6},{1,7},{2,8},{3,9}}"`. This annotation means that
 - for GPU index 0, the valid iterations are 0,1,2,3,4,5,6.
 - for GPU index 1, the valid iterations are 1,2,3,4,5,6,7.
 - for GPU index 2, the valid iterations are 2,3,4,5,6,7,8.
 - for GPU index 3, the valid iterations are 3,4,5,6,7,8,9.

The index in the list denotes the device index and the bounds {start,end} are inclusive.
Copybara import of the project:

--
f5714a8053db0892c950b1e5268f249ff64b2e7b by Shraiysh Vaishay <svaishay@nvidia.com>:

Add collective-permute-valid-iteration-annotator

This patch adds the `collective-permute-valid-iteration-annotator` pass.

If a while loop is known to be a pipelined while loop, has a known trip count
and increments with step=1, then this pass annotates the `collective-permute`
operations within the while loop with valid iterations for each GPU. This is only done
when the source-target pairs of the `collective-permute` operation form a forward or
backward cycle.

For example, if the trip count is 10 (iteration 0 to 9), with step=1, and the source-target
pairs of a `collective-permute` operation are `{{0,1},{1,2},{2,3},{3,0}}`, then this pass
would annotate such operation with `_xla_send_recv_validation="{{0,6},{1,7},{2,8},{3,9}}"`.
This annotation means that

    - for GPU index 0, the valid iterations are 0,1,2,3,4,5,6.
    - for GPU index 1, the valid iterations are 1,2,3,4,5,6,7.
    - for GPU index 2, the valid iterations are 2,3,4,5,6,7,8.
    - for GPU index 3, the valid iterations are 3,4,5,6,7,8,9.

The index in the list denotes the device index and the bounds {start,end} are inclusive.

Merging this change closes #12892

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/12892 from shraiysh:CR-valid-iteration-annotator f5714a8053db0892c950b1e5268f249ff64b2e7b
