Create Hlo Stats Tool

PiperOrigin-RevId: 690501074
tensorflow · Dec 2, 2024 · 202bb4f · 202bb4f
1 parent 94b92ef
commit 202bb4f
Show file tree

Hide file tree

Showing 13 changed files with 820 additions and 0 deletions.
diff --git a/frontend/app/components/hlo_stats/BUILD b/frontend/app/components/hlo_stats/BUILD
@@ -0,0 +1,49 @@
+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 = "hlo_stats",
+    srcs = [
+        "hlo_stats.ts",
+        "hlo_stats_module.ts",
+    ],
+    assets = [
+        ":hlo_stats_css",
+        "hlo_stats.ng.html",
+    ],
+    deps = [
+        "@npm//@angular/core",
+        "@npm//@angular/router",
+        "@npm//@ngrx/store",
+        "@npm//rxjs",
+        "@org_xprof//frontend/app/common/angular:angular_material_divider",
+        "@org_xprof//frontend/app/common/angular:angular_material_expansion",
+        "@org_xprof//frontend/app/common/angular:angular_material_select",
+        "@org_xprof//frontend/app/common/angular:angular_material_tooltip",
+        "@org_xprof//frontend/app/common/constants",
+        "@org_xprof//frontend/app/common/interfaces",
+        "@org_xprof//frontend/app/common/interfaces:chart",
+        "@org_xprof//frontend/app/common/utils",
+        "@org_xprof//frontend/app/components/chart",
+        "@org_xprof//frontend/app/components/chart:chart_options",
+        "@org_xprof//frontend/app/components/chart:data_processor",
+        "@org_xprof//frontend/app/components/chart:default_data_provider",
+        "@org_xprof//frontend/app/components/chart/dashboard",
+        "@org_xprof//frontend/app/components/controls/category_filter",
+        "@org_xprof//frontend/app/components/controls/string_filter",
+        "@org_xprof//frontend/app/components/framework_op_stats/flop_rate_chart",
+        "@org_xprof//frontend/app/services/data_service",
+        "@org_xprof//frontend/app/store",
+    ],
+)
+
+sass_binary(
+    name = "hlo_stats_css",
+    src = "hlo_stats.scss",
+    sourcemap = False,
+    deps = [
+        "@org_xprof//frontend/app/styles:common",
+    ],
+)
diff --git a/frontend/app/components/hlo_stats/hlo_stats.ng.html b/frontend/app/components/hlo_stats/hlo_stats.ng.html
@@ -0,0 +1,100 @@
+<div class="section-container">
+  <div class="row">
+    <div>
+      <div class="mat-headline">Notes</div>
+      <div class="header-content">
+        <div class="description">
+          (1) You can filter the table by HLO op category, HLO op name or Framework op name (eg. JAX op, TF op).
+        </div>
+        <div class="description">
+          (2) In the charts and table below, "IDLE" represents the portion of the total execution
+          time on device that is idle..
+        </div>
+        <div class="description">
+          (3) In the pie charts, the "Other" sector represents the sum of sectors that are too
+          small to be shown individually.
+        </div>
+      </div>
+    </div>
+    <div class="flex-space"></div>
+  </div>
+
+  <br>
+  <mat-expansion-panel
+    [expanded]="showChartSection"
+    (opened)="showChartSection = true"
+    (closed)="showChartSection = false">
+    <mat-expansion-panel-header>
+      <mat-panel-title>
+        Chart Section
+      </mat-panel-title>
+    </mat-expansion-panel-header>
+    <div class="row">
+      <div class="pie-chart">
+        <div class="mat-title">Time per HLO category</div>
+        <chart chartType="PieChart" [dataInfo]="dataInfoCategoryChart"></chart>
+      </div>
+      <div class="pie-chart">
+        <div class="mat-title">Time per HLO operation</div>
+        <chart chartType="PieChart" [dataInfo]="dataInfoOpChart"></chart>
+      </div>
+      <div class="pie-chart">
+        <div class="mat-title" matTooltip="Time spent on rematerialization in HLO or Framework Ops">Time spent on rematerialization</div>
+        <chart chartType="PieChart" [dataInfo]="dataInfoRematerializationChart"></chart>
+      </div>
+      <div class="pie-chart" *ngIf="communicationOps.size">
+        <div class="mat-title">Time per HLO by Replica Groups</div>
+        <mat-select [(value)]="selectedCommOp" (selectionChange)="updateOpReplicaGroupChart()">
+          <mat-option *ngFor="let op of communicationOps" [value]="op">
+          {{op}}
+          </mat-option>
+        </mat-select>
+        <chart chartType="PieChart" [dataInfo]="dataInfoOpReplicaGroupChart"></chart>
+      </div>
+    </div>
+
+    <div class="row">
+      <div class="pie-chart">
+        <div class="mat-title">Time spent on rematerialization per HLO category</div>
+        <chart chartType="PieChart" [dataInfo]="dataInfoRematerializationCategoryChart"></chart>
+      </div>
+      <div class="pie-chart">
+        <div class="mat-title">Time spent on outside compilation</div>
+        <chart chartType="PieChart" [dataInfo]="dataInfoOutsideCompilationChart"></chart>
+      </div>
+
+      <flop-rate-chart [data]="data" [xColumn]="flopRateChartXColumn"
+                       [yColumn]="flopRateChartYColumn" [opType]="opType">
+      </flop-rate-chart>
+    </div>
+  </mat-expansion-panel>
+</div>
+
+<div class="section-container">
+  <div class="row">
+    <category-filter [dataTable]="dataTable" column="category"
+                     all="All" (changed)="updateFilters($event)">
+    </category-filter>
+    <string-filter [dataTable]="dataTable" column="program_id"
+      [value]="programIdSelected"
+      all="All" (changed)="updateFilters($event)">
+    </string-filter>
+    <string-filter [dataTable]="dataTable" column="hlo_op_name"
+      [value]="hloOpNameSelected" [exactMatch]="true"
+      all="All" (changed)="updateFilters($event)">
+    </string-filter>
+    <string-filter [dataTable]="dataTable" column="hlo_op_expression"
+                   all="All" (changed)="updateFilters($event)">
+    </string-filter>
+    <string-filter [dataTable]="dataTable" column="tf_op_name"
+                   all="All" (changed)="updateFilters($event)">
+    </string-filter>
+  </div>
+</div>
+
+<div>
+  <chart chartType="Table"
+         class="table"
+         [dataInfo]="dataInfoForTable">
+  </chart>
+</div>
diff --git a/frontend/app/components/hlo_stats/hlo_stats.scss b/frontend/app/components/hlo_stats/hlo_stats.scss
@@ -0,0 +1,35 @@
+/** CSS for an hlo stats component. */
+
+@import 'frontend/app/styles/common';
+
+:host {
+  display: block;
+}
+
+.section-container {
+  padding: 20px;
+}
+
+.header-content {
+  font-style: italic;
+  margin-top: -16px;
+}
+
+.row {
+  display: flex;
+}
+
+.flex-space {
+  flex: 1;
+}
+
+.pie-chart {
+  height: fit-content;
+  margin: 8px 8px 8px 0;
+  width: 450px;
+}
+
+.table {
+  border: 1px solid $border-color;
+  margin: 0 16px;
+}