-
Notifications
You must be signed in to change notification settings - Fork 45
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[iree-prof-tools] Run iree-vis for mlir asm files. (#245)
1) Refactor extension.ts into 2 ts files. 2) Convert mlir files to temporary graph json file when filename ends with ".mlir" and iree-vis path is set which is empty by default. 3) Remove temp files when webview is closed. 4) iree-vis returns non-zero on error. 5) Reset port number when model-explorer terminal is gone. Signed-off-by: Byungchul Kim <byungchul@google.com>
- Loading branch information
1 parent
fb982cd
commit fabe743
Showing
7 changed files
with
212 additions
and
83 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
98 changes: 27 additions & 71 deletions
98
iree-prof-tools/model-explorer-extension/src/extension.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,86 +1,42 @@ | ||
| import * as vscode from 'vscode'; | ||
| import {convertMlirToJsonIfNecessary} from './mlirUtil'; | ||
| import {WebviewPanelForModelExplorer} from './modelExplorer'; | ||
|
|
||
| export function activate(context: vscode.ExtensionContext) { | ||
| // A random port number of model explorer web server. | ||
| const port = 30080 + Math.floor(Math.random() * 9900); | ||
|
|
||
| // Internal model explorer web server shared by multiple webview panels. | ||
| var internalModelExplorerTerminal: vscode.Terminal|null = null | ||
|
|
||
| async function startModelExplorer() { | ||
| var modelFile = vscode.window.activeTextEditor?.document.fileName; | ||
| if (!modelFile) { | ||
| const openFiles = await vscode.window.showOpenDialog({ | ||
| canSelectFiles: true, | ||
| canSelectMany: false, | ||
| title: 'No active text editor or too large file. Open a model file' | ||
| }); | ||
| modelFile = openFiles?.length == 1 ? openFiles[0].fsPath : undefined; | ||
| context.subscriptions.push( | ||
| vscode.commands.registerCommand('modelExplorer.show', async () => { | ||
| const modelFile = await getModelFileName(); | ||
| if (!modelFile) { | ||
| vscode.window.showInformationMessage('Invalid model file path.'); | ||
| return; | ||
| } | ||
| } | ||
|
|
||
| const config = vscode.workspace.getConfiguration('modelExplorer'); | ||
| const externalUrl = config.get<string>('externalModelExplorerUrl') ?? ''; | ||
| const connectToExternalServer = externalUrl.length > 0; | ||
| const modelExplorerUrl = connectToExternalServer ? externalUrl : `http://localhost:${port}`; | ||
|
|
||
| const panel = vscode.window.createWebviewPanel( | ||
| 'modelExplorer', | ||
| 'Model Explorer', | ||
| vscode.window.activeTextEditor?.viewColumn ?? vscode.ViewColumn.One, | ||
| { // Webview options. | ||
| enableScripts: true | ||
| } | ||
| ); | ||
| if (connectToExternalServer || internalModelExplorerTerminal != null) { | ||
| panel.webview.html = getWebviewContent(modelExplorerUrl, modelFile); | ||
| return; | ||
| } | ||
| const panel = new WebviewPanelForModelExplorer(context); | ||
| context.subscriptions.push(panel); | ||
|
|
||
| // No model explorer is available. Starts one. | ||
| vscode.window.showInformationMessage('Starting a model explorer web server...'); | ||
| internalModelExplorerTerminal = | ||
| vscode.window.createTerminal( | ||
| 'modelExplorerWebServer', | ||
| config.get<string>('internalModelExplorerPath') ?? 'model-explorer', | ||
| ['--no_open_in_browser', `--port=${port}`] | ||
| ); | ||
| vscode.window.onDidCloseTerminal(terminal => { | ||
| if (terminal == internalModelExplorerTerminal) { | ||
| vscode.window.showInformationMessage('Model explorer web server is closed.'); | ||
| internalModelExplorerTerminal = null; | ||
| const modelFileToLoad = await convertMlirToJsonIfNecessary(modelFile); | ||
| if (modelFileToLoad != modelFile) { | ||
| panel.addDisposeCallback(() => { | ||
| vscode.workspace.fs.delete(vscode.Uri.file(modelFileToLoad)); | ||
| }); | ||
| } | ||
| }); | ||
| context.subscriptions.push(internalModelExplorerTerminal); | ||
|
|
||
| // Delay webview rendering to wait for model explorer ready. | ||
| const timeout = setTimeout(() => { | ||
| panel.webview.html = getWebviewContent(modelExplorerUrl, modelFile!!); | ||
| }, 2000); // 2000 is arbitrary. Need a more reliable way. | ||
| panel.startModelExplorer(modelFileToLoad); | ||
| }) | ||
| ); | ||
| } | ||
|
|
||
| panel.onDidDispose(() => { clearTimeout(timeout); }, null, context.subscriptions); | ||
| async function getModelFileName(): Promise<string | undefined> { | ||
| const fileName = vscode.window.activeTextEditor?.document.fileName; | ||
| if (fileName) { | ||
| return fileName; | ||
| } | ||
|
|
||
| context.subscriptions.push( | ||
| vscode.commands.registerCommand('modelExplorer.show', startModelExplorer)); | ||
| } | ||
| const openFiles = await vscode.window.showOpenDialog({ | ||
| canSelectFiles: true, | ||
| canSelectMany: false, | ||
| title: 'No active text editor or too large file. Open a model file' | ||
| }); | ||
|
|
||
| function getWebviewContent(modelExplorerUrl: string, modelFile: string) { | ||
| vscode.window.showInformationMessage(`Loading a model file, ${modelFile}...`); | ||
| const encodedData = encodeURIComponent(`{"models":[{"url":"${modelFile}"}]}`); | ||
| return `<!DOCTYPE html> | ||
| <html lang="en"> | ||
| <head> | ||
| <meta charset="UTF-8"> | ||
| <title>Model Explorer</title> | ||
| </head> | ||
| <body> | ||
| <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> | ||
| </body> | ||
| </html>`; | ||
| } | ||
| return openFiles?.length == 1 ? openFiles[0].fsPath : undefined; | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| import * as vscode from 'vscode'; | ||
|
|
||
| export async function convertMlirToJsonIfNecessary( | ||
| modelFile: string | ||
| ): Promise<string> { | ||
| const config = vscode.workspace.getConfiguration('modelExplorer'); | ||
| const ireeVisPath = config.get<string>('ireeVisPath') ?? ''; | ||
| if (!modelFile.endsWith('.mlir') || ireeVisPath.length == 0) { | ||
| return modelFile; | ||
| } | ||
|
|
||
| vscode.window.showInformationMessage( | ||
| `Generating graph for a model file, ${modelFile}...` | ||
| ); | ||
| const graphJsonFile = modelFile + '.graph.json'; | ||
| const ireeVisTerminal = vscode.window.createTerminal({ | ||
| 'name': 'ireeVisRunner', | ||
| 'shellPath': ireeVisPath, | ||
| 'shellArgs': [ | ||
| `--input_iree_file=${modelFile}`, | ||
| `--output_json_file=${graphJsonFile}` | ||
| ], | ||
| 'hideFromUser': true | ||
| }); | ||
|
|
||
| return new Promise<string>((resolve, reject) => { | ||
| const token = vscode.window.onDidCloseTerminal(terminal => { | ||
| if (terminal == ireeVisTerminal) { | ||
| token.dispose(); | ||
| if (terminal.exitStatus?.code == 0) { | ||
| vscode.window.showInformationMessage( | ||
| `Succeeded to generate a graph file, ${graphJsonFile}.` | ||
| ); | ||
| resolve(graphJsonFile); | ||
| } else { | ||
| vscode.window.showErrorMessage('Failed to generate a graph file.'); | ||
| resolve(modelFile); | ||
| } | ||
| } | ||
| }); | ||
| }); | ||
| } |
106 changes: 106 additions & 0 deletions
106
iree-prof-tools/model-explorer-extension/src/modelExplorer.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,106 @@ | ||
| import * as vscode from 'vscode'; | ||
|
|
||
| // 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; | ||
|
|
||
| export class WebviewPanelForModelExplorer { | ||
| context: vscode.ExtensionContext; | ||
| panel: vscode.WebviewPanel; | ||
| disposeCallbacks: (() => void)[]; | ||
|
|
||
| constructor(context: vscode.ExtensionContext) { | ||
| this.context = context; | ||
| this.panel = vscode.window.createWebviewPanel( | ||
| 'modelExplorer', | ||
| 'Model Explorer', | ||
| vscode.window.activeTextEditor?.viewColumn ?? vscode.ViewColumn.One, | ||
| { // Webview options. | ||
| enableScripts: true, | ||
| retainContextWhenHidden: true | ||
| } | ||
| ); | ||
|
|
||
| this.disposeCallbacks = []; | ||
|
|
||
| this.panel.onDidDispose( | ||
| () => { for (let f of this.disposeCallbacks) { f(); }}, | ||
| null, | ||
| context.subscriptions); | ||
| } | ||
|
|
||
| dispose() { | ||
| this.panel.dispose(); | ||
| } | ||
|
|
||
| addDisposeCallback(f: () => void) { | ||
| this.disposeCallbacks.push(f); | ||
| } | ||
|
|
||
| startModelExplorer(modelFile: string) { | ||
| const config = vscode.workspace.getConfiguration('modelExplorer'); | ||
| const externalUrl = config.get<string>('externalModelExplorerUrl') ?? ''; | ||
| if (externalUrl.length > 0) { | ||
| this.panel.webview.html = getWebviewContent(externalUrl, modelFile); | ||
| return; | ||
| } | ||
|
|
||
| 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. | ||
| vscode.window.showInformationMessage( | ||
| 'Starting a model explorer web server...' | ||
| ); | ||
| internalModelExplorerTerminal = vscode.window.createTerminal( | ||
| 'modelExplorerWebServer', | ||
| config.get<string>('internalModelExplorerPath') ?? 'model-explorer', | ||
| ['--no_open_in_browser', `--port=${internalModelExplorerPort}`] | ||
| ); | ||
| const token = vscode.window.onDidCloseTerminal(terminal => { | ||
| if (terminal == internalModelExplorerTerminal) { | ||
| token.dispose(); | ||
| vscode.window.showInformationMessage( | ||
| 'Model explorer web server is closed.' | ||
| ); | ||
| internalModelExplorerTerminal = undefined; | ||
| internalModelExplorerPort = undefined; | ||
| } | ||
| }); | ||
| this.context.subscriptions.push(internalModelExplorerTerminal); | ||
|
|
||
| // Delay webview rendering to wait for model explorer ready. | ||
| const timeout = setTimeout(() => { | ||
| this.panel.webview.html = getWebviewContent(modelExplorerUrl, modelFile); | ||
| }, 2000); // 2000 is arbitrary. Need a more reliable way. | ||
|
|
||
| this.addDisposeCallback(() => { clearTimeout(timeout); }); | ||
| } | ||
| } | ||
|
|
||
| function getRandomPort(): number { | ||
| return 30080 + Math.floor(Math.random() * 9900); | ||
| } | ||
|
|
||
| function getWebviewContent(modelExplorerUrl: string, modelFile: string) { | ||
| vscode.window.showInformationMessage(`Loading a model file, ${modelFile}...`); | ||
| const encodedData = encodeURIComponent(`{"models":[{"url":"${modelFile}"}]}`); | ||
| return `<!DOCTYPE html> | ||
| <html lang="en"> | ||
| <head> | ||
| <meta charset="UTF-8"> | ||
| <title>Model Explorer</title> | ||
| </head> | ||
| <body> | ||
| <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> | ||
| </body> | ||
| </html>`; | ||
| } |