Create Roofline Model tool in Tensorboard Plugin Profiler

PiperOrigin-RevId: 700161829

frontend/app/common/constants/BUILD

@@ -15,6 +15,9 @@ ts_library(
     srcs = [
         "roofline_model_constants.ts",
     ],
+    deps = [
+        "@npm//@types/google.visualization",
+    ],
 )
 
 ts_library(

frontend/app/common/constants/roofline_model_constants.ts

@@ -53,7 +53,7 @@ export const SCATTER_CHART_OPTIONS = {
     // Ticks have to be explicitly defined for scaling axis evenly.
     ticks: [0, 0.001, 0.01, 0.1, 1, 10, 100, 1000, 10000, 100000, 1000000],
   },
-  legend: {position: 'right'},
+  legend: {position: 'right' as google.visualization.ChartLegendPosition},
   tooltip: {isHtml: true},
   // Be mindful that series is not specified here, otherwise the shallow copy in
   // components of this object could result in overwritten series styles

frontend/app/components/main_page/BUILD

@@ -38,6 +38,7 @@ xprof_ng_module(
         "@org_xprof//frontend/app/components/op_profile",
         "@org_xprof//frontend/app/components/overview_page",
         "@org_xprof//frontend/app/components/pod_viewer",
+        "@org_xprof//frontend/app/components/roofline_model",
         "@org_xprof//frontend/app/components/sidenav",
         "@org_xprof//frontend/app/components/tf_data_bottleneck_analysis",
         "@org_xprof//frontend/app/components/trace_viewer",

frontend/app/components/main_page/main_page_module.ts

@@ -29,6 +29,8 @@ import {OverviewPage} from 'org_xprof/frontend/app/components/overview_page/over
 import {OverviewPageModule} from 'org_xprof/frontend/app/components/overview_page/overview_page_module';
 import {PodViewer} from 'org_xprof/frontend/app/components/pod_viewer/pod_viewer';
 import {PodViewerModule} from 'org_xprof/frontend/app/components/pod_viewer/pod_viewer_module';
+import {RooflineModel} from 'org_xprof/frontend/app/components/roofline_model/roofline_model';
+import {RooflineModelModule} from 'org_xprof/frontend/app/components/roofline_model/roofline_model_module';
 import {SideNavModule} from 'org_xprof/frontend/app/components/sidenav/sidenav_module';
 import {TfDataBottleneckAnalysis} from 'org_xprof/frontend/app/components/tf_data_bottleneck_analysis/tf_data_bottleneck_analysis';
 import {TfDataBottleneckAnalysisModule} from 'org_xprof/frontend/app/components/tf_data_bottleneck_analysis/tf_data_bottleneck_analysis_module';
@@ -69,6 +71,8 @@ export const routes: Routes = [
   {path: 'inference_profile', component: InferenceProfile},
   {path: 'hlo_stats', component: HloStats},
   {path: 'hlo_stats^', component: HloStats},
+  {path: 'roofline_model', component: RooflineModel},
+  {path: 'roofline_model^', component: RooflineModel},
   {path: '**', component: EmptyPage},
 ];
 
@@ -96,6 +100,7 @@ export const routes: Routes = [
     FrameworkOpStatsAdapterModule,
     DcnCollectiveStatsModule,
     HloStatsModule,
+    RooflineModelModule,
     InferenceProfileModule,
     RouterModule.forRoot(routes),
   ],

frontend/app/components/roofline_model/BUILD

@@ -0,0 +1,40 @@
+load("@io_bazel_rules_sass//:defs.bzl", "sass_binary")
+load("//defs:defs.bzl", "xprof_ng_module")
+
+package(default_visibility = ["//frontend:internal"])
+
+xprof_ng_module(
+    name = "roofline_model",
+    srcs = [
+        "roofline_model.ts",
+        "roofline_model_module.ts",
+    ],
+    assets = [
+        ":roofline_model_css",
+        "roofline_model.ng.html",
+    ],
+    deps = [
+        "@npm//@angular/core",
+        "@npm//@angular/router",
+        "@npm//@ngrx/store",
+        "@npm//@types/google.visualization",
+        "@npm//rxjs",
+        "@org_xprof//frontend/app/common/constants:roofline_model_constants",
+        "@org_xprof//frontend/app/common/interfaces",
+        "@org_xprof//frontend/app/common/utils",
+        "@org_xprof//frontend/app/components/chart/table",
+        "@org_xprof//frontend/app/components/controls/category_filter",
+        "@org_xprof//frontend/app/components/controls/string_filter",
+        "@org_xprof//frontend/app/components/roofline_model/operation_level_analysis",
+        "@org_xprof//frontend/app/components/roofline_model/program_level_analysis",
+        "@org_xprof//frontend/app/services/data_service",
+        "@org_xprof//frontend/app/store",
+    ],
+)
+
+sass_binary(
+    name = "roofline_model_css",
+    src = "roofline_model.scss",
+    # stack = True,
+    sourcemap = False,
+)

frontend/app/components/roofline_model/operation_level_analysis/operation_level_analysis.ts

@@ -64,7 +64,7 @@ export class OperationLevelAnalysis
       title: 'Percentage of self time per HLO op category',
       colors: PIE_CHART_PALETTE,
       sliceVisibilityThreshold: 0.01,
-    },
+    } as google.visualization.PieChartOptions,
   };
   dataInfoRooflineScatterChart: ChartDataInfo = {
     data: null,

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

@@ -0,0 +1,56 @@
+<div class="section-container">
+  <div class="block-content">
+    <div class="row">
+      <h2>Device Information</h2>
+      <div class="flex-space"></div>
+    </div>
+  </div>
+  <div class="block-content">
+    <div *ngFor="let info of deviceInfoArray" style="font-size:16px">
+      <div *ngIf="info.display"><b>{{info.label}}:</b> <span> {{info.value}} {{info.unit}} {{info.context}}</span></div>
+    </div>
+  </div>
+</div>
+
+<div class="section-container">
+  <div class="block-content">
+    <h2>Section1: Program-Level Analysis</h2>
+  </div>
+  <div class="block-content">
+    <div class="description">
+      Note:<br>
+      (1) This section provides program-level analysis.<br>
+      (2) A tooltip with extra information will show up if you mouse over a point in the roofline chart.<br>
+      (3) You can choose whether to include infeed and outfeed ops for the analysis.<br>
+      (4) "Total" aggregates all operations throughout the entire profiling session. It includes incomplete steps.<br>
+      (5) "Total (HW)" is based on the hardware performance counters while the others are based on the XLA's cost analysis.
+      It is always calculated including infeed and outfeed ops regardless of the option.
+      The gap between "Total" and "Total (HW)" is due to hardware limitation (e.g., padding).<br>
+      (6) "Average" shows the average step information by aggregating the operations in the complete steps only.
+    </div>
+  </div>
+  <div class="block-content">
+    <program-level-analysis [rooflineModelData]="dataTableProgram" [viewColumns]="columnsIdxProgram" [rooflineSeriesData]="scatterDataProgram" [scatterChartOptions]="scatterChartOptionsProgram" (filterUpdated)="updateDataTableProgram($event)"></program-level-analysis>
+  </div>
+</div>
+
+<div class="section-container">
+  <div class="block-content">
+    <h2>Section2: Operation-Level Analysis</h2>
+  </div>
+  <div class="block-content">
+    <div class="description">
+      Note:<br>
+      (1) This section provides operation-level analysis.<br>
+      (2) A tooltip with extra information will show up if you mouse over a point in the roofline chart.<br>
+      (3) To avoid sluggishness, only the 1000 most time-consuming operations are shown.<br>
+      (4) You can choose whether to include infeed and outfeed ops for the analysis.<br>
+      (5) You can filter data by HLO category, bottleneck resource or HLO name.<br>
+      (6) "IDLE" represents the portion of the total execution time on device that is idle.<br>
+      (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">
+    <operation-level-analysis [selectedOp]="selectedOpName" [rooflineModelData]="dataTableOp" [viewColumns]="columnsIdxOp" [rooflineSeriesData]="scatterDataOp" [scatterChartOptions]="scatterChartOptionsOp" (filterUpdated)="updateDataTableOp($event)"></operation-level-analysis>
+  </div>
+</div>

frontend/app/components/roofline_model/roofline_model.scss

@@ -0,0 +1,39 @@
+.section-container {
+  margin: 20px 20px 0px;
+}
+
+.block-content {
+  padding: 5px;
+}
+
+.row {
+  display: flex;
+}
+
+.flex-space {
+  flex: 1;
+}
+
+.description {
+  font-size: 14px;
+}
+
+.tableHeaderCell {
+  word-wrap: break-word;
+  background-color: azure; //!to hide the scrolled-up text.
+}
+
+.tableTableCell {
+  word-break: break-all;
+}
+
+.opColumnClass {
+  max-height: 200px;
+  overflow-y: auto;
+}
+
+.errorMessage {
+  border: 2px solid;
+  background-color: #ffcccb;
+  color: red;
+}