tvm.relay.ir_pass

This file contains: 1. The set of passes for Relay, which exposes an interface for configuring the

passes and scripting them in Python.
  1. The pass manager for Relay which exposes different granularity of interfaces for users to implement and use passes more conveniently.
class tvm.relay.ir_pass.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.ir_pass.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 and SequentialPass as well.

class tvm.relay.ir_pass.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.

set_pass_context(pass_ctx)

Setup the pass context for analysis and optimizations. This context could be shared by different passes for sequential passes.

Parameters:pass_ctx (PassContext) – The context that is used to help perform a certain pass or a series of passes.
class tvm.relay.ir_pass.PassContext

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.

class tvm.relay.ir_pass.PassInfo(name, opt_level, 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:
  • name (str) – The pass name.
  • opt_level (int) – The optimization level of this pass.
  • required (List[str]) – The list of passes that are required by a certain pass.
class tvm.relay.ir_pass.SequentialPass

A pass that works on a sequence of pass objects. A sequential pass class should be created through sequential_pass.

tvm.relay.ir_pass.all_type_vars(expr, mod=None)

Get all type variables from expression/type e

Parameters:
  • expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type
  • mod (tvm.relay.Module, optional) – The global module
Returns:

free – The list of all type variables in post-DFS order

Return type:

List[tvm.relay.TypeVar]

tvm.relay.ir_pass.all_vars(expr)

Get all vars from expression expr in post-DFS order.

Parameters:expr (tvm.relay.Expr) – The input expression
Returns:free – The list of all variables in post-DFS order.
Return type:List[tvm.relay.Var]
tvm.relay.ir_pass.alpha_equal(lhs, rhs)

Compare two Relay expr for structural equivalence (alpha equivalence).

Parameters:
  • lhs (tvm.relay.Expr) – One of the input Expression.
  • rhs (tvm.relay.Expr) – One of the input Expression.
Returns:

result – True iff lhs is alpha equal to rhs.

Return type:

bool

tvm.relay.ir_pass.alter_op_layout(expr)

Alternate the layouts of operators or replace primitive operators with other expressions. This pass can be used for computing convolution in custom layouts or other general weight pre-transformation.

Parameters:expr (tvm.relay.Expr) – The input expression.
Returns:transformed_expr – Transformed expression with alternated layout.
Return type:tvm.relay.Expr
tvm.relay.ir_pass.backward_fold_scale_axis(expr)

Backward fold axis scaling into weights of conv2d/dense.

Parameters:expr (tvm.relay.Expr) – The input expression, we expect that expr’s types should be fully inferred by infer_type.
Returns:folded_expr – The folded expression after transformation.
Return type:tvm.relay.Expr

Note

It is recommended to call backward_fold_scale_axis before using forward_fold_scale_axis. As backward folding targets common conv-bn pattern.

tvm.relay.ir_pass.bound_type_vars(expr, mod=None)

Get bound type variables from expression/type e

Parameters:
  • expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type
  • mod (tvm.relay.Module, optional) – The global module
Returns:

free – The list of bound type variables in post-DFS order

Return type:

List[tvm.relay.TypeVar]

tvm.relay.ir_pass.bound_vars(expr)

Get bound vars from expression expr in post-DFS order.

Parameters:expr (tvm.relay.Expr) – The input expression
Returns:free – The list of bound variables in post-DFS order.
Return type:List[tvm.relay.Var]
tvm.relay.ir_pass.canonicalize_ops(expr)

Canonicalize special operators to basic operators. This can simplify latter analysis. (e.g. Expand bias_add to expand_dims and broadcast_add.)

Parameters:e (tvm.relay.Expr) – The input Expression
Returns:result – An expression without bias_add
Return type:tvm.relay.Expr
tvm.relay.ir_pass.check_kind(t, mod=None)

Check that the type is well kinded and return the kind. For example, this mean type cannot has tensor of tensor, or is a tuple type of 2 shapes.

Parameters:
  • t (tvm.relay.Type) – The type to check
  • mod (Optional[tvm.relay.Module]) – The global module.
Returns:

kind – the kind of t

Return type:

Kind

Examples

assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Shape)])) == Shape
assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Type)])) == Type
tvm.relay.ir_pass.collect_device_annotation_ops(expr)

