Skip to content

Commit

Permalink
Move the node finding function to op detail data class
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 571051870
  • Loading branch information
lionelfeng authored and copybara-github committed Oct 5, 2023
1 parent 929cc2c commit 3361e25
Showing 1 changed file with 26 additions and 0 deletions.
26 changes: 26 additions & 0 deletions frontend/app/components/op_profile/op_profile_data.ts
Original file line number Diff line number Diff line change
@@ -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. */
Expand Down Expand Up @@ -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;
}
}

0 comments on commit 3361e25

Please sign in to comment.