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: RelayExpr) 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: RelayExpr) 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: RelayExpr) 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.collect_non_negative_expressions(sinfo: StructInfo) List[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: 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.
- tvm.relax.analysis.contains_impure_call(expr: RelayExpr, own_name: Var | GlobalVar | None = None) bool
Check if the given expression (likely a function body) contains any impure calls.
- Parameters:
- Returns:
ret – True if there is an impure call (call to a function that may have visible side effects).
- Return type:
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.
- 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:
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, RelayExpr] | 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:
- 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.
- tvm.relax.analysis.free_vars(expr: RelayExpr) 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:
- tvm.relax.analysis.get_var2val(func: Function) Dict[Var, RelayExpr]
Get a mapping from relax.Var to Expr for each variable in the function.
- 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:
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:
- tvm.relax.analysis.struct_info_base_check(base: StructInfo, derived: StructInfo) BaseCheckResult
Run a base check to see if base subsumes derived.
- Parameters:
base (StructInfo) – The base struct info.
derived (StructInfo) – The derived struct info.
- Returns:
ret – The derived return value struct info.
- Return type:
- tvm.relax.analysis.struct_info_lca(lhs: StructInfo, rhs: StructInfo) StructInfo
Unify the two struct info to their least common ancestor.
- Parameters:
lhs (StructInfo) – The left operand.
rhs (StructInfo) – The right operand.
- Returns:
ret – The corresponding lca result.
- Return type:
- 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:
- 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:
- 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:
- tvm.relax.analysis.well_formed(obj: IRModule | Function, check_struct_info: bool = True) bool
Check if the IRModule is well formed.
- Parameters:
- Returns:
ret – True if the IRModule is well formed, False if not.
- Return type:
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: 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:
Notes
We regards “relax.memory.alloc_tensor/storage” as the results produced by memory planning.