Creating this post as a follow-up for discussions that happened in [TVMC] Add option to dump TIR code to file by lhutton1 · Pull Request #14186 · apache/tvm · GitHub following from:
In the future it seems good to have more granularity in the IR that is being dumped, e.g. print IR before/after named pass.
The proposal is to add new options to TVMC to print the IR before/after/other-condition a named pass from the build pipeline. Initially this will involve adding two new options:
--print-ir-before=<comma-separated-list-of-pass-names>
--print-ir-after=<comma-separated-list-of-pass-names>
For example, tvmc compile --print-ir-after="PartitionGraph" could be used to print the IR after the graph has been partitioned in Relay. This can later be extended with other, more advanced, options. This would be a similar approach to the options used in MLIR.
print is suggested over dump (used by --dump-code) as it gives the user more flexibility, and they can always pipe the output to a named file of their choosing.
This seems relatively simple to implement using a pass instrument that can be attached to the current pass context during a tvmc compile. Using the run_before_pass and run_after_pass methods, the module will be printed if the name of the pass is in a user provided set of names.
cc @PhilippvK
Thank you for following up on this topic. Here are my thoughts:
- Should we differentiate between Relay and TIR Passes here? The expected naming of the passes should probably aligned with the
--disabled-pass argument. (See issues raised here: https://github.com/apache/tvm/pull/14273)
- I am agnostic regarding dumping to file or stdout as long as the printed code can the parsed by a 3rd party script in a straightforward fashion by differentiating logging messages from dumped code:
- Your proposal allows investigating temporary IRs on a very granular level which is very useful. However I would imagine that often a user would just want to look at the Relay code before/after the target-specific
pass_pipeline/partition_function (see https://github.com/apache/tvm/blob/ccc0b9162f2e983a8810e99c903c7141dbec81b6/python/tvm/driver/tvmc/compiler.py#L338) instead of directly before/after the PartitionGraph pass. I think this is a use-case which should be considered as well, either with --print-ir-before/after or using --dump-code logic.
I am looking forward for further discussion.
Thanks for the great ideas @PhilippvK,
Good point, completely agree on alignment here!
- I am agnostic regarding dumping to file or stdout as long as the printed code can the parsed by a 3rd party script in a straightforward fashion by differentiating logging messages from dumped code:
Perhaps we can have some kind of marker to add to the start and end of the dump? e.g.
--- Print "relay.transform.ConvertLayout" ---
--- End print "relay.transform.ConvertLayout" ---
would that make to output easier to process?
- Your proposal allows investigating temporary IRs on a very granular level which is very useful. However I would imagine that often a user would just want to look at the Relay code before/after the target-specific
pass_pipeline/partition_function (see tvm/python/tvm/driver/tvmc/compiler.py at ccc0b9162f2e983a8810e99c903c7141dbec81b6 · apache/tvm · GitHub) instead of directly before/after the PartitionGraph pass. I think this is a use-case which should be considered as well, either with --print-ir-before/after or using --dump-code logic.
That’s a good point, you’re correct that with the proposal so far this wouldn’t be possible, although, I think it would be good to tie into the --print-ir-before/after options. My initial thought was that we could perhaps rely on the name of a Sequential pass to provide this feature, but I’m a little unsure how these interact with the pass instrument currently - I’ll take a closer look when I have a go at implementing it 