tvm.relax.analysis

Relax IR 分析。

class tvm.relax.analysis.BaseCheckResult(value)

返回细粒度基础检查的结果。

注意

基础检查带有细粒度的失败级别。

  • FAIL_L0：lhs 和 rhs 完全没有交集。

  • FAIL_L1：通过查看静态信息得到失败。

  • FAIL_L2：由于未知的符号变量关系而导致失败。

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

从表达式 expr 返回所有全局变量。 :param expr: 表达式。 :type expr: Expr

返回:

ret – expr 中全局变量的列表，按后序 DFS 顺序排列

返回类型:

List[GlobalVar]

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

从表达式 expr 返回所有（局部）变量。 :param expr: 表达式。 :type expr: Expr

返回:

ret – expr 中变量的列表，按后序 DFS 顺序排列

返回类型:

List[relax.Var]

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

从表达式 expr 返回所有绑定变量。绑定变量是在 expr 中声明的所有变量。它们仅在 expr 内部有意义，并且只能在其中使用。 :param expr: 表达式。 :type expr: Expr

返回:

ret – expr 中绑定变量的列表，按后序 DFS 顺序排列

返回类型:

List[relax.Var]

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

收集在非负上下文中使用 TIR 表达式

获取在结构信息使用的上下文中为非负的 TIR 变量。例如，用作张量形状的任何表达式。

返回的列表是去重后的 - 每个 TIR 表达式最多出现一次。列表的顺序是结构信息中出现的顺序。

参数:

sinfo (StructInfo) – 要分析的结构信息对象。

返回:

ret – 可以从 StructInfo 定义的 TIR 变量列表

返回类型:

List[tir.Var]

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

收集其值可以在编译时计算的变量

如果函数具有 kNumInput 属性，则前 kNumInput 个参数在运行时提供，而所有剩余参数可能在编译时已知。此实用程序收集所有变量绑定，这些绑定仅直接或间接地依赖于编译时已知的参数。

参数:

func (Function) – 要分析的 relax.Function

返回:

ret – 可以在编译时计算的变量集合，按其在函数中出现的顺序排列。

返回类型:

List[relax.Var]

tvm.relax.analysis.contains_impure_call(expr: RelaxExpr, own_name: Var | GlobalVar | None = None) → bool

检查给定表达式（可能是函数体）是否包含任何不纯调用。

参数:

  • expr (Expr) – 要检查的表达式。如果 expr 是一个函数，我们将检查函数体。

  • own_name (relax.Var 或 GlobalVar (可选)) – 对于递归函数，分析可以忽略自调用以检查纯度。

返回:

ret – 如果存在不纯调用（调用可能具有可见副作用的函数），则为 True。

返回类型:

bool

注释

依赖于 StructInfo 注释，因此请确保首先对模块进行规范化。此外，嵌套 函数中的不纯调用并不 意味着外部表达式包含不纯调用 - 仅当嵌套函数稍后被调用时才如此。

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

获取可能从输入结构信息中定义的 TIR 变量。返回的列表是去重后的 - 每个 TIR 变量最多出现一次。

参数:

sinfo (StructInfo) – 要分析的结构信息对象。

返回:

ret – 可以从 StructInfo 定义的 TIR 变量列表

返回类型:

List[tir.Var]

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

获取在输入函数中定义的 TIR 变量。返回的列表是去重后的 - 每个 TIR 变量最多出现一次。

参数:

func (Function) – 要分析的函数对象。

返回:

ret – 在输入函数中定义的符号变量列表。

返回类型:

List[relax.Var]

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

从输入导出调用的返回值结构信息。

参数:

  • func_sinfo (FuncStructInfo) – 调用函数的签名。

  • call (relax.Call) – 调用表达式

  • ctx (tvm.relax.BlockBuilder) – 上下文块构建器。

返回:

ret – 导出的返回值结构信息。

返回类型:

StructInfo

注意

这是一个内部导出函数，在这种情况下会忽略 call.op 字段，导出仅取决于 func_sinfo。

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

查找模块中所有递归或相互递归的函数集。

如果两个或多个函数之间存在一些引用循环，则它们是相互递归的。例如，如果有两个函数 A 和 B，如果 A 调用 B 且 B 调用 A，则它们是相互递归的。另一种情况是三个函数 A、B 和 C，其中 A 调用 B，B 调用 C，C 调用 A。

（请注意，函数不必互相调用即可互相引用。例如，如果一个函数返回另一个函数，那仍然是一个可能递归的引用，即使没有调用。）

如果一个函数只是简单递归，而不是与任何其他函数相互递归，它将被报告为一组。

参数:

mod (模块)

返回:

ret – 列表的每个成员都是一个全局函数列表，这些函数相互递归引用。如果一个函数只是简单递归，而不是与任何其他函数相互递归，它将是此列表中的单例。

返回类型:

List[List[GlobalVar]]

tvm.relax.analysis.erase_to_well_defined(sinfo: StructInfo, shape_var_map: Dict[Var, PrimExpr] | None = None, var_map: Dict[Var, RelaxExpr] | None = None) → StructInfo

将 sinfo 擦除为良好定义的形式。

此函数删除 StructInfo 对给定映射中未定义的形状和变量的依赖性。

参数:

  • sinfo (StructInfo) – 输入结构信息。

  • shape_var_map (Dict[tir.Var, tir.PrimExpr]) – 指定已定义的形状变量以及它们应映射到的值。

  • var_map (Dict[relax.Var, Expr]) – 指定已定义的变量以及它们应映射到的值。

返回:

ret – 相应的擦除后的结构信息。

返回类型:

StructInfo

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

