Updated source mapper component to support more compiler passes.

No-op changes as the corresponding data service is not implemented yet.
But the

PiperOrigin-RevId: 825430176

Author: zzzaries

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

@@ -119,7 +119,8 @@
         [stackTrace]="stackTrace"
         [sessionId]="sessionId"
         [programId]="focusedOpProgramId"
-        [opName]="focusedOpName"/>
+        [opName]="focusedOpName"
+        [opCategory]="focusedOpCategory"/>
     </as-split-area>
     </as-split>
   </mat-sidenav-content>

frontend/app/components/op_profile/op_profile_base.ts

@@ -44,6 +44,7 @@ export class OpProfileBase implements OnDestroy, OnInit, OnChanges {
   stackTrace = '';
   focusedOpProgramId = '';
   focusedOpName = '';
+  focusedOpCategory = '';
   showStackTrace = false;
   useUncappedFlops = false;
 
@@ -91,6 +92,7 @@ export class OpProfileBase implements OnDestroy, OnInit, OnChanges {
     this.stackTrace = node?.xla?.sourceInfo?.stackFrame || '';
     this.focusedOpProgramId = node?.xla?.programId || '';
     this.focusedOpName = node?.name || '';
+    this.focusedOpCategory = node?.xla?.category || '';
   }
 
   ngOnChanges(changes: SimpleChanges) {

frontend/app/components/source_mapper/BUILD

@@ -16,6 +16,7 @@ xprof_ng_module(
     deps = [
         "@npm//@angular/core",
         "@npm//@angular/forms",
+        "@npm//@ngrx/store",
         "@npm//rxjs",
         "@org_xprof//frontend/app/common/angular:angular_material_expansion",
         "@org_xprof//frontend/app/common/angular:angular_material_form_field",
@@ -28,6 +29,7 @@ xprof_ng_module(
         "@org_xprof//frontend/app/components/stack_trace_snippet:message",
         "@org_xprof//frontend/app/services/data_service_v2:data_service_v2_interface",
         "@org_xprof//frontend/app/services/source_code_service:source_code_service_interface",
+        "@org_xprof//frontend/app/store",
     ],
 )
 

frontend/app/components/source_mapper/source_mapper.ng.html

@@ -9,16 +9,16 @@
       </mat-select>
     </mat-form-field>
     <stack-trace-snippet
-      [sourceFileAndLineNumber]="sourceFileAndLineNumber"
-      [stackTrace]="stackTrace"
+      [sourceFileAndLineNumber]="adaptedSourceFileAndLineNumber"
+      [stackTrace]="adaptedStackTrace"
       [sourceContextWindow]="sourceContextWindow"
     />
   </div>
 
   <div class="ir-text-container">
     <mat-form-field class="ir-type-select" (click)="$event.stopPropagation()" appearance="outline">
       <mat-label>Compiler Pass</mat-label>
-      <mat-select [(ngModel)]="selectedCompilerPass">
+      <mat-select [ngModel]="selectedCompilerPass" (ngModelChange)="onCompilerPassChange($event)">
         <mat-option *ngFor="let pass of compilerPasses" [value]="pass">
           {{pass}}
         </mat-option>

frontend/app/components/source_mapper/source_mapper.ts

@@ -1,13 +1,17 @@
 import {Component, inject, Input, OnChanges, OnDestroy, SimpleChanges} from '@angular/core';
+import {Store} from '@ngrx/store';
 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 {Address} from 'org_xprof/frontend/app/services/source_code_service/source_code_service_interface';
+import {getTagsState} from 'org_xprof/frontend/app/store/selectors';
 import {ReplaySubject} from 'rxjs';
 import {takeUntil} from 'rxjs/operators';
 
+const CUSTOM_CALL_CATEGORY = 'custom-call';
+
 enum CompilerPass {
   HLO_OPTIMIZED = 'HLO(optimized)',
-  MOSAIC_LLO = 'Mosaic llo',
+  MOSAIC_ORIGINAL = 'Mosaic original',
 }
 
 /**
@@ -27,37 +31,107 @@ export class SourceMapper implements OnDestroy, OnChanges {
   private readonly dataService: DataServiceV2Interface =
       inject(DATA_SERVICE_INTERFACE_TOKEN);
 
+  /**
+   * The source file and line number of the HLO op.
+   * This is used to find the source code snippet.
+   * Processed from xla source info.
+   */
   @Input() sourceFileAndLineNumber: string|undefined = undefined;
-  @Input() stackTrace: string|undefined = undefined;
   /**
-   * The number of lines to show around the stack frame.
+   * The stack trace of the HLO op.
+   * Processed from xla stack frame info.
    */
+  @Input() stackTrace: string|undefined = undefined;
+  // The number of lines to show around the stack frame.
   @Input() sourceContextWindow = 40;
   @Input() sessionId = '';
+  // The program id of the HLO op.
   @Input() programId = '';
+  // The name of the HLO op.
   @Input() opName = '';
+  // The category of the HLO op.
+  @Input() opCategory = '';
 
   sourceCodeSnippetAddresses: readonly Address[] = [];
   hloTextByProgramId = new Map<string, string>();
-  // TODO(yinzz): add CompilerPass.MOSAIC_LLO
-  compilerPasses = [CompilerPass.HLO_OPTIMIZED];
+  mosaicTextByKernelName = new Map<string, string>();
+  mosaicSourceFileAndLineNumberByKernelName = new Map<string, string>();
   selectedCompilerPass = CompilerPass.HLO_OPTIMIZED;
   sourceFileNames: string[] = [];
   selectedSourceFileName = '';
+  tags: string[] = [];
+
+  // TODO: create data service to fetch source info and ir text for selected
+  // kernel (matching with hlo op name) by using llo dump and utils.
+
+  constructor(private readonly store: Store<{}>) {
+    this.store.select(getTagsState)
+        .pipe(takeUntil(this.destroyed))
+        .subscribe((tags: string[]) => {
+          if (!tags || tags.length === 0) return;
+          this.tags = tags;
+        });
+  }
 
   ngOnChanges(changes: SimpleChanges) {
     if (changes['sessionId'] &&
         changes['sessionId'].currentValue !== this.sessionId) {
       this.hloTextByProgramId.clear();
+      this.mosaicTextByKernelName.clear();
+      this.mosaicSourceFileAndLineNumberByKernelName.clear();
     }
     if (changes['programId']) {
       this.maybeUpdateHloTextCache();
     }
+    if (changes['opName'] && changes['opName'].currentValue !== this.opName) {
+      this.maybeUpdateMosaicTextCache();
+      this.maybeUpdateMosaicSourceFileAndLineNumberCache();
+    }
     if (changes['sourceFileAndLineNumber'] || changes['stackTrace']) {
       this.parseSourceFileNames();
     }
   }
 
+  get adaptedStackTrace(): string {
+    switch (this.selectedCompilerPass) {
+      case CompilerPass.HLO_OPTIMIZED:
+        return this.stackTrace || '';
+      case CompilerPass.MOSAIC_ORIGINAL:
+        return this.getPallasKernelStackTrace();
+      default:
+        return '';
+    }
+  }
+
+  get adaptedSourceFileAndLineNumber(): string {
+    switch (this.selectedCompilerPass) {
+      case CompilerPass.HLO_OPTIMIZED:
+        return this.sourceFileAndLineNumber || '';
+      case CompilerPass.MOSAIC_ORIGINAL:
+        return this.pallasKernelSourceFileAndLineNumber;
+      default:
+        return '';
+    }
+  }
+
+  get compilerPasses(): string[] {
+    const basePasses = [CompilerPass.HLO_OPTIMIZED];
+    // llo_debug tag is identifier for llo debug proto captured.
+    // the proto contains llo bundles for kernels, and source info for mosaic
+    // passes.
+    if (this.hasLloDebugTag) {
+      // CustomCall op category is identifier for a pallas kernel.
+      if (this.isCustomCall) {
+        basePasses.push(CompilerPass.MOSAIC_ORIGINAL);
+      }
+    }
+    return basePasses;
+  }
+
+  get hasLloDebugTag() {
+    return this.tags.includes('llo_debug');
+  }
+
   get irTextLink(): string {
     return '';
   }
@@ -70,6 +144,8 @@ export class SourceMapper implements OnDestroy, OnChanges {
     switch (this.selectedCompilerPass) {
       case CompilerPass.HLO_OPTIMIZED:
         return this.hloTextByProgramId.get(this.programId) || '';
+      case CompilerPass.MOSAIC_ORIGINAL:
+        return this.mosaicTextByKernelName.get(this.opName) || '';
       default:
         return '';
     }
@@ -79,6 +155,20 @@ export class SourceMapper implements OnDestroy, OnChanges {
     return this.irText.split('\n');
   }
 
+  get isCustomCall() {
+    return this.opCategory.includes(CUSTOM_CALL_CATEGORY);
+  }
+
+  get pallasKernelSourceFileAndLineNumber() {
+    return this.mosaicSourceFileAndLineNumberByKernelName.get(this.opName) ||
+        '';
+  }
+
+  // Not implemented yet.
+  getPallasKernelStackTrace() {
+    return '';
+  }
+
   isFocusLine(line: string): boolean {
     switch (this.selectedCompilerPass) {
       case CompilerPass.HLO_OPTIMIZED:
@@ -92,7 +182,13 @@ export class SourceMapper implements OnDestroy, OnChanges {
     switch (this.selectedCompilerPass) {
       case CompilerPass.HLO_OPTIMIZED:
         return this.irTextLines.findIndex(
-                   (line) => line.includes(`${this.opName} =`)) ||
+                   (line: string) => line.includes(`${this.opName} =`)) ||
+            0;
+      case CompilerPass.MOSAIC_ORIGINAL:
+        // Assumptions: the MLIR text contains the key word kernel in the kernel
+        // definition line.
+        return this.irTextLines.findIndex(
+                   (line: string) => line.includes('kernel')) ||
             0;
       default:
         return 0;
@@ -150,6 +246,65 @@ export class SourceMapper implements OnDestroy, OnChanges {
         });
   }
 
+  maybeUpdateMosaicTextCache() {
+    if (!this.opName || this.sessionId === '' ||
+        this.selectedCompilerPass !== CompilerPass.MOSAIC_ORIGINAL) {
+      return;
+    }
+    const text = this.mosaicTextByKernelName.get(this.opName);
+    if (text) {
+      return;
+    }
+    this.dataService
+        .getCustomCallText(
+            this.sessionId,
+            '',
+            this.opName,
+            this.programId,
+            )
+        .pipe(takeUntil(this.destroyed))
+        .subscribe((data: string) => {
+          if (data) {
+            this.mosaicTextByKernelName.set(this.opName, data);
+          }
+        });
+  }
+
+  maybeUpdateMosaicSourceFileAndLineNumberCache() {
+    if (!this.opName || this.sessionId === '' ||
+        this.selectedCompilerPass !== CompilerPass.MOSAIC_ORIGINAL) {
+      return;
+    }
+    const sourceFileAndLineNumber =
+        this.mosaicSourceFileAndLineNumberByKernelName.get(this.opName);
+    if (sourceFileAndLineNumber) {
+      return;
+    }
+    this.dataService.getLloSourceInfo(this.sessionId, this.opName)
+        .pipe(takeUntil(this.destroyed))
+        .subscribe((sourceInfo) => {
+          if (sourceInfo) {
+            this.mosaicSourceFileAndLineNumberByKernelName.set(
+                this.opName, sourceInfo);
+          }
+        });
+  }
+
+  onCompilerPassChange(newCompilerPass: CompilerPass) {
+    this.selectedCompilerPass = newCompilerPass;
+    switch (this.selectedCompilerPass) {
+      case CompilerPass.HLO_OPTIMIZED:
+        this.maybeUpdateHloTextCache();
+        break;
+      case CompilerPass.MOSAIC_ORIGINAL:
+        this.maybeUpdateMosaicTextCache();
+        this.maybeUpdateMosaicSourceFileAndLineNumberCache();
+        break;
+      default:
+        break;
+    }
+  }
+
   ngOnDestroy(): void {
     this.destroyed.next();
     this.destroyed.complete();