Changes:

- Only loading & refresh the right section (not whole page) when changing memory types - Change to use <chart> component for the memory allocation chart, so we can hover and see data tooltip - Updated chart titles, add chart axis titles to make its meaning more straightforward - updated graph-search tooltip and placeholder for clearer call to actions PiperOrigin-RevId: 567765930
tensorflow · Sep 26, 2023 · 5fbf6a1 · 5fbf6a1
1 parent e3a1905
commit 5fbf6a1
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 88 deletions.
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
@@ -3,8 +3,8 @@
 <div class="section-container">
   <div class="row">
     <div>
-      <div class="mat-headline">Module name: {{moduleName}}</div>
-      <div class="sub-title">Peak memory allocation is {{peakHeapSizeMiB}} MiB</div>
+      <div class="mat-headline">Module Name: {{moduleName}}</div>
+      <div class="sub-title">Peak memory allocation: {{peakHeapSizeMiB}} MiB</div>
       <div class="description" [hidden]="!paddingOverhead">
         {{paddingOverhead}} MiB total padding overhead
       </div>
@@ -35,33 +35,26 @@
 
 <div [hidden]="!heapSizes.length">
   <mat-divider ></mat-divider>
+  <div class="heap-chart-header"><h2>HLO Ops at Peak Memory Allocation Time</h2></div>
 </div>
 
 <max-heap-chart
    [maxHeap]="maxHeap"
-   [title]="'By Program Order'"
+   [title]="'by Program Order'"
    [selectedIndex]="selectedIndex"
    (selected)="setSelectedHeapObject($event)">
 </max-heap-chart>
 
-<div [hidden]="!maxHeap.length">
-  <mat-divider ></mat-divider>
-</div>
-
 <max-heap-chart
   [maxHeap]="maxHeapBySize"
-  [title]="'By Size'"
+  [title]="'by Buffer Size'"
   [selectedIndex]="selectedIndexBySize"
   (selected)="setSelectedHeapObjectBySize($event)">
 </max-heap-chart>
 
-<div [hidden]="!maxHeapByPaddingSize.length">
-  <mat-divider ></mat-divider>
-</div>
-
 <max-heap-chart
   [maxHeap]="maxHeapByPaddingSize"
-  [title]="'By Padding Size'"
+  [title]="'by Padding Size'"
   [selectedIndex] = "selectedIndexByPaddingSize"
   (selected)="setSelectedHeapObjectByPaddingSize($event)">
 </max-heap-chart>

diff --git a/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.scss b/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.scss
@@ -30,3 +30,7 @@
 .control {
   padding: 20;
 }
+
+.heap-chart-header {
+  margin: 20px;
+}
diff --git a/frontend/app/components/memory_viewer/program_order_chart/BUILD b/frontend/app/components/memory_viewer/program_order_chart/BUILD
@@ -17,6 +17,9 @@ xprof_ng_module(
         "@npm//@angular/core",
         "@npm//@types/google.visualization",
         "@org_xprof//frontend/app/common/interfaces",
+        "@org_xprof//frontend/app/common/interfaces:chart",
+        "@org_xprof//frontend/app/components/chart",
+        "@org_xprof//frontend/app/components/chart:default_data_provider",
     ],
 )
 

diff --git a/frontend/app/components/memory_viewer/program_order_chart/program_order_chart.ng.html b/frontend/app/components/memory_viewer/program_order_chart/program_order_chart.ng.html
@@ -1,10 +1,11 @@
 <div [className]="heapSizes.length ? 'auto-height' : 'hidden-content'">
-  <div class="mat-headline">Working Space Size (MiB) vs Program Order (HLO Sequence)
-    <span><a [href]="timelineUrl" [hidden]="!timelineUrl.length" target="_blank">timeline</a></span>
+  <div class="chart-title-container">
+    <h2>Memory Allocation Size (MiB) vs Program Order (HLO Sequence)</h2>
+    <div class="timeline-link"><a [href]="timelineUrl" [hidden]="!timelineUrl.length" target="_blank">timeline</a></div>
   </div>
   <div class="chart-container">
-    <div class="chart" #peakChart></div>
-    <div class="chart" #chart></div>
-    <div class="chart" #activeChart></div>
+    <chart class="chart" #peakChart [chartType]="AREA_CHART" [dataInfo]="peakChartDataInfo"></chart>
+    <chart class="chart" #activeChart [chartType]="AREA_CHART" [dataInfo]="activeChartDataInfo"></chart>
+    <chart class="chart" #heapChart [chartType]="LINE_CHART" [dataInfo]="heapChartDataInfo"></chart>
   </div>
 </div>
diff --git a/frontend/app/components/memory_viewer/program_order_chart/program_order_chart.scss b/frontend/app/components/memory_viewer/program_order_chart/program_order_chart.scss
@@ -13,6 +13,16 @@
   height: 380px;
 }
 
