tvm_tirx.ir.transform

tvm_tirx.ir.transform#

IR passes, pass contexts, sequential pipelines, and pass instrumentation.

class tvm_tirx.ir.transform.ModulePass(pass_info)#

A pass whose callback receives and returns a complete IRModule.

Create instances with module_pass(), which accepts a function or a class defining transform_module. Call after = pass_object(before) to execute the callback under the current PassContext. Its callback, rather than this wrapper class, defines how the IR changes; see module_pass() for a TVMScript before/after example.

class tvm_tirx.ir.transform.Pass(*args: Any, **kwargs: Any)#

Base interface for callable IRModule transformations and analyses.

Obtain a concrete pass from a pass factory or module_pass(), then call after = pass_object(before). The info property exposes its name and optimization level. The base class defines no transformation of its own.

property info#

Get the pass meta.

class tvm_tirx.ir.transform.PassContext(opt_level=2, required_pass=None, disabled_pass=None, instruments=None, config=None)#

Scope optimization levels, pass selection, instruments, and pass options.

Parameters:

  • opt_level (Optional[int]) – Maximum optimization level enabled in sequential pipelines (default 2).

  • required_pass (Optional[Union[List[str], Set[str], Tuple[str]]]) – Pass names to enable regardless of their optimization level.

  • disabled_pass (Optional[Union[List[str], Set[str], Tuple[str]]]) – Pass names to skip in sequential pipelines, even if required.

  • instruments (Optional[Sequence[PassInstrument]]) – The list of pass instrument implementations.

  • config (Optional[Dict[str, Object]]) – Additional configurations for specific passes.

override_instruments(instruments)#

Override instruments within this PassContext.

If there are existing instruments, their exit_pass_ctx callbacks are called. Then switching to new instruments and calling new enter_pass_ctx callbacks.

instrumentsSequence[PassInstrument]

The list of pass instrument implementations.

static current()#

Return the current pass context.

static list_configs()#

List all registered PassContext configuration names and metadata.

Returns:

configs

Return type:

Dict[str, Dict[str, str]]

class tvm_tirx.ir.transform.PassInfo(opt_level, name)#

The class contains the meta data required by a pass. It is the container of information needed by running an optimization or analysis. This class can be extended by adding new members when more meta data is needed.

Parameters:

  • opt_level (int) – The optimization level of this pass.

  • name (str) – The pass name.

class tvm_tirx.ir.transform.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.

Return type:

None

exit_pass_ctx()#

Called when exiting the instrumented context.

Return type:

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 – True to run the pass, or False to skip the pass.

Return type:

bool

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:
Return type:

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:
Return type:

None

tvm_tirx.ir.transform.PrintIR(header='')#

Log the header and module IR, returning the module unchanged.

Parameters:

header (str) – The header to be displayed along with the dump.

Returns:

fpass – A module pass with optimization level zero.

Return type:

Pass

Examples

The observation pass returns the module unchanged.

import tvm_tirx as tvm
from tvm_tirx.script import ir as I, tirx as T

@I.ir_module
class Before:
    @T.function(private=True)
    def main():
        T.evaluate(0)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main():
        T.evaluate(0)

def run_pass():
    return tvm.ir.transform.PrintIR("before lowering")(Before)

tvm.ir.assert_structural_equal(run_pass(), Expected)
class tvm_tirx.ir.transform.Sequential(passes=None, opt_level=0, name='sequential')#

A pass that works on a sequence of pass objects. Multiple passes can be executed sequentially using this class.

Passes run in the given order, subject to the optimization level and the required/disabled pass controls in PassContext.

Parameters:

  • passes (Optional[List[Pass]]) – A sequence of passes candidate for optimization.

  • opt_level (Optional[int]) – The optimization level of this sequential pass. The opt_level of a default sequential pass is set to 0. Child passes use the active PassContext’s optimization level.

  • name (Optional[str]) – The name of the sequential pass.

tvm_tirx.ir.transform.module_pass(pass_func=None, opt_level=None, name=None)#

Decorate a module pass.

This function returns a callback when pass_func is provided. Otherwise, it serves a decorator function.

pass_func can also be a class type with a method transform_module. This function will create a decorated ModulePass using transform_module as the pass function.

Parameters:

  • pass_func (Optional[Callable[(Module, PassContext) ->Module]]) – The transformation function or class.

  • opt_level (int) – The optimization level of this module pass.

  • name (Optional[str]) – The name of the module pass. The name could be empty. In this case, the name of the optimization function will be used as the pass name.

Returns:

create_module_pass – A decorator will be returned if pass_func is not provided, otherwise return the decorated result. The returned decorator has two behaviors depending on the input: A new ModulePass will be returned when we decorate a pass function. A new ModulePass class will be returned when we decorate a class type.

Return type:

Union[Callable, ModulePass]

Examples

Decorate a class for a configurable pass, or a function for a fixed pass. Both forms below simplify a TIRx module.

import tvm_tirx as tvm
from tvm_tirx.script import ir as I, tirx as T

@I.ir_module
class Before:
    @T.function(private=True)
    def main(A: T.Tensor((4,), "int32")):
        for i in T.serial(4):
            A[i] = i // 4

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Tensor((4,), "int32")):
        for i in T.serial(4):
            A[i] = 0

@tvm.ir.transform.module_pass(opt_level=0)
class CustomPipeline:
    def __init__(self, enable_simplify=True):
        self.enable_simplify = enable_simplify

    def transform_module(self, mod, ctx):
        if self.enable_simplify:
            mod = tvm.tirx.transform.StmtSimplify()(mod)
        return mod

pipeline = CustomPipeline()
tvm.ir.assert_structural_equal(pipeline(Before), Expected)
tvm.ir.assert_structural_equal(CustomPipeline(False)(Before), Before)

@tvm.ir.transform.module_pass(opt_level=0)
def simplify(mod, ctx):
    return tvm.tirx.transform.StmtSimplify()(mod)

tvm.ir.assert_structural_equal(simplify(Before), Expected)
tvm_tirx.ir.transform.pass_instrument(pi_cls=None)#

Decorate a pass instrument.

Parameters:

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

Contents