tvm_tirx.ir.transform.pass_instrument

tvm_tirx.ir.transform.pass_instrument(pi_cls=None)

Decorate a class as a pass instrument.

The class does not need to inherit from PassInstrument. It defines any of enter_pass_ctx, exit_pass_ctx, should_run, run_before_pass, and run_after_pass, with the signatures documented on PassInstrument; the decorator returns a PassInstrument subclass that calls them.

Parameters:

pi_cls (class) – Instrument class. See example below.

Returns:

result (type or Callable) – The instrument class, or a decorator when pi_cls is omitted.

Examples

Skip a pass by name. should_run returns False for the pass to skip:

from tvm_tirx import tirx
from tvm_tirx.ir.transform import PassContext, pass_instrument

@pass_instrument
class SkipPass:
    def __init__(self, skipped_name):
        self.skipped_name = skipped_name

    def should_run(self, mod, info):
        return info.name != self.skipped_name


with PassContext(instruments=[SkipPass("tirx.StmtSimplify")]):
    after = tirx.transform.StmtSimplify()(mod)  # returns mod unchanged