XProf OpProfile tools using "per gpu" suffix rather than "per core" f…

…or metrics title for GPU profiling result

PiperOrigin-RevId: 673030685

frontend/app/components/op_profile/op_details/op_details.ng.html

@@ -74,19 +74,19 @@
           </div>
         </div>
         <div class="info" [hidden]="!flopsRate">
-          <div class="title">FLOP rate (per core):</div>
+          <div class="title">{{getTitleByDeviceType('FLOP rate', ':')}}</div>
           <code class="expression">{{flopsRate}}</code>
         </div>
         <div class="info" [hidden]="!bandwidths[memBwType.MEM_BW_TYPE_HBM_RW]">
-          <div class="title">HBM bandwidth (per core):</div>
+          <div class="title">{{getTitleByDeviceType('HBM bandwidth', ':')}}</div>
           <code class="expression">{{bandwidths[memBwType.MEM_BW_TYPE_HBM_RW]}}</code>
         </div>
         <div class="info" [hidden]="!bandwidths[memBwType.MEM_BW_TYPE_SRAM_RD]">
-          <div class="title">On-chip Read bandwidth (per core):</div>
+          <div class="title">{{getTitleByDeviceType('On-chip Read bandwidth', ':')}}</div>
           <code class="expression">{{bandwidths[memBwType.MEM_BW_TYPE_SRAM_RD]}}</code>
         </div>
         <div class="info" [hidden]="!bandwidths[memBwType.MEM_BW_TYPE_SRAM_WR]">
-          <div class="title">On-chip Write bandwidth (per core):</div>
+          <div class="title">{{getTitleByDeviceType('On-chip Write bandwidth', ':')}}</div>
           <code class="expression">{{bandwidths[memBwType.MEM_BW_TYPE_SRAM_WR]}}</code>
         </div>
         <div class="info" [hidden]="!expression">

frontend/app/components/op_profile/op_details/op_details.ts

@@ -3,7 +3,8 @@ import {Store} from '@ngrx/store';
 import {Node} from 'org_xprof/frontend/app/common/interfaces/op_profile.jsonpb_decls';
 import {NavigationEvent} from 'org_xprof/frontend/app/common/interfaces/navigation_event';
 import * as utils from 'org_xprof/frontend/app/common/utils/utils';
-import {getActiveOpProfileNodeState, getCurrentRun, getOpProfileRootNode, getSelectedOpNodeChainState} from 'org_xprof/frontend/app/store/selectors';
+import {getActiveOpProfileNodeState, getCurrentRun, getOpProfileRootNode, getProfilingGeneralState, getSelectedOpNodeChainState} from 'org_xprof/frontend/app/store/selectors';
+import {ProfilingGeneralState} from 'org_xprof/frontend/app/store/state';
 import {Observable, ReplaySubject} from 'rxjs';
 import {takeUntil} from 'rxjs/operators';
 
@@ -60,6 +61,8 @@ export class OpDetails {
   memBwType = utils.MemBwType;
   currentRun = '';
   showUtilizationWarning = false;
+  deviceType = 'TPU';
+
 
   constructor(
       private readonly store: Store<{}>,
@@ -81,13 +84,31 @@ export class OpDetails {
         .subscribe((node: Node|null) => {
           this.rootNode = node || undefined;
         });
+    this.store.select(getProfilingGeneralState)
+        .pipe(takeUntil(this.destroyed))
+        .subscribe((generalState: ProfilingGeneralState|null) => {
+          this.deviceType = (generalState && generalState.deviceType) ?
+              generalState.deviceType :
+              'TPU';
+        });
+
     this.currentRun$.subscribe(run => {
       if (run) {
         this.currentRun = run;
       }
     });
   }
 
+  getTitleByDeviceType(titlePrefix: string, titleSuffix: string) {
+    if (this.deviceType === 'GPU') {
+      return `${titlePrefix} (per gpu)${titleSuffix}`;
+    } else if (this.deviceType === 'TPU') {
+      return `${titlePrefix} (per core)${titleSuffix}`;
+    } else {
+      return `${titlePrefix}${titleSuffix}`;
+    }
+  }
+
   hasValidGraphViewerLink() {
     const aggregatedBy = this.selectedOpNodeChain[0];
     if (aggregatedBy === 'by_category' && this.moduleList.length > 1) {

frontend/app/components/op_profile/op_profile.ts

@@ -4,7 +4,7 @@ import {Store} from '@ngrx/store';
 import {OpProfileProto} from 'org_xprof/frontend/app/common/interfaces/data_table';
 import {NavigationEvent} from 'org_xprof/frontend/app/common/interfaces/navigation_event';
 import {DataService} from 'org_xprof/frontend/app/services/data_service/data_service';
-import {setLoadingStateAction, setOpProfileRootNodeAction} from 'org_xprof/frontend/app/store/actions';
+import {setLoadingStateAction, setOpProfileRootNodeAction, setProfilingDeviceTypeAction} from 'org_xprof/frontend/app/store/actions';
 import {ReplaySubject} from 'rxjs';
 import {takeUntil} from 'rxjs/operators';
 
@@ -47,7 +47,12 @@ export class OpProfile extends OpProfileBase implements OnDestroy {
               message: '',
             }
           }));
-          this.parseData(data as OpProfileProto | null);
+          if (data) {
+            const profileProtoData = data as OpProfileProto;
+            this.store.dispatch(setProfilingDeviceTypeAction(
+                {deviceType: profileProtoData.deviceType}));
+          }
+          this.parseData(data as (OpProfileProto | null));
           this.store.dispatch(
               setOpProfileRootNodeAction({rootNode: this.rootNode || null}));
         });

