tvm.instrument

Common pass instrumentation across IR variants.

class tvm.instrument.PassInstrument

A pass instrument implementation.

To use, a user class can either subclass from PassInstrument directly, or can apply the pass_instrument() wrapper. In either case, the enter_pass_ctx, exit_pass_ctx, should_run, run_before_pass, and run_after_pass methods can be defined to adjust the instrument's behavior. See the no-op implementations in this class definition for more information on each.

enter_pass_ctx()

Called when entering the instrumented context.

Returns

None

exit_pass_ctx()

Called when exiting the instrumented context.

Returns

None

run_after_pass(mod, info)

Instrument after the pass runs.

Called once for each pass that is run while the instrumented context is active.

Parameters

mod : tvm.ir.module.IRModule

The module on which an optimization pass is being run.

info : tvm.transform.PassInfo

The pass information.

Returns

None

run_before_pass(mod, info)

Instrument before the pass runs.

Called once for each pass that is run while the instrumented context is active.

Parameters

mod : tvm.ir.module.IRModule

The module on which an optimization pass is being run.

info : tvm.transform.PassInfo

The pass information.

Returns

None

should_run(mod, info)

Determine whether to run the pass or not.

Called once for each pass that is run while the instrumented context is active.

Parameters

mod : tvm.ir.module.IRModule

The module on which an optimization pass is being run.

info : tvm.transform.PassInfo

The pass information.

Returns

should_run : bool

True to run the pass, or False to skip the pass.

class tvm.instrument.PassPrintingInstrument(print_before_pass_names, print_after_pass_names)

A pass instrument to print if before or print ir after each element of a named pass.

class tvm.instrument.PassTimingInstrument

A wrapper to create a passes time instrument that implemented in C++

static render()

Retrieve rendered time profile result Returns ------- string : string

The rendered string result of time profiles

Examples

timing_inst = PassTimingInstrument()
with tvm.transform.PassContext(instruments=[timing_inst]):
    relay_mod = relay.transform.InferType()(relay_mod)
    relay_mod = relay.transform.FoldScaleAxis()(relay_mod)
    # before exiting the context, get profile results.
    profiles = timing_inst.render()

class tvm.instrument.PrintAfterAll(*args, **kwargs)

Print the name of the pass, the IR, only after passes execute.

class tvm.instrument.PrintBeforeAll(*args, **kwargs)

Print the name of the pass, the IR, only before passes execute.

tvm.instrument._wrap_class_pass_instrument(pi_cls)

Wrap a python class as pass instrument

tvm.instrument.pass_instrument(pi_cls=None)

Decorate a pass instrument.

Parameters

pi_classclass

Instrument class. See example below.

Examples

@tvm.instrument.pass_instrument
class SkipPass:
    def __init__(self, skip_pass_name):
        self.skip_pass_name = skip_pass_name

    # Uncomment to customize
    # def enter_pass_ctx(self):
    #    pass

    # Uncomment to customize
    # def exit_pass_ctx(self):
    #    pass

    # If pass name contains keyword, skip it by return False. (return True: not skip)
    def should_run(self, mod, pass_info)
        if self.skip_pass_name in pass_info.name:
            return False
        return True

    # Uncomment to customize
    # def run_before_pass(self, mod, pass_info):
    #    pass

    # Uncomment to customize
    # def run_after_pass(self, mod, pass_info):
    #    pass

skip_annotate = SkipPass("AnnotateSpans")
with tvm.transform.PassContext(instruments=[skip_annotate]):
    tvm.relay.build(mod, "llvm")