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: Expr) → List[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: Expr) → List[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: Expr) → List[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.computable_at_compile_time(func: Function) → List[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: Expr, own_name: Var | GlobalVar | None = 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: StructInfo) → List[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: Function) → List[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: FuncStructInfo, call: Call, ctx: BlockBuilder) → 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: IRModule) → List[List[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: StructInfo, shape_var_map: Dict[Var, PrimExpr] | None = None, var_map: Dict[Var, Expr] | None = None) → 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.estimate_memory_usage(mod: IRModule | 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.

tvm.relax.analysis.free_symbolic_vars(func: Function) → List[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: Expr) → List[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: StructInfo) → 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: Function) → Dict[Var, Expr]

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: 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: Function) → Dict[str, List[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: Function) → 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: StructInfo, derived: StructInfo) → 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: StructInfo, rhs: StructInfo) → 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: PrimFunc, write_buffer_transforms: List[IndexMap | Callable]) → Dict[Block, Dict[Block | Buffer, 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: StructInfo) → List[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: DataflowBlock) → Dict[Var, List[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(mod: IRModule, check_struct_info: bool = True) → bool

Check if the IRModule is well formed.

Parameters:

  • mod (tvm.IRModule) – The input IRModule.

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