How to get graph representation of Relay program?

Hi All,

I am wondering how can I get the graph representation of relay program ? I’d like to get the graph representation in some kind of data structure where I can modify the graph in a way that I want.

You get the graph as a return value from the relay.build() function along with library and params.

1 Like

Thank you @cv_ruddha, and may be I should pose my question more clearly.

I’d like to modify the Relay IR, and I believe Relay IR is represented as a Relay Module.

I used relay.analysis.post_order_visit() function to traverse the Relay IR, but I can not modify the Relay IR nodes and create a new Relay IR module. So, my questions is how can I traverse convert Relay IR to another Relay IR?

The graph you get from the relay.build is in a json format, you can make changes to the json string and then deploy it onto any hardware target. Once you make a change in the json string you are essentially changing the Relay IR.

1 Like

Thank you @cv_ruddha. truly appreciate the prompt response.

Q1) It looks like we can do all Relay Level IR optimizations on the JSON file. Then, why we need all the RELAY IR? Isn’t the Relay IR is supposed to be the medium to do optimizations? Everything I read so far says that we need RELAY IR because it allows us to do optimization. Please correct me and help me understand.

Q2) I think i’d still want to do get some kind of representation of Relay IR into some kind of graph so that I can do some optimizations. I am not sure if JSON file is the right place to do that optimizartion.

You will have to trace the functions being used to optimize relay IR. For example, common sub expression elimination is an inbuilt relay transform. See where and how it is implemented. I know that to create your own transform you need to register your function using the TVM_GLOBAL_REGISTER Macro. I haven’t worked much on building my own transforms yet so I am not sure how that works. If anyone else in the forum knows about it please attach it to this thread.

1 Like

Thank you @cv_ruddha. Yes, now I think we are on the same page.

I have looked at the examples and such as common sub expr elimination from src\relay\pass. I believe I can write my own pass following common sub expr elimination and register my pass, but I am looking into achieving this inside PYTHON.

Thus, how can I get a graph representation of RELAY program inside a python ?

Thank you again for the prompt response and sharing your opinion.

1 Like

Thank you @MarisaKirisame

I saw this python file, but I am looking for ways to get graph representation of Relay IR? How can I get a graph representation of Relay IR and make some modification on that Relay IR ?

Or is it not possible to get the graph representations of Relay IR?

the relay IR is a graph representation. you can print(x) to see the text format. call dir(x) to check it’s member whih can be access through dot. If you want the API for all IR representation, see the expr.py file.

1 Like

Thank you @MarisaKirisame. Super appreciate the response.

What I am trying to do is to get a REAL graph where Vertices are nodes, and also with edges which tells me dependenccies. Does Relay IR is DAG?

For example, let us analyze following relay program.

 # Create a simple Relay program.
    c = relay.const(2.0, "float32")
    x = relay.var("x", "float32")
    y = relay.add(c, c)

    func = relay.Function([x], y)
    mod = relay.Module({"main": func})

Here is what mod, print(mod), dir(mod) gives if in the python console. And as I can see, it does not give any DAG info of Relay IR?

>>>mod
IRModuleNode( {GlobalVar(main): FunctionNode([Var(x, ty=TensorType([], float32))], (nullptr), CallNode(Op(add), [Constant(2.0), Constant(2.0)], (nullptr), []), [], (nullptr))})

>>>print(mod)
v0.0.4
def @main(%x: float32) {
  add(2f, 2f)
}

>>>dir(mod)
['functions', 'global_type_var_map_', 'global_var_map_', 'type_definitions']

func is your RelayIR. It is a graph. you can recover func from Module via mod.functions.

1 Like

Thank you. I posed an example here showing what I am trying to achieve: How to do Relay to Relay transformation by traversing each node?

Using the “func” or module, I am still can not write a simple relay to relay program transformation after visiting each node.