Data visualizations and correlating data between PassInstrumentations

I’ve been prototyping a tool for visualizing how many ops of each dialect exist in a top level module throughout a compilation pipeline: MLIR Pipeline Visualizer (source: MLIR instrumentation, javascript visualization). This uses a PassInstrumentation to dump some statistics after each pass to a JSON file, then loads that JSON file into a stacked bar chart. It flattens the compilation DAG by disabling threading, but I could also see people wanting to view similar graph/node visualizations.

This seems useful for project documentation and general compiler inspection already: when I run it on some typical IREE compilation flows, it

  • lets me visualize IREE’s major compilation phases (import into linalg + flow, codegen, HAL, VM)
    • I imagine running this on other projects would be interesting too :smiley:
  • gives a sense for what work happens in each phase (import at the start shuffles through dialects quite a bit, codegen is more of a steady ramp up over many passes)
  • lets me look for interesting passes like DeduplicateExecutables or places where Canonicalize and SymbolDCE eliminate many ops

I’d really like to correlate this data with some IR (either local or entire module) showing the before/after for each pass, but I’m not sure what the best way to do that would be. The details make this tricky:

  • PassInstrumentations are added to PassManagers, and multiple PassManagers may exist within a single program. IRPrinterInstrumentation is free to dump to the console, but my new instrumentation needs some extra structure to produce a valid JSON file.

    • Rather than try to correlate between IRPrinterInstrumentation and my new instrumentation, I could print IR within my new instrumentation, but then I’d still like to produce a merged view with all passes across all pass managers for the visualization somehow.
  • runBeforePass and runAfterPass are two separate callbacks, and it looks like they can be nested, so tracking “this pass changed the IR from X to Y” would need some bookkeeping (and maybe trust that the Pass *pass, Operation *op arguments match?).

Any suggestions for correlating with IR? I’d also be up for landing this upstream in some form if others would find it useful and have opinions on where it could slot in.

1 Like

That is pretty cool, wanna talk about it next week at the public meeting? If you could put together a few slides and/or walk us through the implementation, then we could discuss the various possibilities with this?