Collect the device annotation ops for the given expression.

Parameters:expr (tvm.relay.Expr) – The input expression.
Returns:ret – A dictionary mapping tvm.relay.Expr to device type where the keys are annotation expressions.
Return type:Dict[tvm.relay.expr, int]
tvm.relay.ir_pass.collect_device_info(expr)

Collect the device allocation map for the given expression. The device ids are propagated from the device_copy operators.

Parameters:expr (tvm.relay.Expr) – The input expression.
Returns:ret – A dictionary mapping tvm.relay.Expr to device type.
Return type:Dict[tvm.relay.expr, int]
tvm.relay.ir_pass.combine_parallel_conv2d(expr, min_num_branches=3)

Combine multiple conv2d into one.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • min_num_branches (int) – The minimum number of parallel branches when the transformation should be applied.
Returns:

transformed_expr – Transformed expression

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.dead_code_elimination(expr)

Remove expressions which does not effect the program result (dead code).

Parameters:e (tvm.relay.Expr) – The input Expression
Returns:result – An expression which is semantically equal to the input expression, but with dead code removed.
Return type:tvm.relay.Expr
tvm.relay.ir_pass.eliminate_common_subexpr(expr, fskip=None)

Eliminate common subexpressions.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • fskip (function) – The callback function that decides whether an expression should be skipped.
Returns:

expr – The output expression.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.eta_expand(expr, mod)

Add abstraction over a function.

Parameters:
  • expr (tvm.relay.Expr) – The input expression, we expect that expr’s types should be fully inferred by infer_type.
  • mod (tvm.relay.Module) – The global module.
Returns:

expanded_expr – The expression after eta expansion.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.fold_constant(expr)

Fold the constant expression in expr.

Parameters:expr (tvm.relay.Expr) – The input expression.
Returns:transformed_expr – The transformed expression.
Return type:tvm.relay.Expr
tvm.relay.ir_pass.forward_fold_scale_axis(expr)

Fold the scaling of axis into weights of conv2d/dense.

Parameters:expr (tvm.relay.Expr) – The input expression, we expect that expr’s types should be fully inferred by infer_type.
Returns:folded_expr – The folded expression after transformation.
Return type:tvm.relay.Expr

Note

It is recommended to call backward_fold_scale_axis before using forward_fold_scale_axis. As backward folding targets common conv-bn pattern.

tvm.relay.ir_pass.free_type_vars(expr, mod=None)

Get free type variables from expression/type e

Parameters:
  • expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type
  • mod (tvm.relay.Module, optional) – The global module
Returns:

free – The list of free type variables in post-DFS order

Return type:

List[tvm.relay.TypeVar]

tvm.relay.ir_pass.free_vars(expr)

Get free Vars from expression expr in Post DFS order.

Parameters:expr (tvm.relay.Expr) – The input expression
Returns:free – The list of free variables in post DFS order.
Return type:List[tvm.relay.Var]

Note

The fact that Vars are post-DFS ordred are useful in neural networks: usually this means weights of previous are ordered first.

tvm.relay.ir_pass.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.ir_pass.function_pass(opt_level=2)
def transform(func, ctx):
    return ir_pass.fold_constant(func)

function_pass = transform
assert isinstance(function_pass, ir_pass.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.
tvm.relay.ir_pass.fuse_ops(expr, opt_level=1, mod=None)

Fuse operators in expr together.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • opt_level (int) – The level of fuse optimization.
  • mod (tvm.relay.Module) – The module to perform fusion over.
Returns:

transformed_expr – Transformed expression, containing fused result.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.get_total_mac_number(expr)

Count the number of MACs (multiply-accumulate) of a model

Parameters:expr (tvm.relay.Expr) – The input expression.
Returns:ret – The number of MACs (multiply-accumulate) of a model
Return type:int64
tvm.relay.ir_pass.gradient(expr, mod=None, mode='higher_order')

Transform the input function, returning a function that calculate the original result, paired with gradient of the input.

Parameters:
  • expr (tvm.relay.Expr) – The input expression, which is a Function or a GlobalVar.
  • mod (Optional[tvm.relay.Module]) –
  • mode (Optional[String]) – The mode of the automatic differentiation algorithm. ‘first_order’ only work on first order code, but will not produce reference nor closure. ‘higher_order’ work on all code using reference and closure.
Returns:

expr – The transformed expression.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.graph_equal(lhs, rhs)

Compare two Relay expr for data-flow equivalence. The difference between this and alpha-equality is that variables are not expected to match between lhs and rhs; they are treated as sources and are mapped between each other.

Parameters:
  • lhs (tvm.relay.Expr) – One of the input Expression.
  • rhs (tvm.relay.Expr) – One of the input Expression.
Returns:

result – True iff lhs is data-flow equivalent to rhs.

Return type:

bool

tvm.relay.ir_pass.infer_type(expr, mod=None)

Infer the type of expr under the context of mod.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • mod (Optional[tvm.relay.Module]) – The global module.
Returns:

checked_expr – The checked expression.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.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 (Union[Callable, ModulePass]) – 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.

    Examples

  • ——–

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

  • the module.

  • .. code-block:: python – @relay.ir_pass.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, ir_pass.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.ir_pass.partial_evaluate(expr)

Evaluate the static fragment of the code.

Parameters:expr (tvm.relay.Expr) – The input expression.
Returns:expr – The output expression.
Return type:tvm.relay.Expr
tvm.relay.ir_pass.post_order_visit(expr, fvisit)

Recursively visit the ir in post DFS order node, apply fvisit. Each node is guaranteed to be visited only once.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • fvisit (function) – The visitor function to be applied.
tvm.relay.ir_pass.rewrite_annotated_ops(expr, fallback_device)

Rewrite the annotated program where annotation operators, e.g. on_deivce, mark which device an expression should be scheduled to. This pass helps heterogeneous execution where different operators may need to be allocated on various devices.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • fallback_device (int) – The fallback device type. It is also used as the default device for operators with no annotated device.
Returns:

transformed_expr – Transformed expression with cross device data copy operators.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.sequential_pass(passes=None, opt_level=2, name='sequential_pass', required=None, disabled=None)

Create a sequential pass using a defined optimization function from Python. 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.
  • disabled (Optional[List[str]]) – A list of disabled passes.
Returns:

ret – A sequential pass built through pass_func.

Return type:

Pass

tvm.relay.ir_pass.simplify_inference(expr)

Simplify the data-flow graph for inference phase.

Parameters:e (tvm.relay.Expr) – The input Expression
Returns:result – An expression which is semantically equal to the input expression, but with some simplification
Return type:tvm.relay.Expr
tvm.relay.ir_pass.structural_hash(value)

Hash a Relay expression structurally.

Parameters:expr (tvm.relay.Expr or tvm.relay.Type) – The expression to hash.
Returns:result – The hash value
Return type:int
tvm.relay.ir_pass.to_a_normal_form(expr, mod=None)

Turn Graph Normal Form expression into A Normal Form Expression.

The scope of the root expression is the global scope.

The scope of any non root expression is the least common ancestor of all it’s scope.

Values are ordered by post-DFS order in each scope.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • mod (Optional[tvm.relay.Module]) – The global module.
Returns:

expr – The output expression.

Return type:

tvm.relay.Expr

tvm.relay.ir_pass.to_graph_normal_form(expr)

Turn A Normal Form expression into Graph Normal Form expression :param expr: The input expression :type expr: tvm.relay.Expr

Returns:expr – The output expression
Return type:tvm.relay.Expr
tvm.relay.ir_pass.well_formed(expr)

Check that each Var is only bound once (well formed).

Parameters:expr (tvm.relay.Expr) – The input expression
Returns:well_form – Whether the input expression is well formed
Return type:bool