tvm.relay.transform

Relay pass transformation infrastructure.

tvm.relay.transform.build_config(opt_level=2, fallback_device=cpu(0), required_pass=None, disabled_pass=None)

Configure the build behavior by setting config variables.

Parameters:
  • opt_level (int, optional) –

    Optimization level. The optimization pass name and level are as the following:

    OPT_PASS_LEVEL = {
        "SimplifyInference": 0,
        "OpFusion": 1,
        "FoldConstant": 2,
        "CombineParallelConv2D": 3,
        "FoldScaleAxis": 3,
        "AlterOpLayout": 3,
        "CanonicalizeOps": 3,
        "EliminateCommonSubexpr": 3,
    }
    
  • fallback_device (int, str, or tvm.TVMContext, optional) – The fallback device. It is also used as the default device for operators without specified device during heterogeneous execution.
  • required_pass (set of str, optional) – Optimization passes that are required regardless of optimization level.
  • disabled_pass (set of str, optional) – Optimization passes to be disabled during optimization.
Returns:

pass_context – The pass context for optimizations.

Return type:

PassContext

tvm.relay.transform.module_pass(pass_func=None, opt_level=None, name=None, required=None)

Create a module pass. This function returns a callback when pass_func is provided. Otherwise, it returns the created module level pass using the given optimization function.

Parameters:
  • pass_func (Optional[Callable[(Module/Function, PassContext) ->) – Module/Function]] The implemented optimization pass.
  • 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.
Returns:

create_module_pass – The callable that will create a module pass is returned when pass_func is not passed in. Otherwise, a ModulePass object will be directly created.

Return type:

Union[Callable, ModulePass]

Examples

The following code creates a module level pass and adds an abs function to the module.

@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 = relay.Module({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.
tvm.relay.transform.function_pass(pass_func=None, opt_level=None, name=None, required=None)

Create a function pass. This function returns a callback when pass_func is provided. Otherwise, it returns the created function pass using the given optimization function.

Parameters:
  • pass_func (Optional[Callable[(Module/Function, PassContext) ->) – Module/Function]] The implemented optimization pass.
  • opt_level (int) – The optimization level of this module pass.
  • name (Optional[str]) – The name of the function 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.
Returns:

create_function_pass – The callable that will create a function pass is returned when pass_func is not passed in. Otherwise, a FunctionPass object will be created.

Return type:

Union[Callable, FunctionPass]

Examples

The following code creates a function level pass that performs constant folding.

@relay.transform.function_pass(opt_level=2)
def transform(func, ctx):
    return ir_pass.fold_constant(func)

function_pass = transform
assert isinstance(function_pass, transform.FunctionPass)
assert function_pass.info.opt_level == 2

# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = function_pass(m)
# Now constant folding should have been applied to every function in
# the provided module m. And the updated module will be returned.
class tvm.relay.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.

info

Get the pass meta.

class tvm.relay.transform.PassInfo(opt_level, name, required=None)

The class that 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.
class tvm.relay.transform.PassContext(opt_level=2, fallback_device=cpu(0), required_pass=None, disabled_pass=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_level : Optional[int]
The optimization level of this pass.
fallback_device : Optional[Union[int, str, TVMContext]]
The fallback device type. It is also used as the default device for operators that are not annotated during heterogeneous execution.
required_pass : Optional[Union[List[str], Set[str], Tuple[str]]]
The list of passes that are required by a certain pass.
disabled_pass : Optional[Union[List[str], Set[str], Tuple[str]]]
The list of passes that are disabled.
static current()

Return the current pass context.

class tvm.relay.transform.ModulePass

A pass that works on tvm.relay.Module. 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.relay.transform.FunctionPass

A pass that works on each tvm.relay.Function in a module. A function pass class should be created through function_pass.

class tvm.relay.transform.Sequential(passes=None, opt_level=2, name='sequential', required=None)

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

Some typical usage of the sequential pass are: 1. Users provide a list of passes for optimization. 2. Only an optimization level is provided so that the backend system has

to glob all passes at this level and below to perform the optimizations.

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.
  • name (Optional[str]) – The name of the sequential pass.
  • required (Optional[List[str]]) – The list of passes that the sequential pass is dependent on.