Hey @SebastianBoblestETAS sorry I missed your past message. in thinking about this a bit more, there are a couple of places this could be added.
Right now we generate a map ir_module_by_target in tvm.build. We export a similar thing in tvm.relay.build but it’s not present on all ExecutorFactory, so we can’t use it downstream.
Model Library Format currently will export the TIR representation of this using its own printer. If we had the IRModule from tvm.relay.build, we could also export the TIR representation as well (e.g. in src/tir-N.txt). Most ideally, we would first add the IRModule to tvm.relay.build output, and then modify the Model Library Format artifact to also contain that output.
Two caveats here:
- Model Library Format can’t handle heterogeneous execution, so this might not work for you.
- There might be some limitations with TVMScript cc @junrushao @shingjan but based on initial discussion with them we think this should work fine.
Maybe you could try doing the first bit and see if it’s useful to you and permits the approach you described above (that’s indeed what I was suggesting you do, analyze the TIR to construct the producer/consumer relationship between tensors). If so we could explore adding it to Model Library Format.