tvm.transform¶

Common pass infrastructure across IR variants.

class tvm.transform.ModulePass¶: A pass that works on tvm.IRModule. Users don’t need to interact with this class directly. Instead, a module pass should be created through module_pass, because the design of the module_pass API is flexible enough to handle the creation of a module pass in different manners. In addition, all members of a module pass can be accessed from the base class. The same rule applies to FunctionPass as well.

class tvm.transform.Pass¶

The base class of all passes. All methods here are just simple wrappers that are implemented in the backend. They are defined for users to conveniently interact with the base class.

property info¶: Get the pass meta.

class tvm.transform.PassContext(opt_level=2, required_pass=None, disabled_pass=None, instruments=None, config=None, trace=None, trace_stack=None, make_traceable=None, num_evals=0, tuning_api_database=None)¶

The basis where a Relay optimization/analysis runs on. Each pass context contains a number of auxiliary information that is used to help an optimization pass. Such information includes the error reporter to record the errors of during the optimization, etc.

opt_levelOptional[int]: The optimization level of this pass.
required_passOptional[Union[List[str], Set[str], Tuple[str]]]: The list of passes that are required by a certain pass.
disabled_passOptional[Union[List[str], Set[str], Tuple[str]]]: The list of passes that are disabled.
instrumentsOptional[Sequence[PassInstrument]]: The list of pass instrument implementations.
configOptional[Dict[str, Object]]: Additional configurations for specific passes.
trace: Optional[relax.tuning.Trace]: Initial trace for trace mode.
trace_stack: Optional[List[relax.tuning_api.Trace]]: Initial trace stack for trace mode.
make_traceable: Optional[List[str]]: List of passes to make traceable.
num_evals: int: initial number of evaluations conducted in the pipeline.

tuning_api_database: Optional[relax.tuning_api.JSONDatabase]

static current()¶: Return the current pass context.

get_current_trace()¶: Get the trace on the top of the stack.

get_trace_stack()¶: Get the current trace stack.

get_trace_stack_size()¶: Get the size of current stack.

get_tuning_api_database()¶: Get tuning api database.

inc_num_evals(num: int)¶: Increment the number of evaluations conducted in the pipeline.

static list_configs()¶

List all registered PassContext configuration names and metadata.

Returns:: configs
Return type:: Dict[str, Dict[str, str]]

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.

pop_trace(return_current=True)¶: Pop a topmost trace from the stack. :returns: Trace :rtype: Optional[relax.tuning.Trace]

push_trace(trace)¶: Push a trace into the stack.

set_num_evals(num: int)¶: Set the number of evaluations conducted in the pipeline.

class tvm.transform.PassInfo(opt_level, name, required=None, traceable=False)¶

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.
required (List[str]) – The list of passes that are required by a certain pass.

tvm.transform.PrintIR(header='', show_meta_data=False)¶

A special trace pass that prints the header and IR.

Parameters:

header (str) – The header to be displayed along with the dump.
show_meta_data (bool) – A boolean flag to indicate if meta data should be printed.

Return type:

The pass

class tvm.transform.Sequential(passes=None, opt_level=0, name='sequential', required=None, traceable=False)¶

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

Note that users can also provide a series of passes that they don’t want to apply when running a sequential pass. Pass dependency will be resolved in the backend as well.

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. Note that some of the passes within the Sequantial may still not be executed if their opt_level is higher than the provided opt_level.
name (Optional[str]) – The name of the sequential pass.
required (Optional[List[str]]) – The list of passes that the sequential pass is dependent on.

tvm.transform.module_pass(pass_func=None, opt_level=None, name=None, required=None, traceable=False)¶

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.
required (Optional[List[str]]) – The list of passes that the module pass is dependent on.
traceable (Boolean) – Boolean variable whether the module pass is traceable

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

The following code block decorates a module pass class.

@relay.transform.module_pass
class CustomPipeline:
    def __init__(self, enable_fold):
        self.enable_fold = enable_fold
        self.cse = relay.transform.EliminateCommonSubexpr()
        self.const_fold = relay.transform.FoldConstant()

    def transform_module(self, mod, ctx):
        mod = self.cse(mod, ctx)
        if self.enable_fold:
            mod = self.const_fold(mod, ctx)
        return mod

# create an instance of customized pipeline
pipeline = CustomPipeline(enable_fold=False)
assert isinstance(pipeline, transform.ModulePass)
# run the pipeline.
output_module = pipeline(input_module)

The following code creates a module pass by decorating a user defined transform function.

@relay.transform.module_pass(opt_level=2)
def transform(mod, ctx):
    tp = relay.TensorType((10,), "float32")
    x = relay.var("x", tp)
    gv = relay.GlobalVar("var")
    func = relay.Function([x], relay.abs(x))
    new_mod = tvm.IRModule({gv: func})
    new_mod.update(mod)
    return new_mod

module_pass = transform
assert isinstance(module_pass, transform.ModulePass)
assert module_pass.info.opt_level == 2

# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = module_pass(m)
# Now a function abs should be added to the module m.