Skip to content
Snippets Groups Projects
Commit ec715ef9 authored by Jane Liu's avatar Jane Liu Committed by TensorFlower Gardener
Browse files

PR #10636: Offloading 1/3: Add annotation for copy-start/copy-done

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

Add stream id in the backend_config of copy-start instruction. The stream id is obtained from hlo_query::NextChannelId().
The corresponding copy-done instruction which is the use of copy-start instruction will be traversed and added the stream id in the backend_config too. This part is automatically done by the subsequent AnnotateStreamAttributesForUsers() existing in the function.
The bool data member copy_start_done_ is used to differentiate copy-start/copy-done from other collective instructions and go through two different paths.
https://github.com/openxla/xla/pull/10450 is split and the current PR is the first 1 out of 3 PRs.
Copybara import of the project:

--
ff99c161a634b868d4265204d10d0b80adf2e772 by Jane Liu <janeliu@nvidia.com>:

Add annotation of stream id for copy-start and its use of copy-done instruction

--
64c746a5de6aa4c9370034ed712192dfd78e3a6e by Jane Liu <janeliu@nvidia.com>:

Enable the annotator for copy-start/copy-done in gpu compiler

--
7972b987e45cd165834483c4350e953094b5dbe8 by Jane Liu <janeliu@nvidia.com>:

Add the annotator pass after HLO rematerialization pass

--
2a9e317aac85e4de3f0ff5f1558885408ab4562d by Jane Liu <janeliu@nvidia.com>:

Add the dependency in BUILD

--
db5ed79f0ec90a1032252beff5108e4f01b61c4a by Jane Liu <janeliu@nvidia.com>:

Use a function to annotate copy-start and add description

--
b4b3932bfedfabc7f3e22a8766ecbc1fa1188402 by Jane Liu <janeliu@nvidia.com>:

remove the bool var copy_start from the StreamAttributeAnnotator class

--
11148fe63588f58c85ab2530bda64807bf400476 by Jane Liu <janeliu@nvidia.com>:

Fixes according to the code review

Merging this change closes #10636

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/10636 from zhenying-liu:offloading/annotator 11148fe63588f58c85ab2530bda64807bf400476
PiperOrigin-RevId: 617496055
parent d72b16ff
No related merge requests found
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment