Export TIR to json

Development pre-RFC
6

@SebastianBoblestETAS I think this could be helpful for folks consuming TIR outside of TVM itself. A couple of thoughts:

  1. Right now a related problem is that there isn’t currently a way for tvm.relay.build to return the IRModules which contain TIR itself. You can see that in the BuildOutput struct. I believe here you tried to save the runtime.Module–these are distinct from the IRModule and don’t contain TIR. As a first cut, it would be great to add those to the BuildOutput (possibly behind an option) from where they are here. Note that tvm::build currently does return those IRModules and you can print TIR from them.
  2. As @tqchen mentioned, you might be able to simply call tvm.ir.save_json on the IRModule once we have that returned from tvm.relay.build. Would this accomplish what you want?
  3. I’m not sure whether the graph-level IRModule generated by AOT is included in lowered_funcs IRModule from my link earlier. You might need to check on that and return it from AOTExecutorCodegen.
  4. Since TIR is also used internally in settings ranging from post-scheduling to pre-codegen, there are a few different possible forms of TIR. For example, once LowerTVMBuiltin pass runs, TIR gets significantly more complex as the details of the calling convention are included. If we’re going to commit a standard TIR export, we’ll need to choose a suitable setting from which that’s done. I’d suggest the simplest setting (e.g. before any TIR transforms are run), but unsure if that provides the level of detail you need.
  5. We could also opt not to standardize and allow folks to configure the point in time at which the IRModule itself is extracted for returning or saved to JSON in the pipeline.

Anyway, there are a few different thoughts here; I’d be interested in your responses and perhaps we can come to a more specified design from there.