[XLA] Fix a bug in SchedulingAnnotationCrossesOverlapLimit.
Before this code change, we checked whether scheduling each instruction in a given group individually would cross any limits. This check was not strong enough because i) async instructions in the group are supposed to overlap each other and ii) total resource usage of the group can still exceed the limit while individual usages do not. For example, if the all-gather limit is 1, we should not allow a group with 2 async all-gathers to be scheduled. With this code change, we compute the "accumulated" resource usage of the annotation group and compare that against the limit. PiperOrigin-RevId: 726684858
Showing
- ci/official/envs/linux_x86_cuda 1 addition, 1 deletionci/official/envs/linux_x86_cuda
- third_party/xla/xla/service/BUILD 1 addition, 0 deletionsthird_party/xla/xla/service/BUILD
- third_party/xla/xla/service/hlo_runner_pjrt.cc 2 additions, 3 deletionsthird_party/xla/xla/service/hlo_runner_pjrt.cc
- third_party/xla/xla/service/latency_hiding_scheduler.cc 39 additions, 3 deletionsthird_party/xla/xla/service/latency_hiding_scheduler.cc
- third_party/xla/xla/service/latency_hiding_scheduler.h 2 additions, 0 deletionsthird_party/xla/xla/service/latency_hiding_scheduler.h
- third_party/xla/xla/service/latency_hiding_scheduler_test.cc 9 additions, 23 deletionsthird_party/xla/xla/service/latency_hiding_scheduler_test.cc
Please register or sign in to comment