tvm.relax.analysis

Relax IR analysis.

class tvm.relax.analysis.BaseCheckResult(value)

Return result of fine-grained base check.

Note

Base check comes with fine-grained fail levels.

  • FAIL_L0: The lhs and rhs have no intersection at all.

  • FAIL_L1: We get the failure by looking at static information.

  • FAIL_L2: We get the failure due to unknown symbolic variable relations.

tvm.relax.analysis.all_global_vars(expr: tvm.ir.expr.RelayExpr) → List[tvm.ir.expr.GlobalVar]

Return all global variables from expression expr. :param expr: The expression. :type expr: Expr

Returns

ret – List of global vars in expr, in post-DFS order

Return type

List[GlobalVar]

tvm.relax.analysis.all_vars(expr: tvm.ir.expr.RelayExpr) → List[tvm.relax.expr.Var]

Return all (local) variables from expression expr. :param expr: The expression. :type expr: Expr

Returns

ret – List of vars in expr, in post-DFS order

Return type

List[relax.Var]

tvm.relax.analysis.bound_vars(expr: tvm.ir.expr.RelayExpr) → List[tvm.relax.expr.Var]

Return all bound variables from expression expr. Bound variables are all variables that are declared in the expr. They only have meaning inside that expr, and can only be used in it. :param expr: The expression. :type expr: Expr

Returns

ret – List of bound vars in expr, in post-DFS order

Return type

List[relax.Var]

tvm.relax.analysis.collect_non_negative_expressions(sinfo: tvm.relax.expr.StructInfo) → List[tvm.ir.expr.PrimExpr]

Collect TIR expressions used in non-negative contexts

Get TIR variables that are non-negative within the context where the struct info is used. For example, any expression used as a tensor shape.

The returned list is deduplicated - each TIR expression will appear at most once. The order of the list is in the order of occurrence within the struct info.

Parameters

sinfo (StructInfo) – The struct info object to be analyzed.

Returns

ret – The list of TIR variables that can be defined from the StructInfo

Return type

List[tir.Var]

tvm.relax.analysis.computable_at_compile_time(func: tvm.relax.expr.Function) → List[tvm.relax.expr.Var]

Collect variables whose value can be computed at compile-time

If a function has the kNumInput attribute, then the first kNumInput parameters are provided at run-time, while all remaining parameters may be known at compile-time. This utility collects all variable bindings that only depend, directly or indirectly, on the parameters known at compile-time.

Parameters

func (Function) – The relax.Function to analyze

Returns

ret – The set of variables that can be computed at compile-time, in order of their occurrence within the function.

Return type

List[relax.Var]

tvm.relax.analysis.contains_impure_call(expr: tvm.ir.expr.RelayExpr, own_name: Optional[Union[tvm.relax.expr.Var, tvm.ir.expr.GlobalVar]] = None) → bool

Check if the given expression (likely a function body) contains any impure calls.

Parameters

  • expr (Expr) – The expression to be examined. If expr is a function, we check the body.

  • own_name (relax.Var or GlobalVar (optional)) – For a recursive function, the analysis can ignore the self-calls for checking purity.

Returns

ret – True if there is an impure call (call to a function that may have visible side effects).

Return type

bool

Notes

Relies on StructInfo annotations, so ensure that the module has been normalized first. Also, an impure call in a nested function does not mean that the outer expression contains an impure call–it only does if the nested function is later called.

tvm.relax.analysis.definable_tir_vars_in_struct_info(sinfo: tvm.relax.expr.StructInfo) → List[tvm.tir.expr.Var]

Get the TIR variables that may be defined from input struct info. The returned list is deduplicated - each TIR variable will appear at most once.

Parameters

sinfo (StructInfo) – The struct info object to be analyzed.

Returns

ret – The list of TIR variables that can be defined from the StructInfo

Return type

List[tir.Var]

tvm.relax.analysis.defined_symbolic_vars(func: tvm.relax.expr.Function) → List[tvm.relax.expr.Var]

Get the TIR variables that defined in the input function. The returned list is deduplicated - each TIR variable will appear at most once.

Parameters

func (Function) – The function object to be analyzed.

Returns

ret – The list of symbolic variables that are defined in the input function.

Return type

List[relax.Var]

tvm.relax.analysis.derive_call_ret_struct_info(func_sinfo: tvm.relax.struct_info.FuncStructInfo, call: tvm.relax.expr.Call, ctx: tvm.relax.block_builder.BlockBuilder) → tvm.relax.expr.StructInfo

