[MLIR] MLIR Visualization Project

Hi, I am interested in working on the Open Project titled “MLIR Visualizations”.

MLIR allows for representing multiple levels of abstraction all together in the same IR/function. Visualizing MLIR modules therefore requires going beyond visualizing a graph of nodes all at the same level (which is not trivial in and of itself!), nor is it specific to Machine Learning. Beyond visualizing a MLIR module, there is also visualizing MLIR itself that is of interest. In particular, visualizing the rewrite rules, visualizing the matching process (including the failure to match, sort of like https://www.debuggex.com/ but for declarative rewrites), considering effects of rewrites over time, etc.

The visualizations should all be built with open source components but whether standalone (e.g., combining with, say, GraphViz to generate offline images) or dynamic tools (e.g., displayed in browser) is open for discussion. It should be usable completely offline in either case.

We will be working with interested students to refine the exact project based on interests given the wide scope of potential approaches. And open to proposals within this general area.

(Source: Open Projects - MLIR)

I would like to know the current status of this project and the work being done to achieve it.

I am currently going through the Toy Turorials’ chapters to get myself familiarized with the workings of MLIR.

Because of my interests and background, this project seemed to be a good place to start as I try to get acquainted with the LLVM community here. I worked on enhancing the Dask library’s visualization models as a part of the Google Summer of Code this summer.

@jpienaar

1 Like

Hey,

Welcome! Currently there is a Graphviz generator for showing one visualization (use-def chains) for MLIR modules, some work on going to make it more customizable and support some ability to filter what is shown. We’ve had some previous efforts in the general direction that was more focussed on syntax highlighting aspect. And conceptually the work recently presented at ODM showing visualizing op counts across passes is another instance of this topic (I don’t think that code is upstream yet, but could be mistaken).

Best,

Jacques

Hey!

If MLIR already supports the visualizations you talked about, it’s super cool. If there is any further/ongoing work in this domain, I would love to know more and contribute to it.

Where can I find this generator? Could you provide the link so that I can have a better idea of it? Can I be of any help here?

Additionally, I have some ideas for the following:

Visualizing the Rewrite Rules and Pattern Matching DAGs

  • Rewriting rules modifies the operating code. They have an input/output DAG pair that may be graphically represented. The useful feature here would be to highlight the differences (the removal of some DAG nodes) with a specific color or any other graph properties.
  • Different graph features could be used to show pattern matches (the similarity) between DAGs.

Graphviz is a wonderful tool for this as well, and it integrates nicely with the existing codebase.

It is in llvm-project/ViewOpGraph.cpp at main · llvm/llvm-project · GitHub primarily. It has mostly been worked on as/when there was an itch to scratch :slight_smile:

That would be very nice! I have some local version of a basic recursive dump (a la LLVM’s dumpr), which I find useful when looking at and trying to write patterns. That is orthogonal to what exists and would be useful.

We can coordinate on that a bit more, but the broad strokes sounds good and orthogonal to any other work going on here that I’m aware of.

1 Like

Related, GitHub - lutzroeder/netron: Visualizer for neural network, deep learning, and machine learning models is a nice visualizer that supports many common ML frameworks and formats. It also have an online app for opening models for visualization at https://netron.app/. It would be really cool if MLIR can be supported there too. But maybe not for all dialects and need to program javascript. :slight_smile: Just wanted to share the pointer, I’ll leave to @jpienaar to make the suggestions.

1 Like

Hey, I have come up with a little example to start the discussion.

Sample Program: Add two constant numbers 0 and 2.

Pre-optimization

func @main() -> i32 {
  %x = constant 0 : i32
  %y = constant 2 : i32
  %z = addi %x, %y : i32
  return %z : i32
}

Post-optimization

func @main() -> i32 {
  %z = constant 2 : i32
  return %z : i32
}

image

To compare the DAGs produced for these two modules, we can generate their respective DAGs and then utilize colors to compare them side by side (for eg: I used red to signify the past, and green to signify post optimizations).

The design choices can be discussed in detail.

Graphviz Source Code
digraph {
    node [shape=box fontname="Helvetica" style=filled fillcolor=white]
    edge [fontname="Helvetica"]
    
    /*
        ~ Unoptimized
        func @main() -> i32 {
          %x = constant 0 : i32
          %y = constant 2 : i32
          %z = addi %x, %y : i32
          return %z : i32
        }
    */
    
    constant1 [label="constant 0 : i32" fillcolor="#FF8484"]
    constant2 [label="constant 2 : i32" fillcolor="#FF8484"]
    addi [label="addi %x, %y : i32" fillcolor="#FF8484"]
    return1 [label="return %z : i32"]
    
    constant1 -> addi [label="%x"]
    constant2 -> addi [label="%y"]
    addi -> return1 [label="%z"]
    
    /*
        ~ Optimized
        func @main() -> i32 {
          %z = constant 2 : i32
          return %z : i32
        }
    */
    
    constant3 [label="constant 2 : i32" fillcolor="#84FF84"]
    return2 [label="return %z : i32"]
    
    constant3 -> return2 [label="%z"] 
}

Does this look good?

How would you like me to move forward in this direction?

1 Like

Would the intention be to just show what changed in general or show which ops led to which? E.g., if we had 2 such addi’s how would this appear, if we had 3 different patterns that all operated here before we got to the end state how would that be shown.

I think it would be nice to explore “pen” and paper as you have done here, and thinking about the rewrite drivers (greedy vs dialect conversion) and how these could be tracked in practice, that should show what is possible here or gaps.

I want to highlight the changes made by optimizers in a clean, intuitive, and robust way.

The goal is to share all the useful information of the operation textually/graphically.

I will spend some more time working out not-so-basic examples and get back to you in some days. I will get a much better idea of what the “Useful Information” is and how we can visualize them by playing with more DAGs.