Add dvfs time scale multiplier processing logic for roofline model.

PiperOrigin-RevId: 805986876

Author: zzzaries

frontend/app/common/constants/roofline_model_constants.ts

@@ -12,7 +12,7 @@ export const SCATTER_CHART_AXIS = {
   minX: 0.00001,
   maxX: 100000,
   minY: 0.00001,
-  maxY: 1000000,
+  maxY: 10000000,
 };
 
 /** scatter base options for roofline chart */
@@ -191,6 +191,13 @@ export const DEVICE_INFO = [
     unit: 'Flop/byte',
     display: false,
   },
+  {
+    id: 'time_scale_multiplier',
+    label: 'Time Scale Multiplier',
+    type: 'number',
+    unit: '',
+    display: true,
+  },
 ];
 
 /**

frontend/app/components/roofline_model/BUILD

@@ -19,6 +19,10 @@ xprof_ng_module(
         "@npm//@ngrx/store",
         "@npm//@types/google.visualization",
         "@npm//rxjs",
+        "@org_xprof//frontend/app/common/angular:angular_material_icon",
+        "@org_xprof//frontend/app/common/angular:angular_material_progress_bar",
+        "@org_xprof//frontend/app/common/angular:angular_material_slide_toggle",
+        "@org_xprof//frontend/app/common/angular:angular_material_tooltip",
         "@org_xprof//frontend/app/common/classes",
         "@org_xprof//frontend/app/common/constants:roofline_model_constants",
         "@org_xprof//frontend/app/common/interfaces",
@@ -39,4 +43,5 @@ sass_binary(
     src = "roofline_model.scss",
     # stack = True,
     sourcemap = False,
+    deps = ["@org_xprof//frontend/app/styles:common"],
 )

frontend/app/components/roofline_model/roofline_model.ng.html

@@ -1,4 +1,4 @@
-<div class="section-container">
+<div class="section-container" style="position: sticky; top: 0; z-index: 1; background: white;">
   <div class="block-content">
     <div class="row">
       <h2>Device Information</h2>
@@ -10,6 +10,19 @@ <h2>Device Information</h2>
     <div *ngFor="let info of deviceInfoArray" style="font-size:16px">
       <div *ngIf="info.display"><b>{{info.label}}:</b> <span> {{deviceInfoText(info)}}</span></div>
     </div>
+    <!--tslint:disable-next-line:need higher angular version to deprecate ngIf-->
+    <div class="control" *ngIf="hasValidTimeScaleMultiplier()">
+      <div class="control-title">
+        Apply Scaling Factor
+        <mat-icon
+          class="tooltip-icon"
+          matTooltip="Show flops utilization scaled with the evaluated device capability."
+          matTooltipPosition="above">
+          info
+        </mat-icon>
+      </div>
+      <mat-slide-toggle [checked]="applyScalingFactor" (toggleChange)="toggleScalingFactor()"></mat-slide-toggle>
+    </div>
   </div>
 </div>
 
@@ -30,9 +43,13 @@ <h2>Section1: Program-Level Analysis</h2>
       (6) "Average" shows the average step information by aggregating the operations in the complete steps only.
     </div>
   </div>
-  <div class="block-content">
+  <div [hidden]="loadingAnalysis" class="block-content">
     <program-level-analysis #programLevelAnalysis [rooflineModelData]="dataTableProgram" [viewColumns]="columnsIdxProgram" [rooflineSeriesData]="scatterDataProgram" [scatterChartOptions]="scatterChartOptionsProgram" (filterUpdated)="updateDataTableProgram($event)"></program-level-analysis>
   </div>
+  <div class="progress-container" [hidden]="!loadingAnalysis">
+    <div class="loading-message">Loading analysis...</div>
+    <mat-progress-bar color="primary" mode="indeterminate" aria-label="refreshing program level roofline analysis"></mat-progress-bar>
+  </div>
 </div>
 
 <div class="section-container">
@@ -51,7 +68,11 @@ <h2>Section2: Operation-Level Analysis</h2>
       (7) Ops with zero FLOP (e.g., data formatting ops like reshape, IDLE, etc.) do not show up in the roofline chart.<br>
     </div>
   </div>
-  <div class="block-content">
+  <div [hidden]="loadingAnalysis" class="block-content">
     <operation-level-analysis #opLevelAnalysis [selectedOp]="selectedOpName" [rooflineModelData]="dataTableOp" [viewColumns]="columnsIdxOp" [rooflineSeriesData]="scatterDataOp" [scatterChartOptions]="scatterChartOptionsOp" (filterUpdated)="updateDataTableOp($event)"></operation-level-analysis>
   </div>
+  <div class="progress-container" [hidden]="!loadingAnalysis">
+    <div class="loading-message">Loading analysis...</div>
+    <mat-progress-bar color="primary" mode="indeterminate" aria-label="refreshing operation level roofline analysis"></mat-progress-bar>
+  </div>
 </div>

frontend/app/components/roofline_model/roofline_model.scss

@@ -1,3 +1,9 @@
+@import 'frontend/app/styles/common';
+
+:host {
+  display: block;
+}
+
 .section-container {
   margin: 20px 20px 0px;
 }
@@ -37,3 +43,21 @@
   background-color: #ffcccb;
   color: red;
 }
+
+.tooltip-icon {
+  transform: scale(0.8);
+}
+
+.control-title {
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+}
+
+.control {
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+  justify-content: left;
+  gap: 10px;
+}

frontend/app/components/roofline_model/roofline_model.ts

@@ -8,7 +8,7 @@ import {getGigaflopsReadableString, setLoadingState} from 'org_xprof/frontend/ap
 import {DATA_SERVICE_INTERFACE_TOKEN, DataServiceV2Interface} from 'org_xprof/frontend/app/services/data_service_v2/data_service_v2_interface';
 import {setCurrentToolStateAction} from 'org_xprof/frontend/app/store/actions';
 import {ReplaySubject} from 'rxjs';
-import {takeUntil} from 'rxjs/operators';
+import {take, takeUntil} from 'rxjs/operators';
 
 import {OperationLevelAnalysis} from './operation_level_analysis/operation_level_analysis';
 import {ProgramLevelAnalysis} from './program_level_analysis/program_level_analysis';
@@ -27,6 +27,7 @@ declare interface DeviceIndicators {
   hasCmem: boolean;
   hasMegacore: boolean;
   isGpu: boolean;
+  timeScaleMultiplier: number;
 }
 type ColumnIdxArr = Array<number|google.visualization.ColumnSpec>;
 
@@ -61,6 +62,9 @@ export class RooflineModel implements OnDestroy {
   @ViewChild('opLevelAnalysis') opLevelAnalysis?: OperationLevelAnalysis;
 
   host = '';
+  applyScalingFactor = false;
+  loadingAnalysis = false;
+
   // Device Information section data
   deviceInfoArray: DeviceInfoData[] = [];
   // Some critical indicators
@@ -69,6 +73,7 @@ export class RooflineModel implements OnDestroy {
     hasCmem: false,
     hasMegacore: false,
     isGpu: false,
+    timeScaleMultiplier: 1.0,
   };
 
   // dataTableRaw from the raw roofline model data
@@ -178,6 +183,20 @@ export class RooflineModel implements OnDestroy {
         });
   }
 
+  updateAnalysis() {
+    this.loadingAnalysis = true;
+    const params = new Map<string, string|boolean>();
+    if (this.applyScalingFactor) {
+      params.set('apply_time_scale_multiplier', this.applyScalingFactor);
+    }
+    this.dataService.getData(this.sessionId, this.tool, this.host, params)
+        .pipe(take(1))
+        .subscribe((data) => {
+          this.parseData(data as RooflineModelData[]);
+          this.loadingAnalysis = false;
+        });
+  }
+
   parseData(data?: RooflineModelData[]) {
     if (!google?.visualization) {
       console.log('gviz lib is not loaded yet.');
@@ -203,6 +222,11 @@ export class RooflineModel implements OnDestroy {
     this.processScatterDataOp();
   }
 
+  hasValidTimeScaleMultiplier(): boolean {
+    return this.deviceIndicators.timeScaleMultiplier > 0 &&
+        this.deviceIndicators.timeScaleMultiplier !== 1;
+  }
+
   /** parse the device information from the original dataset */
   parseDeviceInfoData(dataTableRaw: google.visualization.DataTable) {
     this.deviceIndicators = {
@@ -211,6 +235,8 @@ export class RooflineModel implements OnDestroy {
       hasMegacore: !!Number(dataTableRaw.getTableProperty('megacore')),
       isGpu: dataTableRaw.getTableProperty('device_type')
                  .startsWith(NVIDIA_GPU_TYPE_PREFIX),
+      timeScaleMultiplier:
+          Number(dataTableRaw.getTableProperty('time_scale_multiplier')) || 1,
     };
 
     this.deviceInfoArray = DEVICE_INFO.reduce(
@@ -250,13 +276,25 @@ export class RooflineModel implements OnDestroy {
             curInfo.context +=
                 '(if yes, the analysis assumes Megacore where an HLO runs on both TensorCores utilizing the full chip\'s resources so that the rooflines are twice higher)';
             curInfo.value = this.deviceIndicators.hasMegacore ? 'Yes' : 'No';
+          } else if (
+              cur.id === 'time_scale_multiplier' &&
+              !this.hasValidTimeScaleMultiplier()) {
+            curInfo.display = false;
           }
         }
-        const value = this.dataTableRaw!.getTableProperty(cur.id);
+        let value = this.dataTableRaw!.getTableProperty(cur.id);
+        value = cur.type === 'number' ? Number(value) : value;
+        if ([
+              'peak_flop_rate', 'peak_vmem_read_bw', 'peak_vmem_write_bw'
+            ].includes(cur.id)) {
+          curInfo.value = this.applyScalingFactor ?
+              (value * this.deviceIndicators.timeScaleMultiplier).toFixed(2) :
+              value;
+        }
         acc.push({
           // convert numeric value to numbers, as some ridge numbers will be
           // used as axis values in chart
-          value: cur.type === 'number' ? Number(value) : value,
+          value,
           // put cur at last to overwrite with preprocessed data
           ...curInfo,
         });
@@ -1185,6 +1223,11 @@ export class RooflineModel implements OnDestroy {
     }
   }
 
+  toggleScalingFactor() {
+    this.applyScalingFactor = !this.applyScalingFactor;
+    this.updateAnalysis();
+  }
+
   ngOnDestroy() {
     setLoadingState(false, this.store);
     this.destroyed.next();

frontend/app/components/roofline_model/roofline_model_module.ts

@@ -1,5 +1,9 @@
 import {CommonModule} from '@angular/common';
 import {NgModule} from '@angular/core';
+import {MatIconModule} from '@angular/material/icon';
+import {MatProgressBarModule} from '@angular/material/progress-bar';
+import {MatSlideToggleModule} from '@angular/material/slide-toggle';
+import {MatTooltipModule} from '@angular/material/tooltip';
 import {TableModule} from 'org_xprof/frontend/app/components/chart/table/table_module';
 import {CategoryFilterModule} from 'org_xprof/frontend/app/components/controls/category_filter/category_filter_module';
 import {ExportAsCsvModule} from 'org_xprof/frontend/app/components/controls/export_as_csv/export_as_csv_module';
@@ -20,6 +24,10 @@ import {RooflineModel} from './roofline_model';
     StringFilterModule,
     ProgramLevelAnalysisModule,
     OperationLevelAnalysisModule,
+    MatTooltipModule,
+    MatSlideToggleModule,
+    MatIconModule,
+    MatProgressBarModule,
   ],
   exports: [RooflineModel],
 })

