[XLA:MSA] Fix the AllocationRequest for window prefetch #21736
+115
−76
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_718439738
branch
3 times, most recently
from
0e7d18d to
4de2457
Compare
We currently have two performance issues in window prefetch. The first is that we specified the created WindowPrefetchedAllocation to consume CopyResources by using non-zero shape. This could consume all resources and interfere with the prefetching decision of other tensors. In fact, we don't have prefetching implemented yet, so we could specify using zero CopyResource. Another issue is that the generated allocation from window prefetch currently spans too long in time. Its earliest prefetch time is set to the operand's define time, and its end is the use time. We should keep the earliest prefetch time to be as close as the use time as possible. This is for keeping the interference to the prefetching of other tensors to be minimal. We updated WindowPrefetch() to simply just allocate chunk for the exposed span vmem at this moment. Since we don't call Prefetch() from WindowPrefetch(), we can simplify the data structure of AllocationRequest and PrefetchContext a bit. PiperOrigin-RevId: 718439738
copybara-service
bot
force-pushed
the
test_718439738
branch
from
4de2457 to
2f1d65c
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
[XLA:MSA] Fix the AllocationRequest for window prefetch
We currently have two performance issues in window prefetch.
The first is that we specified the created WindowPrefetchedAllocation to consume CopyResources by using non-zero shape. This could consume all resources and interfere with the prefetching decision of other tensors. In fact, we don't have prefetching implemented yet, so we could specify using zero CopyResource.
Another issue is that the generated allocation from window prefetch currently spans too long in time. Its earliest prefetch time is set to the operand's define time, and its end is the use time. We should keep the earliest prefetch time to be as close as the use time as possible. This is for keeping the interference to the prefetching of other tensors to be minimal.
We updated WindowPrefetch() to simply just allocate chunk for the exposed span vmem at this moment. Since we don't call Prefetch() from WindowPrefetch(), we can simplify the data structure of AllocationRequest and PrefetchContext a bit.