MLIR Tree Viewer

A VS Code extension that renders MLIR (.mlir / .ttadapter) files as an interactive, collapsible tree with dependency tracing and on-the-fly memory inspection.

Features

• Tree View — All operations rendered as a nested tree, with dialect color tags (linalg, memref, scf, tensor, func, arith, …).
• Dependency Tracing — Click an op in Dep Trace mode to dim everything outside its producer/consumer chain. Trace depth is configurable (1 / 2 / 3 / full).
• Pin Mode — Mark multiple ops simultaneously without dimming the rest. Useful for cross-referencing ops from different parts of the file.
• Mem Info Bar — Toggle the on-demand memory inspector that follows the selected op:
  • Static-shape buffers show bytes, shape, and element type (e.g. 4.0 KB (128x16 f16)).
  • Dynamic-shape buffers show rank and dynamic-dim count in orange italic.
  • Sticks to the top of the tree panel when scrolling large files.
• Minimap Navigator — Right-edge vertical bar highlights the current selection, producers, consumers, and pins. Click to jump.
• Switch to Text — One-click toggle to the built-in text editor without closing the file.

Installation

From VS Code Marketplace

Search for MLIR Tree Viewer in the Extensions view (Ctrl+Shift+X), or visit: https://marketplace.visualstudio.com/items?itemName=Mlir-Visualizer.mlir-tree-viewer

From VSIX

code --install-extension mlir-tree-viewer-1.0.1.vsix

Activating

Open any .mlir or .ttadapter file. The custom editor takes over automatically — no command palette step needed.

Usage

1. Open a .mlir or .ttadapter file.

2. Click any op in the tree to select it. The minimap updates with the selection.

3. Use the toolbar:

   Button	Effect
   Dep Trace	Toggle dependency-tracing mode. Click an op to highlight its full producer/consumer chain.
   Pin	Toggle pin mode. Click ops to mark them; click again to unpin.
   Mem Info	Toggle the memory info bar. When on, clicking an op shows its result type and size at the top.
   Switch to Text	Open the file in the default text editor.

4. The Trace dropdown next to Dep Trace sets how many levels of dependencies to follow.

Requirements

• VS Code 1.80.0 or later
• MLIR files with standard dialects (func, memref, tensor, linalg, scf, arith, bufferization, etc.)

Project Structure

extension/
├── package.json              # Extension manifest (name, publisher, contributes)
├── tsconfig.json             # TypeScript config
├── src/
│   ├── extension.ts          # Activation + custom editor registration
│   ├── customEditor.ts       # CustomTextEditorProvider (loads/parses the file)
│   └── webview/
│       └── generator.ts      # Inlined HTML/CSS/JS for the tree view webview
└── out/                      # Compiled output (generated)

The webview (HTML, CSS, JS) is fully embedded in generator.ts and served via acquireVsCodeApi — no external assets at runtime.

Development

# Install deps
npm install

# One-shot build
npm run compile

# Watch mode (auto-recompile on save)
npm run compile -- -w

# Run a debug instance
# Open this folder in VS Code, press F5
# → opens Extension Development Host with the dev build loaded
# → edit src/, save, hit Ctrl+R in the dev host to reload

# Package the VSIX
npm run package        # produces mlir-tree-viewer-<version>.vsix

# Publish to Marketplace
vsce login Mlir-Visualizer
vsce publish

Testing the analyzer logic

The MLIR parser and shape analyzer live in generator.ts alongside the webview code. To sanity-check them in pure Node without launching VS Code:

node -e "
const { getHtml } = require('./out/webview/generator');
const fs = require('fs');
const mlir = fs.readFileSync('../HSTU.mlir', 'utf-8');
const html = getHtml({ asWebviewUri: u => u, cspSource: '' }, null, mlir, 'HSTU.mlir');
const scripts = [...html.matchAll(/<script>([\s\S]*?)<\/script>/g)];
scripts.forEach(([_, code], i) => { new Function(code); console.log('script', i, 'OK'); });
"

License

MIT