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
Showing
- third_party/xla/xla/service/gpu/BUILD 1 addition, 0 deletionsthird_party/xla/xla/service/gpu/BUILD
- third_party/xla/xla/service/gpu/gpu_compiler.cc 1 addition, 0 deletionsthird_party/xla/xla/service/gpu/gpu_compiler.cc
- third_party/xla/xla/service/gpu/stream_attribute_annotator.cc 22 additions, 1 deletion...d_party/xla/xla/service/gpu/stream_attribute_annotator.cc
- third_party/xla/xla/service/gpu/stream_attribute_annotator_test.cc 44 additions, 0 deletions...ty/xla/xla/service/gpu/stream_attribute_annotator_test.cc
Please register or sign in to comment