tvm.relay.ir_pass

The set of passes for Relay.

Exposes an interface for configuring the passes and scripting them in Python.

tvm.relay.ir_pass.all_type_vars(expr)

Get all type variables from expression/type e

Parameters:expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type
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)

Get bound type variables from expression/type e

Parameters:expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type
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. 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 (tvm.relay.Module, optional) – The global module
Returns:

well_kinded – whether the input type is well kinded.

Return type:

bool

Examples

assert not check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Shape)]))
assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.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)

Fold multiple conv2d into one.

Parameters:expr (tvm.relay.Expr) – The input expression.
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.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)

Get free type variables from expression/type e

Parameters:expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type
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.fuse_ops(expr, opt_level=1)

Fuse operators in expr together.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.
  • opt_level (int) – The level of fuse optimization.
Returns:

transformed_expr – Transformed expression, containing fused result.

Return type:

tvm.relay.Expr

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

.

Parameters:
  • expr (tvm.relay.Expr) – The input expression, which is a Function or a GlobalVar.
  • mod (Optional[tvm.relay.Module]) – The global module.
Returns:

ret – A function that calculate the original result paired with gradient.

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.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.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.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