tvm_tirx.ir.transform.PassInstrument

tvm_tirx.ir.transform.PassInstrument#

class tvm_tirx.ir.transform.PassInstrument#

Bases: Object

A pass instrument implementation.

An instrument observes the passes that run inside a PassContext, for example to print, time, or skip them. Install it with PassContext(instruments=[...]); it then sees every pass that runs while the context is active, whether the pass is called directly or by a Sequential.

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.

Examples

Record the name of every pass that runs:

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

@pass_instrument
class RecordNames:
    def __init__(self):
        self.names = []

    def run_before_pass(self, mod, info):
        self.names.append(info.name)


recorder = RecordNames()
with PassContext(instruments=[recorder]):
    tirx.transform.StmtSimplify()(mod)
print(recorder.names)  # ['tirx.StmtSimplify']
enter_pass_ctx()#

Called when entering the instrumented context.

Returns:

None

exit_pass_ctx()#

Called when exiting the instrumented context.

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

should_run (bool) – True to run the pass, or False to skip the pass.

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

None

property name#