Derive the call’s ret value struct info from inputs.

Parameters

  • func_sinfo (FuncStructInfo) – The call’s function signature.

  • call (relax.Call) – The call expression

  • ctx (tvm.relax.BlockBuilder) – The context block builder.

Returns

ret – The derived return value struct info.

Return type

StructInfo

Note

This is an internal derivation function, call.op field is ignored in this case and the derivation only depends on func_sinfo.

tvm.relax.analysis.detect_recursion(mod: tvm.ir.module.IRModule) → List[List[tvm.ir.expr.GlobalVar]]

Find all sets of recursive or mutually recursive functions in the module.

Two or more functions are mutually recursive if there is some cycle of references among them. For example, if there are two functions A and B, they are mutually recursive if A calls B and B calls A. Another case would be with three functions A, B, and C, where A calls B, B calls C, and C calls A.

(Note that functions do not have to call each other to reference each other. For example, if a function returns another function, that is still a reference that could potentially be recursive, even without a call.)

If a function is simply recursive and not mutually recursive with any other, it will be reported as a group by itself.

Parameters

mod (The module) –

Returns

ret – Each member of the list is a list of global functions that references each other mutually recursively. If a function is simply recursive and not mutually recursive with any other, it will be a singleton in this list.

Return type

List[List[GlobalVar]]

tvm.relax.analysis.erase_to_well_defined(sinfo: tvm.relax.expr.StructInfo, shape_var_map: Optional[Dict[tvm.tir.expr.Var, tvm.ir.expr.PrimExpr]] = None, var_map: Optional[Dict[tvm.relax.expr.Var, tvm.ir.expr.RelayExpr]] = None) → tvm.relax.expr.StructInfo

Erase sinfo into a well defined form.

This function removes the StructInfo’s dependencies on shape and vars that are not defined in given maps.

Parameters

  • sinfo (StructInfo) – The input struct info.

  • shape_var_map (Dict[tir.Var, tir.PrimExpr]) – Specifies the defined shape vars and the values they should map to.

  • var_map (Dict[relax.Var, Expr]) – Specifies the defined vars and the values they should map to.

Returns

ret – The corresponding erased struct info.

Return type

StructInfo

tvm.relax.analysis.free_symbolic_vars(func: tvm.relax.expr.Function) → List[tvm.relax.expr.Var]

Get the TIR variables that are used but not defined in the input function. The returned list is deduplicated - each TIR variable will appear at most once.

Parameters

func (Function) – The function object to be analyzed.

Returns

ret – The list of symbolic variables that are used but not defined in the input function.

Return type

List[relax.Var]

tvm.relax.analysis.free_vars(expr: tvm.ir.expr.RelayExpr) → List[tvm.relax.expr.Var]

Return all free variables from expression expr. Free variables are variables that are not bound by a VarBinding or a function parameter in the expression. :param expr: The expression. :type expr: Expr

Returns

ret – List of free vars in expr, in post-DFS order

Return type

List[relax.Var]

tvm.relax.analysis.get_static_type(sinfo: tvm.relax.expr.StructInfo) → tvm.ir.type.Type

Get the corresponding static type from a StructInfo.

Parameters

sinfo (StructInfo) – The input struct info.

Returns

ret – The corresponding static type.

Return type

Type

tvm.relax.analysis.get_var2val(func: tvm.relax.expr.Function) → Dict[tvm.relax.expr.Var, tvm.ir.expr.RelayExpr]

Get a mapping from relax.Var to Expr for each variable in the function.

Parameters

func (Function) – The input function to be analyzed.

Returns

A mapping from relax.Var to Expr.

Return type

Dict[relax.Var, Expr]

tvm.relax.analysis.has_reshape_pattern(func: tvm.tir.function.PrimFunc) → bool

Check if the given PrimFunc is essentially doing a reshape operation. The reshape operation also includes expand_dims, squeeze, flatten, etc.

Here the allowed reshape pattern is: for example, assume the operation is B[l_0, l_1, …, l_b] = A[r_0, r_1, …, r_a], we check if we can prove that the flattened index of l_0, …, l_b under buffer B equals to the flattened index of r_0, …, r_a under buffer A.

Parameters

func (tir.PrimFunc) – The function to be examined.

Returns

ret – A boolean indicating if the given PrimFunc is doing a reshape.

Return type

bool

Notes

According to the description above, the returned result can only be false-negative and cannot be false-positive, since whenever we cannot prove the equality, we return false. This property guarantees the safety of this function.

