[iree-prof-tools] Make model-explorer interacting with editor. (#246)

1) Updated README.md with new mode-explorer package name.
2) "Explore Model: Focus" command moves focus on a node in the model
   explorer whose node label matches with the word at the current
   cursor's position in the editor.
3) When a node is focused in the model explorer, corresponding areas in
   the editor are highlighted.

Signed-off-by: Byungchul Kim <byungchul@google.com>

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

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": {

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<string, WebviewPanelForModelExplorer>();
+
+  // 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<string | undefined> {
   const fileName = vscode.window.activeTextEditor?.document.fileName;
   if (fileName) {

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<string, any>;
+  private byLabel: Map<string, any[]>;
+
+  constructor(graphCollection: any) {
+    this.byId = new Map<string, any>();
+    this.byLabel = new Map<string, any[]>();
+    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<NodeMap | undefined> {
+  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;
+}

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<string> {

iree-prof-tools/model-explorer-extension/src/modelExplorer.ts

@@ -1,22 +1,30 @@
 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;
 
 // 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,35 +34,51 @@ 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<string>('externalModelExplorerUrl') ?? '';
     if (externalUrl.length > 0) {
       this.panel.webview.html = getWebviewContent(externalUrl, modelFile);
       return;
     }
 
+    // If an internal model explorer has already been running, reuse it.
     internalModelExplorerPort = internalModelExplorerPort ?? getRandomPort();
     const modelExplorerUrl = `http://localhost:${internalModelExplorerPort}`;
     if (internalModelExplorerTerminal != null) {
       this.panel.webview.html = getWebviewContent(modelExplorerUrl, modelFile);
       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) {
   <title>Model Explorer</title>
 </head>
 <body>
-  <iframe src="${modelExplorerUrl}/?data=${encodedData}&renderer=webgl&show_open_in_new_tab=0"
+  <iframe id="model-explorer-iframe",
+          src="${modelExplorerUrl}/?data=${encodedData}&renderer=webgl&show_open_in_new_tab=0"
           style="position: absolute; top: 0; left: 0; width: 100%; height: 100%; border: none;">
   </iframe>
+  <script>
+    const vscode = acquireVsCodeApi();
+    const iframeWindow = document.getElementById('model-explorer-iframe').contentWindow;
+    window.addEventListener('message', event => {
+      const message = event.data;
+      console.log('Got message: ' + JSON.stringify(message, null, 2) + ' from ' + event.origin);
+      if (event.origin.startsWith('vscode-webview:')) {
+        iframeWindow.postMessage(message, '*');
+      } else {
+        vscode.postMessage(message);
+      }
+    });
+  </script>
 </body>
 </html>`;
 }