Add a graphType parameter to downloadHloProto and related frontend components to support requesting different HLO graphs.

frontend/app/common/constants/constants.ts

@@ -82,6 +82,11 @@ export const GRAPH_CONFIG_KEYS = [
  */
 export const GRAPH_TYPE_DEFAULT = 'xla';
 
+/**
+ * The graph type for viewing the original HLO graph.
+ */
+export const GRAPH_TYPE_ORIGINAL_HLO = 'original_hlo';
+
 /** The tools generated by HLO proto. */
 export const HLO_TOOLS =
     ['memory_viewer', 'graph_viewer'];

frontend/app/components/controls/download_hlo/download_hlo.ts

@@ -1,5 +1,6 @@
 import {Component, inject, Input, OnDestroy} from '@angular/core';
 import {ActivatedRoute} from '@angular/router';
+import {GRAPH_TYPE_DEFAULT} from 'org_xprof/frontend/app/common/constants/constants';
 import {FileExtensionType} from 'org_xprof/frontend/app/common/constants/enums';
 import {DATA_SERVICE_INTERFACE_TOKEN, DataServiceV2Interface} from 'org_xprof/frontend/app/services/data_service_v2/data_service_v2_interface';
 import {ReplaySubject} from 'rxjs';
@@ -32,6 +33,8 @@ export class DownloadHlo implements OnDestroy {
   @Input() moduleName: string = '';
   /** Includes metadata in the proto. */
   @Input() showMetadata: boolean = false;
+  /** The graph type to download. */
+  @Input() graphType: string = GRAPH_TYPE_DEFAULT;
 
   readonly downloadMenuItems = DOWNLOAD_HLO_PROTO_MENU_ITEMS;
   private readonly destroyed = new ReplaySubject<void>(1);
@@ -53,6 +56,7 @@ export class DownloadHlo implements OnDestroy {
     this.dataService
         .downloadHloProto(
             this.sessionId,
+            this.graphType,
             this.moduleName,
             type,
             this.showMetadata,

frontend/app/components/graph_viewer/graph_viewer.ng.html

@@ -32,6 +32,7 @@
     </div>
     <div class="hlo-text-view-container">
       <hlo-text-view
+        [graphType]="graphType"
         [sessionId]="sessionId"
         [moduleName]="selectedModule"
         [showMetadata]="showMetadata"

frontend/app/components/graph_viewer/graph_viewer.ts

@@ -4,7 +4,7 @@ import {ActivatedRoute, Params, Router} from '@angular/router';
 import {Store} from '@ngrx/store';
 import {Throbber} from 'org_xprof/frontend/app/common/classes/throbber';
 import {GRAPH_CENTER_NODE_COLOR, GRAPH_OP_COLORS} from 'org_xprof/frontend/app/common/constants/colors';
-import {DIAGNOSTICS_DEFAULT, GRAPH_CONFIG_KEYS, GRAPH_TYPE_DEFAULT, GRAPHVIZ_PAN_ZOOM_CONTROL} from 'org_xprof/frontend/app/common/constants/constants';
+import {DIAGNOSTICS_DEFAULT, GRAPH_CONFIG_KEYS, GRAPH_TYPE_DEFAULT, GRAPH_TYPE_ORIGINAL_HLO, GRAPHVIZ_PAN_ZOOM_CONTROL} from 'org_xprof/frontend/app/common/constants/constants';
 import {OpProfileProto} from 'org_xprof/frontend/app/common/interfaces/data_table';
 import {Diagnostics} from 'org_xprof/frontend/app/common/interfaces/diagnostics';
 import {GraphConfigInput, GraphTypeObject, GraphViewerQueryParams} from 'org_xprof/frontend/app/common/interfaces/graph_viewer';
@@ -68,6 +68,7 @@ export class GraphViewer implements OnDestroy {
   graphvizUri = '';
   graphTypes: GraphTypeObject[] = [
     {label: 'Hlo Graph', value: GRAPH_TYPE_DEFAULT},
+    {label: 'Original Hlo Graph', value: GRAPH_TYPE_ORIGINAL_HLO},
   ];
   loadingGraph = false;
   loadingModuleList = false;

frontend/app/components/graph_viewer/hlo_text_view/hlo_text_view.ts

@@ -1,5 +1,6 @@
 import {Component, inject, Input, OnDestroy} from '@angular/core';
 import {Throbber} from 'org_xprof/frontend/app/common/classes/throbber';
+import {GRAPH_TYPE_DEFAULT} from 'org_xprof/frontend/app/common/constants/constants';
 import {FileExtensionType} from 'org_xprof/frontend/app/common/constants/enums';
 import {DATA_SERVICE_INTERFACE_TOKEN, DataServiceV2Interface} from 'org_xprof/frontend/app/services/data_service_v2/data_service_v2_interface';
 import {ReplaySubject} from 'rxjs';
@@ -32,6 +33,8 @@ const TOGGLE_BUTTON_ITEMS: ToggleButtonItems[] = [
   styleUrls: ['./hlo_text_view.scss'],
 })
 export class HloTextView implements OnDestroy {
+  /** The graph type. */
+  @Input() graphType = GRAPH_TYPE_DEFAULT;
   /** The hlo module name. */
   @Input() moduleName = '';
   /** Includes metadata in the proto. */
@@ -50,9 +53,10 @@ export class HloTextView implements OnDestroy {
   loading = false;
   loadingMessage = '';
   downloadedTextParams = {
+    graphType: GRAPH_TYPE_DEFAULT,
     moduleName: '',
     showMetadata: false,
-    textType: '', // SHORT_TEXT | LONG_TEXT
+    textType: '',  // SHORT_TEXT | LONG_TEXT
   };
 
   downloadHloText(type: string) {
@@ -62,6 +66,7 @@ export class HloTextView implements OnDestroy {
     this.dataService
         .downloadHloProto(
             this.sessionId,
+            this.graphType,
             this.moduleName,
             type,
             this.showMetadata,
@@ -72,6 +77,7 @@ export class HloTextView implements OnDestroy {
           this.loading = false;
           this.hloText = data as string;
           this.downloadedTextParams = {
+            graphType: this.graphType,
             moduleName: this.moduleName,
             showMetadata: this.showMetadata,
             textType: type,
@@ -83,7 +89,10 @@ export class HloTextView implements OnDestroy {
     const metadataMessage = this.downloadedTextParams.showMetadata
       ? 'with metadata '
       : '';
-    return `(Loaded: hlo ${this.downloadedTextParams.textType} ${metadataMessage}for module ${this.downloadedTextParams.moduleName})`;
+    return `(Loaded: hlo ${this.downloadedTextParams.textType} ${
+        metadataMessage}for module ${
+        this.downloadedTextParams.moduleName}) of graph type ${
+        this.downloadedTextParams.graphType}`;
   }
 
   setLoadingMessage(type: string, showMetadata: boolean) {

frontend/app/services/data_service_v2/data_service_v2.ts

@@ -259,6 +259,7 @@ export class DataServiceV2 implements DataServiceV2Interface {
   // hlo proto map by program id.
   downloadHloProto(
       sessionId: string,
+      graphType: string,
       moduleName: string,
       type: string,
       showMetadata: boolean,
@@ -274,6 +275,7 @@ export class DataServiceV2 implements DataServiceV2Interface {
                        .set('run', sessionId)
                        .set('tag', tool)
                        .set('host', host)
+                       .set('graph_type', graphType)
                        .set('module_name', moduleName)
                        .set('type', type)
                        .set('show_metadata', String(showMetadata));

frontend/app/services/data_service_v2/data_service_v2_interface.ts

@@ -79,6 +79,7 @@ export interface DataServiceV2Interface {
 
   downloadHloProto(
       sessionId: string,
+      graphType: string,
       moduleName: string,
       type: string,
       showMetadata: boolean,