i am very new to TVM , and trying to get relay IR graph as json format. i have created a simple add function to understand the work flow of tvm relay and looks like this
x = relay.var("x") y = tvm.relay.var("y") add_op = tvm.relay.add(x, y) add_function = relay.Function([x,y] , add_op) mod = tvm.IRModule.from_expr(add_function) viz = relay_viz.RelayVisualizer(mod) print(viz) viz.render("output.json") executor = tvm.relay.create_executor(kind="vm", mod=mod)
could someone pleae tell me how to get this graph in json format. my actual goal is to modify this relay graph nodes in some moanner and proceed to opt. and compilation normally. for that i need to visualize and get graph in json format first.
thanks
Hi @yogeesh
Welcome to join the community and thank you for the asking.
As far as I know, the graph you get from the relay.build
is in a json format, therefore, you can change the part you want from it. Hope the following example works for you.
x = relay.var("x")
y = tvm.relay.var("y")
add_op = tvm.relay.add(x, y)
f = relay.Function([x,y], add_op)
mod = tvm.IRModule()
mod["main"] = f
with tvm.transform.PassContext(opt_level=3):
json, lib, params = relay.build(mod, target="llvm")
Reference: https://github.com/apache/tvm/blob/main/python/tvm/relay/build_module.py#L132
Best regards,