+.chart-title-container {
+  display: flex;
+  align-items: center;
+}
+
+.timeline-link {
+  padding: 10px;
+  cursor: pointer;
+}
+
 .chart {
   position: absolute;
   top: 0;

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
@@ -1,6 +1,8 @@
 import {Component, ElementRef, Input, OnChanges, OnInit, SimpleChanges, ViewChild} from '@angular/core';
-
 import {BufferAllocationInfo} from 'org_xprof/frontend/app/common/interfaces/buffer_allocation_info';
+import {ChartDataInfo, ChartType} from 'org_xprof/frontend/app/common/interfaces/chart';
+import {SimpleDataTableOrNull} from 'org_xprof/frontend/app/common/interfaces/data_table';
+import {DefaultDataProvider} from 'org_xprof/frontend/app/components/chart/default_data_provider';
 
 /** A program order chart view component. */
 @Component({
@@ -22,20 +24,31 @@ export class ProgramOrderChart implements OnChanges, OnInit {
   @Input() activeInfo?: BufferAllocationInfo;
 
   /** Optional timeline URL. */
-  @Input() timelineUrl: string = '';
+  @Input() timelineUrl = '';
 
   @ViewChild('activeChart', {static: false}) activeChartRef!: ElementRef;
-  @ViewChild('chart', {static: false}) chartRef!: ElementRef;
-  @ViewChild('peakChart', {static: false}) peakChartRef!: ElementRef;
-
   activeChart: google.visualization.AreaChart|null = null;
-  chart: google.visualization.LineChart|null = null;
-  peakChart: google.visualization.AreaChart|null = null;
-  maxSize: number = 0;
-  maxOrder: number = 0;
+
+  maxSize = 0;
+  maxOrder = 0;
+  activeChartDataInfo: ChartDataInfo = {
+    data: null,
+    dataProvider: new DefaultDataProvider(),
+  };
+  peakChartDataInfo: ChartDataInfo = {
+    data: null,
+    dataProvider: new DefaultDataProvider(),
+  };
+  heapChartDataInfo: ChartDataInfo = {
+    data: null,
+    dataProvider: new DefaultDataProvider(),
+  };
+
+  readonly AREA_CHART = ChartType.AREA_CHART;
+  readonly LINE_CHART = ChartType.LINE_CHART;
 
   ngOnInit() {
-    this.loadGoogleChart();
+    this.updateCharts();
   }
 
   ngOnChanges(changes: SimpleChanges) {
@@ -51,32 +64,26 @@ export class ProgramOrderChart implements OnChanges, OnInit {
   }
 
   drawActiveChart() {
-    if (!this.activeChart) {
-      return;
-    }
-
     if (!this.activeInfo) {
-      this.activeChart.clearChart();
+      if (this.activeChart) {
+        this.activeChart.clearChart();
+      }
       return;
     }
 
     const dataTable = google.visualization.arrayToDataTable([
-      ['X', 'Size'],
+      ['Schedule', 'Size'],
       [this.activeInfo.alloc, this.activeInfo.size],
       [this.activeInfo.free, this.activeInfo.size],
     ]);
 
-    const options = {
+    const options: google.visualization.AreaChartOptions = {
       areaOpacity: 0.7,
       backgroundColor: 'transparent',
       chartArea: {
-        left: 50,
-        right: 50,
-        width: '90%',
-        height: '90%',
+        height: '80%',
       },
       colors: [this.activeInfo.color || ''],
-      focusTarget: 'none',
       hAxis: {
         baselineColor: 'transparent',
         gridlines: {color: 'transparent'},
@@ -99,13 +106,15 @@ export class ProgramOrderChart implements OnChanges, OnInit {
       lineWidth: 2,
     };
 
-    this.activeChart.draw(
-        dataTable, options as google.visualization.AreaChartOptions);
+    this.activeChartDataInfo = {
+      ...this.activeChartDataInfo,
+      data: dataTable.toJSON() as SimpleDataTableOrNull,
+      options,
+    };
   }
 
   drawChart() {
-    if (!this.chart || !this.heapSizes.length ||
-        !this.unpaddedHeapSizes.length) {
+    if (!this.heapSizes.length || !this.unpaddedHeapSizes.length) {
       return;
     }
 
@@ -120,28 +129,26 @@ export class ProgramOrderChart implements OnChanges, OnInit {
     this.maxSize = Math.round(this.maxSize * 1.1);
 
     const dataTable = new google.visualization.DataTable();
-    dataTable.addColumn('number', 'X');
-    dataTable.addColumn('number', 'Size');
-    dataTable.addColumn('number', 'Unpadded Size');
+    dataTable.addColumn('number', 'Schedule');
+    dataTable.addColumn('number', 'Heap Size');
+    dataTable.addColumn('number', 'Heap Size');
     dataTable.addRows(data);
 
-    const options = {
+    const options: google.visualization.LineChartOptions = {
       backgroundColor: 'transparent',
       chartArea: {
-        left: 50,
-        right: 50,
-        width: '90%',
-        height: '90%',
+        height: '80%',
       },
-      focusTarget: 'none',
       hAxis: {
+        title: 'Program Order',
         baselineColor: 'transparent',
         viewWindow: {
           min: 0,
           max: this.maxOrder,
         },
       },
       vAxis: {
+        title: 'Allocated Heap Size',
         baselineColor: 'transparent',
         viewWindow: {
           min: 0,
@@ -151,35 +158,43 @@ export class ProgramOrderChart implements OnChanges, OnInit {
       legend: {position: 'top'},
     };
 
-    this.chart.draw(
-        dataTable, options as google.visualization.LineChartOptions);
+    this.heapChartDataInfo = {
+      ...this.heapChartDataInfo,
+      data: dataTable.toJSON() as SimpleDataTableOrNull,
+      options,
+    };
   }
 
   drawPeakChart() {
-    if (!this.peakChart || !this.peakInfo) {
+    if (!this.peakInfo) {
       return;
     }
 
     const peakWidth = Math.max(Math.round(this.maxOrder / 50), 1);
     const peakAlloc =
         Math.max(Math.round(this.peakInfo.alloc - peakWidth / 2), 0);
     const peakFree = Math.min(peakAlloc + peakWidth, this.maxOrder);
-    const dataTable = google.visualization.arrayToDataTable([
-      ['X', 'Size'],
-      [peakAlloc, this.peakInfo.size],
-      [peakFree, this.peakInfo.size],
+    const dataTable = new google.visualization.DataTable();
+    dataTable.addColumn('number', 'Schedule');
+    dataTable.addColumn('number', 'Allocated Size');
+    dataTable.addColumn({type: 'string', role: 'tooltip'});
+    dataTable.addRows([
+      [
+        peakAlloc, this.peakInfo.size,
+        `peak memory allocation: ${this.peakInfo.size}`
+      ],
+      [
+        peakFree, this.peakInfo.size,
+        `peak memory allocation: ${this.peakInfo.size}`
+      ],
     ]);
 
-    const options = {
+    const options: google.visualization.AreaChartOptions = {
       backgroundColor: 'transparent',
       chartArea: {
-        left: 50,
-        right: 50,
-        width: '90%',
-        height: '90%',
+        height: '80%',
       },
       colors: ['#00ff00'],
-      focusTarget: 'none',
       hAxis: {
         baselineColor: 'transparent',
         gridlines: {color: 'transparent'},
@@ -202,28 +217,16 @@ export class ProgramOrderChart implements OnChanges, OnInit {
       lineWidth: 0,
     };
 
-    this.peakChart.draw(
-        dataTable, options as google.visualization.AreaChartOptions);
+    this.peakChartDataInfo = {
+      ...this.peakChartDataInfo,
+      data: dataTable.toJSON() as SimpleDataTableOrNull,
+      options,
+    };
   }
 
-  loadGoogleChart() {
-    if (!google || !google.charts) {
-      setTimeout(() => {
-        this.loadGoogleChart();
-      }, 100);
-    }
-
-    google.charts.safeLoad({'packages': ['corechart']});
-    google.charts.setOnLoadCallback(() => {
-      this.activeChart =
-          new google.visualization.AreaChart(this.activeChartRef.nativeElement);
-      this.chart =
-          new google.visualization.LineChart(this.chartRef.nativeElement);
-      this.peakChart =
-          new google.visualization.AreaChart(this.peakChartRef.nativeElement);
-      this.drawChart();
-      this.drawPeakChart();
-      this.drawActiveChart();
-    });
+  updateCharts() {
+    this.drawActiveChart();
+    this.drawPeakChart();
+    this.drawChart();
   }
 }
diff --git a/frontend/app/components/memory_viewer/program_order_chart/program_order_chart_module.ts b/frontend/app/components/memory_viewer/program_order_chart/program_order_chart_module.ts
@@ -1,7 +1,14 @@
 import {NgModule} from '@angular/core';
+import {ChartModule} from 'org_xprof/frontend/app/components/chart/chart';
 
 import {ProgramOrderChart} from './program_order_chart';
 
-@NgModule({declarations: [ProgramOrderChart], exports: [ProgramOrderChart]})
+@NgModule({
+  declarations: [ProgramOrderChart],
+  imports: [
+    ChartModule,
+  ],
+  exports: [ProgramOrderChart]
+})
 export class ProgramOrderChartModule {
 }