diff --git a/frontend/app/components/op_profile/op_profile_data.ts b/frontend/app/components/op_profile/op_profile_data.ts index 557d3ebe..44e6762a 100644 --- a/frontend/app/components/op_profile/op_profile_data.ts +++ b/frontend/app/components/op_profile/op_profile_data.ts @@ -1,4 +1,5 @@ import {Node} from 'org_xprof/frontend/app/common/interfaces/op_profile.jsonpb_decls'; +import {ProfileOrNull} from 'org_xprof/frontend/app/common/interfaces/data_table'; import * as utils from 'org_xprof/frontend/app/common/utils/utils'; /** An op profile data class. */ @@ -36,4 +37,29 @@ export class OpProfileData { this.flopsUtilizationPercent = undefined; } } + + // TODO: Make this function stateful or move to util. + findNode(opProfile: ProfileOrNull, moduleName: string, nodeName: string): Node + |null|undefined { + if (!opProfile || !opProfile.byProgram) return null; + for (const topLevelNode of opProfile.byProgram.children!) { + // Find the program id from OpProfile by the selected XLA module. + if (topLevelNode.name === moduleName) { + const node = this.findNodeHelper(topLevelNode.children, nodeName); + if (node) return node; + } + } + return null; + } + + private findNodeHelper(children: Node[]|null|undefined, name: string): Node + |null|undefined { + if (!children) return null; + for (const node of children) { + if (node.name === name) return node; + const findChildren = this.findNodeHelper(node.children, name); + if (findChildren) return findChildren; + } + return null; + } }