wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/analysis/analysis.py
+621
View File
@@ -0,0 +1,621 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=no-else-return, invalid-name
# pylint: disable=unidiomatic-typecheck
"""
This file contains the set of passes for Relax, which exposes an interface for
configuring the passes and scripting them in Python.
"""
from collections.abc import Callable
from enum import IntEnum
import tvm
from tvm import IRModule, tirx
from tvm.ir import Call, Type
from tvm.relax.expr import Binding, DataflowBlock, Expr, Function, GlobalVar, Var
from tvm.relax.type import FuncType
from tvm.tirx import Buffer, IndexMap, PrimFunc, SBlock
from . import _ffi_api
def get_static_type(ty: Type) -> Type:
"""Get the corresponding static type from a Type.
Parameters
----------
ty : Type
The input type.
Returns
-------
ret : Type
The corresponding static type.
"""
return _ffi_api.GetStaticType(ty) # type: ignore
def erase_to_well_defined(
ty: Type,
shape_var_map: dict[tirx.Var, tirx.Expr] | None = None,
var_map: dict[Var, Expr] | None = None,
) -> Type:
"""Erase ty into a well defined form.
This function removes the Type's dependencies on shape and vars that
are not defined in given maps.
Parameters
----------
ty : Type
The input type.
shape_var_map : Dict[tirx.Var, tirx.Expr]
Specifies the defined shape vars and the values they should map to.
var_map : Dict[Var, Expr]
Specifies the defined vars and the values they should map to.
Returns
-------
ret : Type
The corresponding erased type.
"""
shape_var_map = {} if shape_var_map is None else shape_var_map
var_map = {} if var_map is None else var_map
return _ffi_api.EraseToWellDefined(ty, shape_var_map, var_map) # type: ignore
class BaseCheckResult(IntEnum):
"""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.
"""
FAIL_L0 = 0
FAIL_L1 = 1
FAIL_L2 = 2
PASS = 3
def type_base_check(base: Type, derived: Type) -> BaseCheckResult:
"""Run a base check to see if base subsumes derived.
Parameters
----------
base: Type
The base type.
derived: Type
The derived type.
Returns
-------
ret : Type
The derived return value type.
"""
return _ffi_api.TypeBaseCheck(base, derived) # type: ignore
def derive_call_ret_type(func_ty: FuncType, call: Call, ctx: "tvm.relax.BlockBuilder") -> Type:
"""Derive the call's ret value type from inputs.
Parameters
----------
func_ty: FuncType
The call's function signature.
call: Call
The call expression
ctx: tvm.relax.BlockBuilder
The context block builder.
Returns
-------
ret : Type
The derived return value type.
Note
----
This is an internal derivation function, call.op field is
ignored in this case and the derivation only depends on func_ty.
"""
return _ffi_api.DeriveCallRetType(func_ty, call, ctx) # type: ignore
def type_lca(lhs: Type, rhs: Type) -> Type:
"""Unify the two type to their least common ancestor.
Parameters
----------
lhs: Type
The left operand.
rhs: Type
The right operand.
Returns
-------
ret : Type
The corresponding lca result.
"""
return _ffi_api.TypeLCA(lhs, rhs) # type: ignore
def tir_vars_in_type(ty: Type) -> list[tirx.Var]:
"""Get the TIR variables that appear in the input type.
The returned list is deduplicated - each TIR variable will appear at most once.
Parameters
----------
ty : Type
The type object to be analyzed.
Returns
-------
ret : List[tirx.Var]
The list of TIR variables that appear in the input type.
"""
return _ffi_api.TIRVarsInType(ty) # type: ignore
def definable_tir_vars_in_type(ty: Type) -> list[tirx.Var]:
"""Get the TIR variables that may be defined from input type.
The returned list is deduplicated - each TIR variable will appear at most once.
Parameters
----------
ty : Type
The type object to be analyzed.
Returns
-------
ret : List[tirx.Var]
The list of TIR variables that can be defined from the Type
"""
return _ffi_api.DefinableTIRVarsInType(ty) # type: ignore
def collect_non_negative_expressions(ty: Type) -> list[tirx.Expr]:
"""Collect TIR expressions used in non-negative contexts
Get TIR variables that are non-negative within the context where
the type 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 type.
Parameters
----------
ty : Type
The type object to be analyzed.
Returns
-------
ret : List[tirx.Var]
The list of TIR variables that can be defined from the Type
"""
return _ffi_api.CollectNonNegativeExpressions(ty) # type: ignore
def 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 : List[Var]
The list of symbolic variables that are defined in the input function.
"""
return _ffi_api.DefinedSymbolicVars(func) # type: ignore
def 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 : List[Var]
The list of symbolic variables that are used but not defined in the input function.
"""
return _ffi_api.FreeSymbolicVars(func) # type: ignore
def 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.
Parameters
----------
expr: Expr
The expression.
Returns
-------
ret: List[Var]
List of bound vars in expr, in post-DFS order
"""
return _ffi_api.bound_vars(expr)
def 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.
Parameters
----------
expr: Expr
The expression.
Returns
-------
ret: List[Var]
List of free vars in expr, in post-DFS order
"""
return _ffi_api.free_vars(expr)
def all_vars(expr: Expr) -> list[Var]:
"""
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
"""
return _ffi_api.all_vars(expr)
def used_vars(expr: Expr) -> list[Var]:
"""
Return all variables used in an expression.
This function collects all variable references within the given expression,
which is useful for analyzing variable dependencies.
Parameters
----------
expr: Expr
The expression to analyze.
Returns
-------
ret: List[Var]
List of variables used in the expression.
"""
return _ffi_api.used_vars(expr) # type: ignore
def all_global_vars(expr: Expr) -> list[GlobalVar]:
"""
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
"""
return _ffi_api.all_global_vars(expr)
def 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.relax.Expr
The input expression.
fvisit : function
The visitor function to be applied.
"""
return _ffi_api.post_order_visit(expr, fvisit) # type: ignore
def has_reshape_pattern(func: tirx.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 : tirx.PrimFunc
The function to be examined.
Returns
-------
ret : bool
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.
"""
return _ffi_api.has_reshape_pattern(func) # type: ignore
def 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 : Var or GlobalVar (optional)
For a recursive function, the analysis can ignore the self-calls
for checking purity.
Returns
-------
ret : bool
True if there is an impure call
(call to a function that may have visible side effects).
Notes
-----
Relies on Type 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*.
"""
return _ffi_api.contains_impure_call(expr, own_name)
def get_var2val(func: Function) -> dict[Var, Expr]:
"""
Get a mapping from Var to Expr for each variable in the function.
Parameters
----------
func : Function
The input function to be analyzed.
Returns
-------
Dict[Var, Expr]
A mapping from Var to Expr.
"""
return _ffi_api.get_var2val(func) # type: ignore
def 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
-------
Dict[Var, List[Var]]
A mapping from variable definition to its uses.
"""
return _ffi_api.udchain(dfb) # type: ignore
def name_to_binding(func: Function) -> dict[str, list[Binding]]:
"""Return a map from variable name to its bindings."""
return _ffi_api.name_to_binding(func) # type: ignore
def 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
-------
Function
The function with unused variables removed.
"""
return _ffi_api.remove_all_unused(func) # type: ignore
def well_formed(obj: IRModule | Function, check_ty: bool = True) -> None:
"""Check if the IRModule is well formed, raising on the first violation.
Raises an error (seeded with the offending node so a pass runner can report a
precise access path) on the first well-formedness violation. Use
:func:`check_well_formed` for a boolean answer.
Parameters
----------
obj : Union[tvm.IRModule, Function]
The input IRModule or relax.Function.
check_ty : bool
A boolean flag indicating if the property "every Expr must
have defined type information" will be checked.
Note
----
By default the type information is always checked. It is only in test cases
where `check_ty` might be false, so that other well-formed requirements
will be well tested and will not be blocked by not having type information.
"""
_ffi_api.well_formed(obj, check_ty) # type: ignore
def check_well_formed(obj: IRModule | Function, check_ty: bool = True) -> bool:
"""Return whether the IRModule or Function is well formed.
Wraps :func:`well_formed`, returning False instead of raising on the first violation.
Parameters
----------
obj : Union[tvm.IRModule, Function]
The input IRModule or relax.Function.
check_ty : bool
A boolean flag indicating if the property "every Expr must
have defined type information" will be checked.
Returns
-------
ret: bool
True if the IRModule is well formed, False if not.
"""
return _ffi_api.check_well_formed(obj, check_ty) # type: ignore
def _get_prim_func_default_dtype(func: PrimFunc):
"""Detect default index dtype from function buffer map"""
for _, v in func.buffer_map.items():
for value in v.shape:
return value.ty
return "int64"
def suggest_layout_transforms(
func: PrimFunc, write_buffer_transforms: list[IndexMap | Callable]
) -> dict[SBlock, dict[SBlock | 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: Dict[SBlock, Dict[Union[SBlock, 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.
"""
write_buffer_index_maps = []
default_index_dtype = _get_prim_func_default_dtype(func)
for transform in write_buffer_transforms:
if callable(transform):
transform = IndexMap.from_func(transform, index_dtype=default_index_dtype)
assert isinstance(transform, IndexMap)
write_buffer_index_maps.append(transform)
return _ffi_api.suggest_layout_transforms(func, write_buffer_index_maps) # type: ignore
def detect_recursion(mod: tvm.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: 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.
"""
return _ffi_api.detect_recursion(mod) # type: ignore
def 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: List[Var]
The set of variables that can be computed at compile-time, in
order of their occurrence within the function.
"""
return _ffi_api.computable_at_compile_time(func) # type: ignore