[XLA] Add stack trace breakdown to HloLiveRange::ToString for peak memory usage
#94954
+717
−2
Conversation
This file contains hidden or 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
113f9ff to
7bfd1e7
Compare
f0f9320 to
5ef0ffe
Compare
67a5b2c to
a765b82
Compare
7127b19 to
5dedef2
Compare
667f369 to
532df3a
Compare
…memory usage
This change enhances the `HloLiveRange::ToString` output with a hierarchical stack trace visualization of live buffers at the peak memory usage moment. The breakdown groups buffer sizes by call stack, making it easier to identify high-memory-use paths in nested computations.
Key additions:
* Introduced `FormatStackTraceBreakdown` in a new file `hlo_stack_trace.{h,cc}` to encapsulate stack trace logic.
* Refactored `ToString` to delegate stack tree formatting to this helper.
* Added unit test `hlo_stack_trace_test.cc` to verify formatting, including nested computations and buffer attribution.
This improves interpretability of peak memory use, helping developers and memory profiling tools pinpoint memory-intensive subcomputations.
PiperOrigin-RevId: 786270002
532df3a to
89e79b5
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] Add stack trace breakdown to
HloLiveRange::ToStringfor peak memory usageThis change enhances the
HloLiveRange::ToStringoutput with a hierarchical stack trace visualization of live buffers at the peak memory usage moment. The breakdown groups buffer sizes by call stack, making it easier to identify high-memory-use paths in nested computations.Key additions:
FormatStackTraceBreakdownin a new filehlo_stack_trace.{h,cc}to encapsulate stack trace logic.ToStringto delegate stack tree formatting to this helper.hlo_stack_trace_test.ccto verify formatting, including nested computations and buffer attribution.This improves interpretability of peak memory use, helping developers and memory profiling tools pinpoint memory-intensive subcomputations.