Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Visualization of Relay IR #8668

Closed
wants to merge 22 commits into from
Closed

Visualization of Relay IR #8668

wants to merge 22 commits into from

Conversation

kueitang
Copy link
Contributor

@kueitang kueitang commented Aug 5, 2021

Hi,
This utility is designed for users who are prone to snap-shot the relay IR graph on terminal.
The AST structure format covers basic relay function/expressions such as Function, Call, Let, Var, GlobalVar, If, Tuple, Constant currently.

For example:
A Relay IRModule:

x = relay.var("x")
y = relay.const(1)
z = relay.add(x, y)
z = relay.multiply(x, z)
mod = tvm.ir.IRModule().from_expr(z)

with:

viz_res = retv.ASTVisualization()
mod = viz_res(mod)
res = viz_res.get_output()

will be showed as:

== The AST view of the IRModule is ==
@main([Var(x)])
   `--(call)
      |--multiply
      |--x
      `--(call)
         |--add
         |--x
         `--1

It is welcome to have your thoughts on the utility.

Anyways, thanks. :)

Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.

Copy link
Contributor

@leandron leandron left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @kueitang, this seems interesting as a way to see Relay.

I think this feature could be integrated with @chiwwang's ongoing work at #8448. Have you seen that PR?

I suggest you two work together to define the correct APIs needed, so that we have one way to implement visualizers for Relay using a common foundation. What do you think?

@kueitang
Copy link
Contributor Author

kueitang commented Aug 5, 2021

Thanks @leandron!
Yes, I have seen #8448. It uses an interface (Plottor) to seperate frontend and backend.
It is believed that the design uses only node and edge information to plot a graph.
As for this PR, there are more information derived from depth that I need to record while traversing.
It seems a little hard to unify the current implementation to that interface, or it needs to be refactored.

@kueitang kueitang force-pushed the ast-dump branch 3 times, most recently from b0bfea6 to fb24fd3 Compare August 5, 2021 13:18
@chiwwang
Copy link
Contributor

chiwwang commented Aug 6, 2021

Yes I consider graphs just as nodes and edges. So the Plotter interface is very simple.
I think depth, or "rank" can be obtained by a topological sort, but indeed maybe we should consider interfaces on different abstraction levels, for easier implementing other kinds of visualizers.
I only have a rough thoughts that we should allow customized visitors in a different interface.
Need some time to chew it over.

-Add a AST dump pass
-It provides a snap shot to the relay IR graph
@kueitang
Copy link
Contributor Author

hi @tqchen,
would you like to give it any thoughts? Thanks a lot! :)

@comaniac
Copy link
Contributor

If the high-level design or the interface needs more feedback, it would be better and more efficient to file a formal RFC at the same time.

https://github.com/apache/tvm-rfcs

@leandron
Copy link
Contributor

leandron commented Aug 10, 2021

If the high-level design or the interface needs more feedback, it would be better and more efficient to file a formal RFC at the same time.

https://github.com/apache/tvm-rfcs

Agree with @comaniac. With an RFC, perhaps we could also find alignment to have this and other features e.g. #8448, to be built on top of the same interface, and also other parts consuming this (I can imagine tvmc as a consumer) to also be able to integrate better with a single and organized interface.

@chiwwang
Copy link
Contributor

Hi @kueitang
Maybe we can re-use this RFC, https://discuss.tvm.apache.org/t/rfc-visualizing-relay-program-as-graph/4825
There are some valuable thoughts there :)

Back to interfaces of Relay visualization, how do you think about a backend-specific render-callback? Specifically, in the constructor of RelayVisualizer in #8448 :

class RelayVisualizer:
    """Relay IR Visualizer"""

    def __init__(
        self, relay_mod, relay_param=None, plotter_be=PlotterBackend.BOKEH, render_cb=_dft_render_cb
    ):

I try to port this PR with the interface in the following commit in the file _terminal.py.
chiwwang@4d1e14a
(I also make the callback coupled with the backend....but this seems not a good idea. A common render-callback can save efforts of implementing a new backend.)

The output looks like:

`@main([Var(x)])
  `--Call multiply
    |--Var(Input)
    `--Call add
      |--Var(Input)
      `--Constshape: (), dtype: int32`

I don't know if I should consider multiply/add as an argument of Call, so I put it in Call string.
This might be fixed by modifying tmv.ir.Op branch in the callback.

So, currently interfaces of relay-viz in my mind is:

  1. class Plotter and class Graph in plotter.py (one plotter can have lots of graphs)
  2. If needed, a backend-specific render-callback can overwrite the default one, with the prototype
