Show hlo instruction name corresponding to each program order in memo…

…ry viewer. PiperOrigin-RevId: 687094536
tensorflow · Oct 18, 2024 · 6a1c3a6 · 6a1c3a6
1 parent 5ac49e3
commit 6a1c3a6
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 1 deletion.
diff --git a/frontend/app/common/interfaces/memory_viewer_preprocess.jsonpb_decls.d.ts b/frontend/app/common/interfaces/memory_viewer_preprocess.jsonpb_decls.d.ts
@@ -43,6 +43,7 @@ export interface BufferAllocation {
 export interface PreprocessResult {
   heapSizes?: /* double */ number[];
   unpaddedHeapSizes?: /* double */ number[];
+  hloInstructionNames?: string[];
   maxHeap?: HeapObject[];
   maxHeapBySize?: HeapObject[];
   logicalBufferSpans?: {[key: /* int32 */ string]: BufferSpan};

diff --git a/frontend/app/components/memory_viewer/memory_usage/memory_usage.ts b/frontend/app/components/memory_viewer/memory_usage/memory_usage.ts
@@ -24,6 +24,7 @@ export class MemoryUsage {
   indefiniteMemoryUsageBytes: MemoryUsageBytes;
   heapSizes: number[];
   unpaddedHeapSizes: number[];
+  hloInstructionNames: string[];
   maxHeap: HeapObject[];
   maxHeapBySize: HeapObject[];
   maxHeapByPaddingSize: HeapObject[];
@@ -55,6 +56,7 @@ export class MemoryUsage {
     this.indefiniteMemoryUsageBytes = {padded: 0, unpadded: 0};
     this.heapSizes = [];
     this.unpaddedHeapSizes = [];
+    this.hloInstructionNames = [];
     this.maxHeap = [];
     this.maxHeapBySize = [];
     this.maxHeapByPaddingSize = [];
@@ -115,6 +117,7 @@ export class MemoryUsage {
     this.peakHeapSizePosition = (preprocess.peakHeapSizePosition || 0);
     this.heapSizes = preprocess.heapSizes || [];
     this.unpaddedHeapSizes = preprocess.unpaddedHeapSizes || [];
+    this.hloInstructionNames = preprocess.hloInstructionNames || [];
     if (preprocess.logicalBufferSpans) {
       for (const [key, value] of Object.entries(
                preprocess.logicalBufferSpans)) {

diff --git a/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ng.html b/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ng.html
@@ -28,6 +28,7 @@
 <program-order-chart
   [heapSizes]="heapSizes"
   [unpaddedHeapSizes]="unpaddedHeapSizes"
+  [hloInstructionNames]="hloInstructionNames"
   [peakInfo]="peakInfo"
   [activeInfo]="activeInfo"
   [timelineUrl]="timelineUrl">

diff --git a/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ts b/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ts
@@ -50,6 +50,7 @@ export class MemoryViewerMain implements OnDestroy, OnChanges {
   selectedIndexBySize: number = -1;
   selectedIndexByPaddingSize: number = -1;
   unpaddedHeapSizes: number[] = [];
+  hloInstructionNames: string[] = [];
   hasTrace = false;
   diagnostics: Diagnostics = {info: [], warnings: [], errors: []};
 
@@ -164,6 +165,7 @@ export class MemoryViewerMain implements OnDestroy, OnChanges {
         (this.usage.hloTempFragmentation * 100.0).toFixed(2);
     this.heapSizes = this.usage.heapSizes || [];
     this.unpaddedHeapSizes = this.usage.unpaddedHeapSizes || [];
+    this.hloInstructionNames = this.usage.hloInstructionNames || [];
     this.peakInfo = {
       size: utils.bytesToMiB(this.usage.peakHeapSizeBytes),
       alloc: this.usage.peakHeapSizePosition + 1,

diff --git a/frontend/app/components/memory_viewer/program_order_chart/program_order_chart.ts b/frontend/app/components/memory_viewer/program_order_chart/program_order_chart.ts
@@ -18,6 +18,9 @@ export class ProgramOrderChart implements OnChanges, OnInit {
   /** The unpadded heap size list. */
   @Input() unpaddedHeapSizes: number[] = [];
 
+  /** The HLO instruction name corresponding to each program order point. */
+  @Input() hloInstructionNames: string[] = [];
+
   /** The peak buffer allocation information. */
   @Input() peakInfo?: BufferAllocationInfo;
 
@@ -125,14 +128,23 @@ export class ProgramOrderChart implements OnChanges, OnInit {
     for (let i = 0; i < this.heapSizes.length; i++) {
       this.maxSize = Math.max(
           this.maxSize, Math.max(this.heapSizes[i], this.unpaddedHeapSizes[i]));
-      data.push([i, this.heapSizes[i], this.unpaddedHeapSizes[i]]);
+      const tooltip = `<div>
+        Program Order: ${i}<br>Size: ${this.heapSizes[i].toFixed(1)}<br>
+        Unpadded Size: ${this.unpaddedHeapSizes[i].toFixed(1)}<br>
+        HLO instruction: ${this.hloInstructionNames[i]}
+        </div>`;
+      data.push([i,
+                this.heapSizes[i], tooltip,
+                this.unpaddedHeapSizes[i],tooltip]);
     }
     this.maxSize = Math.round(this.maxSize * 1.1);
 
     const dataTable = new google.visualization.DataTable();
     dataTable.addColumn('number', 'Schedule');
     dataTable.addColumn('number', 'Size');
+    dataTable.addColumn({type: 'string', role: 'tooltip', 'p': {'html': true}});
     dataTable.addColumn('number', 'Unpadded Size');
+    dataTable.addColumn({type: 'string', role: 'tooltip', 'p': {'html': true}});
     dataTable.addRows(data);
 
     const options: google.visualization.LineChartOptions = {
@@ -157,6 +169,7 @@ export class ProgramOrderChart implements OnChanges, OnInit {
         },
       },
       legend: {position: 'top'},
+      tooltip: {isHtml: true},
     };
 
     this.heapChartDataInfo = {