From 00c5a79c32412995e0f3957e2f99e06b47b8b032 Mon Sep 17 00:00:00 2001 From: Bryan Massoth Date: Fri, 11 Jul 2025 14:04:23 -0700 Subject: [PATCH] Adding extractor for metadata embedded in hlo expression for debugging. PiperOrigin-RevId: 782102451 --- frontend/app/components/op_profile/op_details/BUILD | 1 + .../op_profile/op_details/op_details.ng.html | 4 ++++ .../components/op_profile/op_details/op_details.ts | 11 ++++++++++- .../op_profile/op_details/op_details_module.ts | 2 ++ plugin/xprof/protobuf/op_profile.proto | 3 +++ xprof/convert/op_profile_builder.cc | 8 ++++++++ 6 files changed, 28 insertions(+), 1 deletion(-) diff --git a/frontend/app/components/op_profile/op_details/BUILD b/frontend/app/components/op_profile/op_details/BUILD index 5b7d945da..cb78b5a80 100644 --- a/frontend/app/components/op_profile/op_details/BUILD +++ b/frontend/app/components/op_profile/op_details/BUILD @@ -14,6 +14,7 @@ xprof_ng_module( "op_details.ng.html", ], deps = [ + "//third_party/javascript/ngx_json_viewer", "@npm//@angular/common", "@npm//@angular/core", "@npm//@ngrx/store", diff --git a/frontend/app/components/op_profile/op_details/op_details.ng.html b/frontend/app/components/op_profile/op_details/op_details.ng.html index 2764f8161..9ba610ad4 100644 --- a/frontend/app/components/op_profile/op_details/op_details.ng.html +++ b/frontend/app/components/op_profile/op_details/op_details.ng.html @@ -119,6 +119,10 @@
Long Op Name:
{{provenance}} +
+
Xprof Kernel Metadata:
+ +
Total Time Sum:
{{rawTimeMs}} diff --git a/frontend/app/components/op_profile/op_details/op_details.ts b/frontend/app/components/op_profile/op_details/op_details.ts index 27b81dbf2..3181f060f 100644 --- a/frontend/app/components/op_profile/op_details/op_details.ts +++ b/frontend/app/components/op_profile/op_details/op_details.ts @@ -1,11 +1,11 @@ import {Component, EventEmitter, inject, Input, Output} from '@angular/core'; 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 {DATA_SERVICE_INTERFACE_TOKEN} from 'org_xprof/frontend/app/services/data_service_v2/data_service_v2_interface'; import {getActiveOpProfileNodeState, getCurrentRun, getOpProfileRootNode, getProfilingGeneralState, getSelectedOpNodeChainState} from 'org_xprof/frontend/app/store/selectors'; import {ProfilingGeneralState} from 'org_xprof/frontend/app/store/state'; +import {Node} from 'org_xprof/frontend/app/common/interfaces/op_profile.jsonpb_decls'; import {Observable, ReplaySubject} from 'rxjs'; import {takeUntil} from 'rxjs/operators'; @@ -51,6 +51,7 @@ export class OpDetails { .fill(''); programId = ''; expression = ''; + xprofKernelMetadata: unknown = null; provenance = ''; sourceFile = ''; sourceLine = -1; @@ -250,6 +251,14 @@ export class OpDetails { this.programId = this.node.xla?.programId || ''; this.expression = this.node.xla?.expression || ''; + if (!!this.node.xla?.xprofKernelMetadata) { + try { + this.xprofKernelMetadata = + JSON.parse(this.node.xla?.xprofKernelMetadata); + } catch (e) { + this.xprofKernelMetadata = this.node.xla?.xprofKernelMetadata; + } + } this.provenance = this.node.xla?.provenance || ''; this.sourceFile = this.node.xla?.sourceInfo?.fileName || ''; this.sourceLine = this.node.xla?.sourceInfo?.lineNumber || -1; diff --git a/frontend/app/components/op_profile/op_details/op_details_module.ts b/frontend/app/components/op_profile/op_details/op_details_module.ts index 185c3c080..58ddda4ca 100644 --- a/frontend/app/components/op_profile/op_details/op_details_module.ts +++ b/frontend/app/components/op_profile/op_details/op_details_module.ts @@ -4,6 +4,7 @@ import {MatButtonModule} from '@angular/material/button'; import {MatCardModule} from '@angular/material/card'; import {MatIconModule} from '@angular/material/icon'; import {MatTooltipModule} from '@angular/material/tooltip'; +import {NgxJsonViewerModule} from 'google3/third_party/javascript/ngx_json_viewer/src/ngx-json-viewer.module'; import {OpDetails} from './op_details'; @@ -16,6 +17,7 @@ import {OpDetails} from './op_details'; MatIconModule, MatButtonModule, MatTooltipModule, + NgxJsonViewerModule, ], exports: [OpDetails] }) diff --git a/plugin/xprof/protobuf/op_profile.proto b/plugin/xprof/protobuf/op_profile.proto index 5a6e1e777..43b71d83b 100644 --- a/plugin/xprof/protobuf/op_profile.proto +++ b/plugin/xprof/protobuf/op_profile.proto @@ -65,6 +65,9 @@ message Node { string semantics = 3; // What the dimension represents, e.g. "spatial". } } + // Xprof kernel metadata extracted from the HLO expression's frontend + // attributes. + optional string xprof_kernel_metadata = 10; } } diff --git a/xprof/convert/op_profile_builder.cc b/xprof/convert/op_profile_builder.cc index 4e39c7246..71dc655d6 100644 --- a/xprof/convert/op_profile_builder.cc +++ b/xprof/convert/op_profile_builder.cc @@ -52,6 +52,11 @@ void PopulateSymbolNode(const OpMetrics& op_metrics, Node* node) { Node::XLAInstruction& xla = *node->mutable_xla(); xla.set_program_id(op_metrics.hlo_module_id()); xla.set_expression(op_metrics.long_name()); + auto kernel_metadata = + tsl::profiler::ExtractXprofKernelMetadata(op_metrics.long_name()); + if (kernel_metadata.ok() && !kernel_metadata->empty()) { + xla.set_xprof_kernel_metadata(*kernel_metadata); + } xla.set_fingerprint(op_metrics.fingerprint()); xla.set_category(op_metrics.category()); xla.set_provenance(op_metrics.provenance()); @@ -101,6 +106,9 @@ void CopySymbolDetailsToDeduplicatedNode(Node* top_child_node, const Node::XLAInstruction& top_child_node_xla = top_child_node->xla(); xla.set_program_id(top_child_node_xla.program_id()); xla.set_expression(top_child_node_xla.expression()); + if (top_child_node_xla.has_xprof_kernel_metadata()) { + xla.set_xprof_kernel_metadata(top_child_node_xla.xprof_kernel_metadata()); + } xla.set_fingerprint(top_child_node_xla.fingerprint()); xla.set_category(top_child_node_xla.category()); if (IsFusion(top_child_node_xla.category())) return;