def render_cb(graph, node_to_id, relay_param)
    """
    Parameters
    ----------
    graph : class plotter.Graph
    node_to_id : Dict[relay.expr, int]
    relay_param : Dict[string, NDarray]
    """

Additionally, the dict node_to_id ensure its order of keys is in the post-order of a relay GlobalVar.
How do you think this kind of integration?
Thanks :)

kueitang and others added 4 commits August 31, 2021 02:38
    * BOKEH BACKEND
        * Scalable Text for Node type
        * Add arrow-head and TODO
        * Use pydot. Remove networkx and pygraphviz dependency
        * Add interactive legend and show information based on zoom level
    * Support multiple GlobalVar. One global var, one graph
    * test reserved post order. try terminal viz by kueitang
…ntegrated two backends * Add a RenderCallback class for reusable nodes type processing
@kueitang
Copy link
Contributor Author

  1. Users can visualize relay by a single entry point, RelayVisualizer. They could choose different renderers by options(backend in RelayVisualizer constructor)
  2. For those who want to implement new renderer, it is expected to implement Graph, Plotter (in plotter.py). We provide default RenderCallback (render_callbak.py) for different relay ops to trigger Graph interface.
    Specific method for a certain op could be overrided if needed, like TermRenderCallback.

@chiwwang @leandron

@chiwwang
Copy link
Contributor

chiwwang commented Oct 25, 2021

Hi @areusch, @PhilippvK, @jroesch, @leandron,
I finish an implementation, with tests, docs, and a tutorial.
You might want to take a look. :)
In short, RelayVisualizer class serve as a single entry point for different renderer.
All it does is using interfaces defined in plotter.py and node_edge_gen.py.
A tutorial show a simple use-case and some possible customization.
Thanks!

@huajsj
Copy link
Contributor

huajsj commented Dec 2, 2021

Thanks @chiwwang.



def get_plotter_and_generator(backend):
"""Specify the Plottor and its NodeEdgeGenerator"""
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need "Parameters" and "return".

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.



class VizNode:
"""Node carry information used by `plotter.Graph` interface."""
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need parameters docstring

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.


class VizEdge:
"""Edges for `plotter.Graph` interface."""

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

def _function_node(
self,
node: relay.Expr,
_: Dict[str, tvm.runtime.NDArray], # relay_param
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove relay_param

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

@chiwwang
Copy link
Contributor

chiwwang commented Dec 6, 2021

Thanks @huajsj for reviewing!

@huajsj
Copy link
Contributor

huajsj commented Dec 15, 2021

@kueitang @chiwwang, do you think it is possible to split this PR into couple small patches? a smaller patch may help to accelerate the process and more review friendly.

@chiwwang
Copy link
Contributor

Sure, I will figure out how to split this PR to smaller patches.
I will raise new PRs and then close this one.
Thanks for suggestions @huajsj !

@jroesch jroesch added status: need review needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it labels Jan 19, 2022
chiwwang added a commit to chiwwang/tvm that referenced this pull request Jan 27, 2022
This PR follows apache#8668, with splitting
out interfaces class and terminal ast-dump implementation.

This visualizer is aimed for quick look-then-fix, so the interface is
simple. Despite that, customization is still possbile through
implementing interfaces defined in `interface.py` or overriding existent
implementations inside a renderer module, like `terminal.py`.

A tutorial is also provided in this PR.

A graphviz renderer will also be contributed after this PR.
@chiwwang
Copy link
Contributor

A follow-up PR: #10085
@kueitang please help to close this thread. Thanks!

@AndrewZhaoLuo
Copy link
Contributor

Closing in favor of #10085

masahi pushed a commit that referenced this pull request Feb 22, 2022
* RelayViz interface and terminal ast-dump.

This PR follows #8668, with splitting
out interfaces class and terminal ast-dump implementation.

This visualizer is aimed for quick look-then-fix, so the interface is
simple. Despite that, customization is still possbile through
implementing interfaces defined in `interface.py` or overriding existent
implementations inside a renderer module, like `terminal.py`.

A tutorial is also provided in this PR.

A graphviz renderer will also be contributed after this PR.

* lint and typo
pfk-beta pushed a commit to pfk-beta/tvm that referenced this pull request Apr 11, 2022
* RelayViz interface and terminal ast-dump.

This PR follows apache#8668, with splitting
out interfaces class and terminal ast-dump implementation.

This visualizer is aimed for quick look-then-fix, so the interface is
simple. Despite that, customization is still possbile through
implementing interfaces defined in `interface.py` or overriding existent
implementations inside a renderer module, like `terminal.py`.

A tutorial is also provided in this PR.

A graphviz renderer will also be contributed after this PR.

* lint and typo
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it status: need review
Projects
None yet
Development

Successfully merging this pull request may close these issues.