tvm.relax.analysis#
Relax IR analysis.
- class tvm.relax.analysis.BaseCheckResult(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]#
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)[源代码]#
Return all global variables from expression expr. Parameters ---------- expr: Expr
The expression.
Returns#
- ret: List[GlobalVar]
List of global vars in expr, in post-DFS order
- 参数:
expr (RelayExpr)
- 返回类型:
List[GlobalVar]
- tvm.relax.analysis.all_vars(expr)[源代码]#
Return all (local) variables from expression expr. Parameters ---------- expr: Expr
The expression.
Returns#
- ret: List[Var]
List of vars in expr, in post-DFS order
- tvm.relax.analysis.bound_vars(expr)[源代码]#
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. Parameters ---------- expr: Expr
The expression.
Returns#
- ret: List[Var]
List of bound vars in expr, in post-DFS order
- tvm.relax.analysis.collect_non_negative_expressions(sinfo)[源代码]#
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#
- sinfoStructInfo
The struct info object to be analyzed.
Returns#
ret : List[tir.Var]
The list of TIR variables that can be defined from the StructInfo
- 参数:
sinfo (StructInfo)
- 返回类型:
List[PrimExpr]
- tvm.relax.analysis.computable_at_compile_time(func)[源代码]#
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: List[Var]
The set of variables that can be computed at compile-time, in order of their occurrence within the function.
- tvm.relax.analysis.contains_impure_call(expr, own_name=None)[源代码]#
Check if the given expression (likely a function body) contains any impure calls.
Parameters#
- exprExpr
The expression to be examined. If expr is a function, we check the body.
- own_nameVar or GlobalVar (optional)
For a recursive function, the analysis can ignore the self-calls for checking purity.
Returns#
- retbool
True if there is an impure call (call to a function that may have visible side effects).
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)[源代码]#
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#
- sinfoStructInfo
The struct info object to be analyzed.
Returns#
ret : List[tir.Var]
The list of TIR variables that can be defined from the StructInfo
- 参数:
sinfo (StructInfo)
- 返回类型:
- tvm.relax.analysis.defined_symbolic_vars(func)[源代码]#
Get the TIR variables that defined in the input function. The returned list is deduplicated - each TIR variable will appear at most once.
Parameters#
- funcFunction
The function object to be analyzed.
Returns#
- retList[Var]
The list of symbolic variables that are defined in the input function.
- tvm.relax.analysis.derive_call_ret_struct_info(func_sinfo, call, ctx)[源代码]#
Derive the call's ret value struct info from inputs.
Parameters#
- func_sinfo: FuncStructInfo
The call's function signature.
- call: Call
The call expression
- ctx: tvm.relax.BlockBuilder
The context block builder.
Returns#
- retStructInfo
The derived return value struct info.
Note#
This is an internal derivation function, call.op field is ignored in this case and the derivation only depends on func_sinfo.
- 参数:
func_sinfo (FuncStructInfo)
call (Call)
ctx (BlockBuilder)
- 返回类型:
- tvm.relax.analysis.detect_recursion(mod)[源代码]#
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: List[List[GlobalVar]]
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.
- tvm.relax.analysis.erase_to_well_defined(sinfo, shape_var_map=None, var_map=None)[源代码]#
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#
- sinfoStructInfo
The input struct info.
- shape_var_mapDict[tir.Var, tir.PrimExpr]
Specifies the defined shape vars and the values they should map to.
- var_mapDict[Var, Expr]
Specifies the defined vars and the values they should map to.
Returns#
- retStructInfo
The corresponding erased struct info.
- 参数:
sinfo (StructInfo)
- 返回类型:
- tvm.relax.analysis.estimate_memory_usage(mod)[源代码]#
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#
- modUnion[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#
- eststr
The estimation information, in the form of a string.
Notes#
We regards "relax.memory.alloc_tensor/storage" as the results produced by memory planning.
- tvm.relax.analysis.free_symbolic_vars(func)[源代码]#
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#
- funcFunction
The function object to be analyzed.
Returns#
- retList[Var]
The list of symbolic variables that are used but not defined in the input function.
- tvm.relax.analysis.free_vars(expr)[源代码]#
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. Parameters ---------- expr: Expr
The expression.
Returns#
- ret: List[Var]
List of free vars in expr, in post-DFS order
- tvm.relax.analysis.get_static_type(sinfo)[源代码]#
Get the corresponding static type from a StructInfo.
Parameters#
- sinfoStructInfo
The input struct info.
Returns#
- retType
The corresponding static type.
- 参数:
sinfo (StructInfo)
- 返回类型:
Type
- tvm.relax.analysis.get_var2val(func)[源代码]#
Get a mapping from Var to Expr for each variable in the function.
Parameters#
- funcFunction
The input function to be analyzed.
Returns#
- Dict[Var, Expr]
A mapping from Var to Expr.
- tvm.relax.analysis.has_reshape_pattern(func)[源代码]#
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#
- functir.PrimFunc
The function to be examined.
Returns#
- retbool
A boolean indicating if the given PrimFunc is doing a reshape.
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.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#
- exprtvm.relay.Expr
The input expression.
- fvisitfunction
The visitor function to be applied.
- tvm.relax.analysis.remove_all_unused(func)[源代码]#
It removes: 1. Unused local VarBindings in a DataflowBlock. 2. Unused DataflowBlocks in a function.
Parameters#
- funcFunction
The input function to be analyzed.
Notes#
For IRModule-wise DCE, use py:func:tvm.relax.transform.DeadCodeElimination.
Returns#
- Function
The function with unused variables removed.
- tvm.relax.analysis.struct_info_base_check(base, derived)[源代码]#
Run a base check to see if base subsumes derived.
Parameters#
- base: StructInfo
The base struct info.
- derived: StructInfo
The derived struct info.
Returns#
- retStructInfo
The derived return value struct info.
- 参数:
base (StructInfo)
derived (StructInfo)
- 返回类型:
- tvm.relax.analysis.struct_info_lca(lhs, rhs)[源代码]#
Unify the two struct info to their least common ancestor.
Parameters#
- lhs: StructInfo
The left operand.
- rhs: StructInfo
The right operand.
Returns#
- retStructInfo
The corresponding lca result.
- 参数:
lhs (StructInfo)
rhs (StructInfo)
- 返回类型:
- tvm.relax.analysis.suggest_layout_transforms(func, write_buffer_transforms)[源代码]#
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: Dict[Block, Dict[Union[Block, Buffer], IndexMap]]
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.
- tvm.relax.analysis.tir_vars_in_struct_info(sinfo)[源代码]#
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#
- sinfoStructInfo
The struct info object to be analyzed.
Returns#
- retList[tir.Var]
The list of TIR variables that appear in the input struct info.
- 参数:
sinfo (StructInfo)
- 返回类型:
- tvm.relax.analysis.udchain(dfb)[源代码]#
Analyze the variable use-def chain in a dataflow block.
Parameters#
- dfbDataflowBlock
The dataflow block to analyze
Returns#
- Dict[Var, List[Var]]
A mapping from variable definition to its uses.
- 参数:
dfb (DataflowBlock)
- 返回类型:
- tvm.relax.analysis.well_formed(obj, check_struct_info=True)[源代码]#
Check if the IRModule is well formed.
Parameters#
- objUnion[tvm.IRModule, Function]
The input IRModule or relax.Function.
- check_struct_infobool
A boolean flag indicating if the property "every Expr must have defined structure info" will be checked.
Returns#
- ret: bool
True if the IRModule is well formed, False if not.
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.