获取在输入函数中使用但未定义的 TIR 变量。返回的列表是去重后的 - 每个 TIR 变量最多出现一次。

参数:

func (Function) – 要分析的函数对象。

返回:

ret – 在输入函数中使用但未定义的符号变量列表。

返回类型:

List[relax.Var]

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

从表达式 expr 返回所有自由变量。自由变量是在表达式中未被 VarBinding 或函数参数绑定的变量。 :param expr: 表达式。 :type expr: Expr

返回:

ret – expr 中自由变量的列表，按后序 DFS 顺序排列

返回类型:

List[relax.Var]

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

从 StructInfo 获取相应的静态类型。

参数:

sinfo (StructInfo) – 输入结构信息。

返回:

ret – 相应的静态类型。

返回类型:

Type

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

获取从 relax.Var 到 Expr 的映射，用于函数中的每个变量。

参数:

func (Function) – 要分析的输入函数。

返回:

从 relax.Var 到 Expr 的映射。

返回类型:

Dict[relax.Var, Expr]

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

检查给定的 PrimFunc 本质上是否在执行 reshape 操作。reshape 操作还包括 expand_dims、squeeze、flatten 等。

这里的允许的 reshape 模式是：例如，假设操作是 B[l_0, l_1, …, l_b] = A[r_0, r_1, …, r_a]，我们检查是否可以证明缓冲区 B 下 l_0, …, l_b 的扁平化索引等于缓冲区 A 下 r_0, …, r_a 的扁平化索引。

参数:

func (tir.PrimFunc) – 要检查的函数。

返回:

ret – 一个布尔值，指示给定的 PrimFunc 是否正在进行 reshape。

返回类型:

bool

注释

根据上面的描述，返回的结果只能是假阴性，不能是假阳性，因为每当我们无法证明相等性时，我们都返回 false。此属性保证了此函数的安全性。

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

返回从变量名到其绑定的映射。

tvm.relax.analysis.post_order_visit(expr, fvisit)

以后序 DFS 顺序递归访问 ir 节点，应用 fvisit。保证每个节点只被访问一次。

参数:

  • expr (tvm.relax.Expr) – 输入表达式。

  • fvisit (function) – 要应用的访问器函数。

tvm.relax.analysis.remove_all_unused(func: Function) → Function

它移除：1. DataflowBlock 中未使用的局部 VarBinding。 2. 函数中未使用的 DataflowBlock。

参数:

func (Function) – 要分析的输入函数。

注释

对于 IRModule 范围的 DCE，请使用 py:func:tvm.relax.transform.DeadCodeElimination。

返回:

删除了未使用变量的函数。

返回类型:

Function

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

运行基础检查以查看 base 是否包含 derived。

参数:
返回:

ret – 导出的返回值结构信息。

返回类型:

StructInfo

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

将两个结构信息统一到它们的最近公共祖先。

参数:
返回:

ret – 相应的 lca 结果。

返回类型:

StructInfo

tvm.relax.analysis.suggest_layout_transforms(func: PrimFunc, write_buffer_transforms: List[IndexMap | Callable]) → Dict[Block, Dict[Block | Buffer, IndexMap]]

建议 PrimFunc 中块和缓冲区的布局转换。

参数:

  • func (PrimFunc) – 将在其上执行分析并建议转换的 PrimFunc。

  • write_buffer_transforms (List[Union[IndexMap, Callable]) – 输出缓冲区上的布局转换列表。布局转换的数量必须与 PrimFunc 的输出数量匹配。

返回:

ret – func 中每个块的建议转换。对于每个块，返回的值是从对象（块或缓冲区）到其索引映射转换的映射。

返回类型:

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

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

获取输入结构信息中出现的 TIR 变量。返回的列表已去重 - 每个 TIR 变量最多出现一次。

参数:

sinfo (StructInfo) – 要分析的结构信息对象。

返回:

ret – 输入结构信息中出现的 TIR 变量列表。

返回类型:

List[tir.Var]

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

分析数据流块中的变量使用-定义链。

参数:

dfb (DataflowBlock) – 要分析的数据流块

返回:

从变量定义到其用途的映射。

返回类型:

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

tvm.relax.analysis.well_formed(obj: IRModule | Function, check_struct_info: bool = True) → bool

检查 IRModule 是否结构良好。

参数:

  • obj (Union[tvm.IRModule, Function]) – 输入的 IRModule 或 relax.Function。

  • check_struct_info (bool) – 一个布尔标志，指示是否将检查“每个 Expr 必须具有定义的结构信息”属性。

返回:

ret – 如果 IRModule 结构良好，则为 True；否则为 False。

返回类型:

bool

注意

默认情况下，始终检查结构信息。仅在测试用例中 check_struct_info 可能为 false，以便充分测试其他结构良好的要求，并且不会因没有结构信息而受阻。

tvm.relax.analysis.estimate_memory_usage(mod: IRModule | Function) → str

分析函数，用于估计 IRModule 中 Relax 函数的内存使用情况。估计包括内存规划前后需要分配的总内存大小。

结果可能被高估，因为估计是静态的，不考虑控制流（例如“if”和跨函数调用）。它只是累积每个 alloc_tensor 和 alloc_storage 的大小。

此分析函数用于演示内存规划的效果。

参数:

mod (Union[IRModule, Function]) – 要分析其内部函数的输入 IRModule。如果输入是一个 Function，我们将用一个 IRModule 包裹它，并将该函数命名为 “main”。

返回:

est – 估计信息，以字符串形式。

返回类型:

str

注释

我们将 “relax.memory.alloc_tensor/storage” 视为内存规划产生的结果。