frontend/app/services/data_service_v2/data_service_v2.ts

@@ -103,7 +103,7 @@ export class DataServiceV2 implements DataServiceV2Interface {
 
   private getHTTPParamsForDataQuery(
       run: string, tag: string, host: string,
-      parameters: Map<string, string> = new Map()): HttpParams {
+      parameters: Map<string, string|boolean> = new Map()): HttpParams {
     // Update searchparams with the updated run, tag and host.
     // In a Single Page App, we need to update the searchparams with the updated
     // run, tag and host on tool change for consistency.
@@ -130,7 +130,7 @@ export class DataServiceV2 implements DataServiceV2Interface {
 
   getData(
       sessionId: string, tool: string, host: string,
-      parameters: Map<string, string> = new Map()):
+      parameters: Map<string, string|boolean> = new Map()):
       Observable<DataTable|DataTable[]|null> {
     const params =
         this.getHTTPParamsForDataQuery(sessionId, tool, host, parameters);

frontend/app/services/data_service_v2/data_service_v2_interface.ts

@@ -25,7 +25,7 @@ export interface DataServiceV2Interface {
       sessionId: string,
       tool: string,
       host?: string,
-      parameters?: Map<string, string>,
+      parameters?: Map<string, string|boolean>,
       ignoreError?: boolean,
       ): Observable<DataTable|DataTable[]|null>;
 

plugin/xprof/protobuf/roofline_model.proto

@@ -26,6 +26,7 @@ enum RecordType {
 }
 
 // A database of RooflineModel records.
+// Next ID: 17
 message RooflineModelDatabase {
   // The device type.
   optional string device_type = 1;
@@ -63,11 +64,16 @@ message RooflineModelDatabase {
   // Error and warning messages for diagnosing profiling issues.
   optional tensorflow.profiler.Diagnostics diagnostics = 7;
 
+  // The weighted average time scale multiplier over all device op metrics.
+  // It is calculated by the total time of each op if running on default pstate
+  // divided by actual total op time.
+  optional double time_scale_multiplier = 16;
+
   reserved 3, 4, 6;
 }
 
 // There is one RooflineModelRecord for each HLO operation profiled.
-// Next ID: 44
+// Next ID: 45
 message RooflineModelRecord {
   // The record type.
   optional RecordType record_type = 18;
@@ -187,6 +193,10 @@ message RooflineModelRecord {
   // Whether the record is calculated including infeed and outfeed ops.
   optional bool include_infeed_outfeed = 26;
 
+  // Whether the device metrics (eg. flops utilization) is calculated with by
+  // applying the time scale multiplier.
+  optional bool apply_time_scale_multiplier = 44;
+
   // Flops for the record
   optional uint64 flops = 36;
 

xprof/convert/op_metrics_to_record.h

@@ -40,6 +40,16 @@ inline double GigaFlopsPerSecondPerCore(const OpMetrics& metrics) {
       metrics.flops(), tsl::profiler::PicoToNano(metrics.time_ps()));
 }
 
+// Normalized flop rate if running on default pstate.
+// Used to compare with default device peak flop rate to get utilization.
+inline double GigaFlopsPerSecondPerCoreNormalizedOnDvfs(
+    const OpMetrics& metrics) {
+  double time_scale_multiplier = tsl::profiler::SafeDivide(
+      metrics.normalized_time_ps() * 1.0, metrics.time_ps());
+  time_scale_multiplier = time_scale_multiplier ? time_scale_multiplier : 1.0;
+  return GigaFlopsPerSecondPerCore(metrics) * time_scale_multiplier;
+}
+
 inline double GigaModelFlopsPerSecondPerCore(const OpMetrics& metrics) {
   // flops and time_ps are accumulated across all occurrences on all cores.
   // time_ps is used instead of self_time_ps because flops for an op includes
@@ -157,7 +167,8 @@ static inline double GetMemoryPeakBandwidth(const PerfEnv& perf_env,
 
 template <typename Record>
 inline void SetRooflineMetrics(const OpMetrics& metrics, const PerfEnv perf_env,
-                               const RunEnvironment& run_env, Record* record) {
+                               const RunEnvironment& run_env, Record* record,
+                               bool apply_time_scale_factor = false) {
   using ::tensorflow::profiler::MemorySpace;
   using ::tensorflow::profiler::PerformanceInfo;
 
@@ -203,16 +214,19 @@ inline void SetRooflineMetrics(const OpMetrics& metrics, const PerfEnv perf_env,
     // access as HBM access.
     hbm_bytes = metrics.bytes_accessed();
   }
+  int64_t device_time_ps = apply_time_scale_factor
+                               ? metrics.normalized_time_ps()
+                               : metrics.time_ps();
   record->set_hbm_bw(tsl::profiler::GibibytesPerSecond(
       hbm_bytes, tsl::profiler::PicoToNano(metrics.time_ps())));
   record->set_cmem_read_bw(tsl::profiler::GibibytesPerSecond(
-      cmem_read_bytes, tsl::profiler::PicoToNano(metrics.time_ps())));
+      cmem_read_bytes, tsl::profiler::PicoToNano(device_time_ps)));
   record->set_cmem_write_bw(tsl::profiler::GibibytesPerSecond(
-      cmem_write_bytes, tsl::profiler::PicoToNano(metrics.time_ps())));
+      cmem_write_bytes, tsl::profiler::PicoToNano(device_time_ps)));
   record->set_vmem_read_bw(tsl::profiler::GibibytesPerSecond(
-      vmem_read_bytes, tsl::profiler::PicoToNano(metrics.time_ps())));
+      vmem_read_bytes, tsl::profiler::PicoToNano(device_time_ps)));
   record->set_vmem_write_bw(tsl::profiler::GibibytesPerSecond(
-      vmem_write_bytes, tsl::profiler::PicoToNano(metrics.time_ps())));
+      vmem_write_bytes, tsl::profiler::PicoToNano(device_time_ps)));
   record->set_hbm_operational_intensity(
       tsl::profiler::SafeDivide(metrics.flops(), hbm_bytes));
   record->set_cmem_read_operational_intensity(