tvm.relax.analysis.name_to_binding(func: tvm.relax.expr.Function) → Dict[str, List[tvm.relax.expr.Binding]]

Return a map from variable name to its bindings.

tvm.relax.analysis.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.relax.analysis.remove_all_unused(func: tvm.relax.expr.Function) → tvm.relax.expr.Function

It removes: 1. Unused local VarBindings in a DataflowBlock. 2. Unused DataflowBlocks in a function.

Parameters

func (Function) – The input function to be analyzed.

Notes

For IRModule-wise DCE, use py:func:tvm.relax.transform.DeadCodeElimination.

Returns

The function with unused variables removed.

Return type

Function

tvm.relax.analysis.struct_info_base_check(base: tvm.relax.expr.StructInfo, derived: tvm.relax.expr.StructInfo) → tvm.relax.analysis.analysis.BaseCheckResult

Run a base check to see if base subsumes derived.

Parameters
Returns

ret – The derived return value struct info.

Return type

StructInfo

tvm.relax.analysis.struct_info_lca(lhs: tvm.relax.expr.StructInfo, rhs: tvm.relax.expr.StructInfo) → tvm.relax.expr.StructInfo

Unify the two struct info to their least common ancestor.

Parameters
Returns

ret – The corresponding lca result.

Return type

StructInfo

tvm.relax.analysis.suggest_layout_transforms(func: tvm.tir.function.PrimFunc, write_buffer_transforms: List[Union[tvm.tir.function.IndexMap, Callable]]) → Dict[tvm.tir.stmt.Block, Dict[Union[tvm.tir.stmt.Block, tvm.tir.buffer.Buffer], tvm.tir.function.IndexMap]]

Suggest Layout transformations of blocks and buffers in a PrimFunc.

Parameters

  • func (PrimFunc) – PrimFunc on which analysis will be performed and transformations suggested.

  • write_buffer_transforms (List[Union[IndexMap, Callable]) – List of layout transformations on the output buffers. The number of layout transformations must match the number of outputs of the PrimFunc.

Returns

ret – Suggested transforms per block in func. For each block the returned value is a map from the object (block or buffer) to it’s index map transformation.

Return type

Dict[Block, Dict[Union[Block, Buffer], IndexMap]]

tvm.relax.analysis.tir_vars_in_struct_info(sinfo: tvm.relax.expr.StructInfo) → List[tvm.tir.expr.Var]

Get the TIR variables that appear in the input struct info. The returned list is deduplicated - each TIR variable will appear at most once.

Parameters

sinfo (StructInfo) – The struct info object to be analyzed.

Returns

ret – The list of TIR variables that appear in the input struct info.

Return type

List[tir.Var]

tvm.relax.analysis.udchain(dfb: tvm.relax.expr.DataflowBlock) → Dict[tvm.relax.expr.Var, List[tvm.relax.expr.Var]]

Analyze the variable use-def chain in a dataflow block.

Parameters

dfb (DataflowBlock) – The dataflow block to analyze

Returns

A mapping from variable definition to its uses.

Return type

Dict[relax.Var, List[relax.Var]]

tvm.relax.analysis.well_formed(obj: Union[tvm.ir.module.IRModule, tvm.relax.expr.Function], check_struct_info: bool = True) → bool

Check if the IRModule is well formed.

Parameters

  • obj (Union[tvm.IRModule, Function]) – The input IRModule or relax.Function.

  • check_struct_info (bool) – A boolean flag indicating if the property “every Expr must have defined structure info” will be checked.

Returns

ret – True if the IRModule is well formed, False if not.

Return type

bool

Note

By default the structure info is always checked. It is only in test cases where check_struct_info might be false, so that other well-formed requirements will be well tested and will not be blocked by not having structure info.

tvm.relax.analysis.estimate_memory_usage(mod: Union[tvm.ir.module.IRModule, tvm.relax.expr.Function]) → str

Analysis function that estimates the memory usage of Relax functions in an IRModule. The estimation includes the total memory size needed to be allocated before and after memory planning.

The result might be over-estimated, as the estimation is static, which does not consider control flows (such as “if” and cross-function calls). It simply accumulates the size of every alloc_tensor and alloc_storage.

This analysis function is used to demonstrate the effect of memory planning.

Parameters

mod (Union[IRModule, Function]) – The input IRModule whose functions inside are to be analyzed. If the input is a Function, we will wrap it with a IRModule, with the function named “main”.

Returns

est – The estimation information, in the form of a string.

Return type

str

Notes

We regards “relax.memory.alloc_tensor/storage” as the results produced by memory planning.