frontend/app/store/actions.ts

@@ -97,6 +97,12 @@ export const setCurrentRunAction: ActionCreatorAny = createAction(
     props<{currentRun: string}>(),
 );
 
+/** Action to set deviceType */
+export const setProfilingDeviceTypeAction: ActionCreatorAny = createAction(
+    '[App State] Set profiling device type',
+    props<{deviceType: string | null}>(),
+);
+
 /** Action to set run tools map */
 export const setRunToolsMapAction: ActionCreatorAny =
     createAction('[App State] set run - tools map state', props<RunToolsMap>());

frontend/app/store/reducers.ts

@@ -60,6 +60,18 @@ export const reducer: ActionReducer<AppState, Action> = createReducer(
           };
         },
         ),
+    on(
+        actions.setProfilingDeviceTypeAction,
+        (state: AppState, action: ActionCreatorAny) => {
+          return {
+            ...state,
+            profilingGeneralState: {
+              ...state.profilingGeneralState,
+              deviceType: action.deviceType,
+            }
+          };
+        },
+        ),
     on(
         actions.setActivePodViewerInfoAction,
         (state: AppState, action: ActionCreatorAny) => {

frontend/app/store/selectors.ts

@@ -34,6 +34,10 @@ export const getOpProfileRootNode: MemoizedSelectorAny = createSelector(
     getOpProfileState,
     (opProfileState: OpProfileState) => opProfileState.rootNode);
 
+/** Selector for getProfilingGeneralState */
+export const getProfilingGeneralState: MemoizedSelectorAny = createSelector(
+    appState, (appState: AppState) => appState.profilingGeneralState);
+
 /** Selector for PodViewerState */
 export const getPodViewerState: MemoizedSelectorAny =
     createSelector(appState, (appState: AppState) => appState.podViewerState);

frontend/app/store/state.ts

@@ -22,6 +22,11 @@ export interface OpProfileState {
   rootNode: ActiveOpProfileNodeState;
 }
 
+/** General State of the Profiling */
+export interface ProfilingGeneralState {
+  deviceType: string;
+}
+
 /** Type for active pod viewer info state */
 type ActivePodViewerInfoState = AllReduceOpInfo|ChannelInfo|PodStatsRecord|null;
 
@@ -93,6 +98,7 @@ export interface AppState {
   dataRequest: DataRequest;
   runToolsMap: RunToolsMap;
   currentRun: string;
+  profilingGeneralState: ProfilingGeneralState;
 }
 
 /** Initial state of active heap object */
@@ -113,6 +119,11 @@ export const INIT_OP_PROFILE_STATE: OpProfileState = {
   rootNode: INIT_ACTIVE_OP_PROFILE_NODE_STATE,
 };
 
+/** Initial general profiling state */
+export const INIT_PROFILING_GENERAL_STATE: ProfilingGeneralState = {
+  deviceType: 'TPU',
+};
+
 /** Initial state of active pod viewer info */
 const INIT_ACTIVE_POD_VIEWER_INFO_STATE: ActivePodViewerInfoState = null;
 
@@ -172,6 +183,7 @@ export const INIT_APP_STATE: AppState = {
   errorMessage: INIT_ERROR_MESSAGE_STATE,
   runToolsMap: INIT_RUN_TOOLS_MAP,
   currentRun: INIT_CURRENT_RUN,
+  profilingGeneralState: INIT_PROFILING_GENERAL_STATE,
 };
 
 /** Feature key for store */