Target overrides in tir/pass?

Currently all tir passes are self-contained, i.e. the class (typically a subclass of a mutator class) and the member functions are all defined in a single .cc file and run as a unit. This doesn’t allow any of the member functions to be overridden to implement custom processing.

Is it a design choice to disallow this? Or can this be modified?

The modifications would be to create a separate header file, and let targets produce their own objects inherited from the default one, that would then be used in the pass.

We prefer to hide as much as possible by default. If there is a need to reuse a mutator, then we could create a header file(starting from internal header in src), for example https://github.com/apache/incubator-tvm/blob/master/src/tir/pass/storage_access.h

It would also be useful to think about how do we do target overrides. There are a few ways:

  • A0: Target creates its own mutator and pass
  • A1: Target registers certain callback funtions that handles a target related optimizations
  • A2: Target registers certain attributes(e.g.) can an op X be lowered to Y in this target and the pass takes that as an input.

This thread is also related Target and Attributes

One thing that I want to note is that we are going to change the way we refactor TIR passes to consoildate it with the relay via a unified IRModule and Pass (hopefully in a month). Some of the changes have already landed(e.g PrimFunc)

In the new convention:

  • The basic unit among passes will be IRModule containing a collection of tir::PrimFunc
  • Each PrimFunc could have an optional target attribute, indicating the target we intend to lower the funtion to
  • The passes should be able to read from these attributes and do some customization.
1 Like