There are some tradeoffs, for example, given the type system is shared across relay and tir. We might in the future have things like
@tvm.script
class MixedModule:
MyStruct = ADT[Integer, Integer]
def relay_func(x: ty.Tensor):
return relay.call_tir_dest_passing(tir_func, x)
def tir_func(x: ty.handle)
So parsing might need to go beyond the function level. Additionally, we might also need to serialize the meta-data.
Conceptually, the basic unit in the transformation is an IRModule, and it would be nice to have a decorator on a module to support that(although it is also important to decorate a function as well)