TVMC print IR before/after named pass

Development
#1

Creating this post as a follow-up for discussions that happened in https://github.com/apache/tvm/pull/14186 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

#2

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.

#3

Thanks for the great ideas @PhilippvK,

  • 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)

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?

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 :slight_smile:

#4

A variant of this was implemented in: https://github.com/apache/tvm/pull/16261