diff --git a/iree-prof-tools/model-explorer-extension/README.md b/iree-prof-tools/model-explorer-extension/README.md index 0265f53..f9bdb84 100644 --- a/iree-prof-tools/model-explorer-extension/README.md +++ b/iree-prof-tools/model-explorer-extension/README.md @@ -5,17 +5,16 @@ model explorer within VS Code. ## How to run it -Model explorer must be installed before this extension get started. +[Model explorer](https://github.com/google-ai-edge/model-explorer) must be +installed before this extension get started. Please follow the +[instruction](https://github.com/google-ai-edge/model-explorer/wiki/1.-Installation) +to install it. -``` -pip install model-explorer -``` - -This extension is started by executing `Explore Model` command from the Command -Palette on a VS code active text editor of a graph json file. The extension -loads the file of the current focused editor to the model explorer web server. -If the model explorer is not running, it starts one with a random port number -in a terminal. +This extension is started by executing `Explore Model: Start` command from the +Command Palette on a VS code active text editor of a graph json file. The +extension loads the file of the current focused editor to the model explorer +web server. If the model explorer is not running, it starts one with a random +port number in a terminal. Currently, it supports Tensorflow saved_model.pb files, TFLite files, StableHLO MLIR files, IREE MLIR files, and graph json files. For IREE MLIR files, it @@ -32,6 +31,14 @@ setting. ![Model Explorer Settings](model-explorer-settings.png) +The model explorer interacts with the editor of the original model file. When a +node is clicked on the model explorer, corresponding areas in the original +model file are highlighted. + +From a position or a selected area in the oringal model file, executing +`Explore Model: Focus` command from the Command Palette moves the focus on a +corresponding node in the model explorer. + ## How to make changes into it To make changes of this extension, VS Code, Node.js and typescript compiler are diff --git a/iree-prof-tools/model-explorer-extension/package.json b/iree-prof-tools/model-explorer-extension/package.json index d44593d..98c4871 100644 --- a/iree-prof-tools/model-explorer-extension/package.json +++ b/iree-prof-tools/model-explorer-extension/package.json @@ -13,8 +13,12 @@ "contributes": { "commands": [ { - "command": "modelExplorer.show", - "title": "Explore Model" + "command": "modelExplorer.start", + "title": "Explore Model: Start" + }, + { + "command": "modelExplorer.focus", + "title": "Explore Model: Focus" } ], "configuration": { diff --git a/iree-prof-tools/model-explorer-extension/src/extension.ts b/iree-prof-tools/model-explorer-extension/src/extension.ts index 7a0d308..bc6bdbb 100644 --- a/iree-prof-tools/model-explorer-extension/src/extension.ts +++ b/iree-prof-tools/model-explorer-extension/src/extension.ts @@ -3,15 +3,23 @@ import {convertMlirToJsonIfNecessary} from './mlirUtil'; import {WebviewPanelForModelExplorer} from './modelExplorer'; export function activate(context: vscode.ExtensionContext) { + const modelToPanelMap = new Map(); + + // Command to start a model explorer associated to the active text editor. context.subscriptions.push( - vscode.commands.registerCommand('modelExplorer.show', async () => { + vscode.commands.registerCommand('modelExplorer.start', async () => { const modelFile = await getModelFileName(); if (!modelFile) { - vscode.window.showInformationMessage('Invalid model file path.'); + vscode.window.showErrorMessage('Invalid model file path.'); return; } - const panel = new WebviewPanelForModelExplorer(context); + const panel = new WebviewPanelForModelExplorer( + context, + vscode.window.activeTextEditor + ); + modelToPanelMap.set(modelFile, panel); + panel.addDisposeCallback(() => { modelToPanelMap.delete(modelFile); }); context.subscriptions.push(panel); const modelFileToLoad = await convertMlirToJsonIfNecessary(modelFile); @@ -24,8 +32,31 @@ export function activate(context: vscode.ExtensionContext) { panel.startModelExplorer(modelFileToLoad); }) ); + + // Command to focus on a node in the model explorer associated to the active + // text editor. The model explorer must be launched before with start command + // above. + context.subscriptions.push( + vscode.commands.registerCommand('modelExplorer.focus', () => { + const modelFile = vscode.window.activeTextEditor?.document.fileName; + if (!modelFile) { + return; + } + + const panel = modelToPanelMap.get(modelFile); + if (panel) { + panel.focusOnNode(); + } else { + vscode.window.showErrorMessage( + 'Model explorer is not running for ' + modelFile + ); + } + }) + ); } +// Gets the filename of active text editor. If no editors are active, show the +// file open dialog. async function getModelFileName(): Promise { const fileName = vscode.window.activeTextEditor?.document.fileName; if (fileName) { diff --git a/iree-prof-tools/model-explorer-extension/src/graphUtil.ts b/iree-prof-tools/model-explorer-extension/src/graphUtil.ts new file mode 100644 index 0000000..874530a --- /dev/null +++ b/iree-prof-tools/model-explorer-extension/src/graphUtil.ts @@ -0,0 +1,52 @@ +import * as vscode from 'vscode'; + +// Map either by node id or by node label to nodes. +// Note that node ids are unique while node labels are not, i.e. multiple nodes +// may have the same node labels. +// TODO(byungchul): Check if node IDs are unique within a graph collection or +// within a graph. +export class NodeMap { + private byId: Map; + private byLabel: Map; + + constructor(graphCollection: any) { + this.byId = new Map(); + this.byLabel = new Map(); + for (let graph of graphCollection.graphs) { + for (let node of graph.nodes) { + this.byId.set(node.id, node); + const nodes = this.byLabel.get(node.label); + if (nodes) { + nodes.push(node); + } else { + this.byLabel.set(node.label, [node]); + } + } + } + } + + // Gets a node or undefined if not found with node ID. + getNodeById(nodeId: string): any | undefined { + return this.byId.get(nodeId); + } + + // Gets the first node or undefined if not found with node label. + getFirstNodeByLabel(nodeLabel: string): any | undefined { + const nodes = this.byLabel.get(nodeLabel); + return nodes && nodes.length > 0 ? nodes[0] : undefined; + } +} + +// Builds NodeMap from a graph JSON file. +// Returns undefined if it fails to parse graph JSON. +export async function buildNodeMap( + modelFile: string +): Promise { + if (!modelFile.endsWith('.json')) { + return undefined; + } + + const blob = await vscode.workspace.fs.readFile(vscode.Uri.file(modelFile)); + const graphCollection = JSON.parse(blob.toString()); + return graphCollection ? new NodeMap(graphCollection) : undefined; +} diff --git a/iree-prof-tools/model-explorer-extension/src/mlirUtil.ts b/iree-prof-tools/model-explorer-extension/src/mlirUtil.ts index 7c513d3..37f96fa 100644 --- a/iree-prof-tools/model-explorer-extension/src/mlirUtil.ts +++ b/iree-prof-tools/model-explorer-extension/src/mlirUtil.ts @@ -1,5 +1,6 @@ import * as vscode from 'vscode'; +// Converts an IREE MLIR file to a graph JSON file with iree-vis program. export async function convertMlirToJsonIfNecessary( modelFile: string ): Promise { diff --git a/iree-prof-tools/model-explorer-extension/src/modelExplorer.ts b/iree-prof-tools/model-explorer-extension/src/modelExplorer.ts index 20b47de..3b08268 100644 --- a/iree-prof-tools/model-explorer-extension/src/modelExplorer.ts +++ b/iree-prof-tools/model-explorer-extension/src/modelExplorer.ts @@ -1,4 +1,5 @@ import * as vscode from 'vscode'; +import {NodeMap, buildNodeMap} from './graphUtil'; // Internal model explorer web server shared by multiple webview panels. var internalModelExplorerTerminal: vscode.Terminal | undefined = undefined; @@ -6,17 +7,24 @@ var internalModelExplorerTerminal: vscode.Terminal | undefined = undefined; // A random port number of internal model explorer web server. var internalModelExplorerPort: number | undefined = undefined; +// Webview panel to run a model explorer. export class WebviewPanelForModelExplorer { context: vscode.ExtensionContext; + editor: vscode.TextEditor | undefined; panel: vscode.WebviewPanel; disposeCallbacks: (() => void)[]; + nodeMap: NodeMap | undefined = undefined; - constructor(context: vscode.ExtensionContext) { + constructor( + context: vscode.ExtensionContext, + editor: vscode.TextEditor | undefined + ) { this.context = context; + this.editor = editor; this.panel = vscode.window.createWebviewPanel( 'modelExplorer', 'Model Explorer', - vscode.window.activeTextEditor?.viewColumn ?? vscode.ViewColumn.One, + vscode.ViewColumn.Beside, { // Webview options. enableScripts: true, retainContextWhenHidden: true @@ -26,20 +34,35 @@ export class WebviewPanelForModelExplorer { this.disposeCallbacks = []; this.panel.onDidDispose( - () => { for (let f of this.disposeCallbacks) { f(); }}, + () => { for (let f of this.disposeCallbacks) { f(); } }, null, - context.subscriptions); + context.subscriptions + ); } dispose() { this.panel.dispose(); } + // Adds a callback called when this webview panel is disposed. addDisposeCallback(f: () => void) { this.disposeCallbacks.push(f); } - startModelExplorer(modelFile: string) { + // Starts a model explorer within this webview panel. + async startModelExplorer(modelFile: string) { + // Set up a message channel to model explorer in webview for interaction + // with the editor. + this.nodeMap = await buildNodeMap(modelFile); + if (this.nodeMap) { + this.panel.webview.onDidReceiveMessage( + async message => { this.onMessage(message); }, + undefined, + this.context.subscriptions + ); + } + + // If model explorer is an external server, no need to wait for it ready. const config = vscode.workspace.getConfiguration('modelExplorer'); const externalUrl = config.get('externalModelExplorerUrl') ?? ''; if (externalUrl.length > 0) { @@ -47,6 +70,7 @@ export class WebviewPanelForModelExplorer { return; } + // If an internal model explorer has already been running, reuse it. internalModelExplorerPort = internalModelExplorerPort ?? getRandomPort(); const modelExplorerUrl = `http://localhost:${internalModelExplorerPort}`; if (internalModelExplorerTerminal != null) { @@ -54,7 +78,7 @@ export class WebviewPanelForModelExplorer { return; } - // No model explorer is available. Starts one. + // No model explorer is available. Start one. vscode.window.showInformationMessage( 'Starting a model explorer web server...' ); @@ -82,12 +106,87 @@ export class WebviewPanelForModelExplorer { this.addDisposeCallback(() => { clearTimeout(timeout); }); } + + // Called on messages from web explorer in the webview. + private async onMessage(message: any) { + console.debug( + 'Got message from model-explorer: ' + JSON.stringify(message, null, 2) + ); + if (message.cmd == 'model-explorer-node-selected') { + const node = this.nodeMap?.getNodeById(message.nodeId); + if (node) { + this.focusOnText(node.label); + } else { + console.error('Unknown node ID: ' + message.nodeId); + } + } + } + + // Sets focus on text matched with the node label of the current node of mode + // explorer in the webview. + private async focusOnText(nodeLabel: string | undefined) { + if (!nodeLabel || !this.editor) { + return; + } + + // Find all matches with node label. + // TODO(byungchul): Utilize source file location info if exists. + const document = this.editor.document; + const selections: vscode.Selection[] = []; + for (const match of document.getText().matchAll(RegExp(nodeLabel, 'g'))) { + const start = document.positionAt(match.index); + const end = document.positionAt(match.index + match[0].length); + selections.push(new vscode.Selection(start, end)); + } + + if (selections.length > 0) { + // Bring the focus on this editor. + this.editor = await vscode.window.showTextDocument( + this.editor.document, + this.editor.viewColumn + ); + this.editor.selections = selections; + this.editor.revealRange(selections[0], vscode.TextEditorRevealType.AtTop); + } + } + + // Sets focus on a node of model explorer in the webview. + focusOnNode() { + if (!this.editor || !this.nodeMap) { + return; + } + + // Find a word at the current cursor position. + const wordRange = this.editor.document.getWordRangeAtPosition( + this.editor.selection.active + ); + + const word = wordRange ? this.editor.document.getText(wordRange) : ''; + console.debug('Word at cursor = "' + word + '"'); + if (!word) { + return; + } + + const node = this.nodeMap.getFirstNodeByLabel(word); + if (node) { + this.panel.webview.postMessage({ + 'cmd': 'model-explorer-select-node-by-node-id', + 'nodeId': node.id + }); + // Bring the focus on the webview panel. + this.panel.reveal(); + } + } } +// Gets a random port for internal model explorer server. function getRandomPort(): number { return 30080 + Math.floor(Math.random() * 9900); } +// Gets the webview contents, i.e. HTML. +// It wraps model-explorer with iframe, and listens to message both from this +// extension and from the model explorer. function getWebviewContent(modelExplorerUrl: string, modelFile: string) { vscode.window.showInformationMessage(`Loading a model file, ${modelFile}...`); const encodedData = encodeURIComponent(`{"models":[{"url":"${modelFile}"}]}`); @@ -98,9 +197,23 @@ function getWebviewContent(modelExplorerUrl: string, modelFile: string) { Model Explorer - + `; }