tvm.tir.stmt_functor

用于 IR 转换的语句 functor 实用工具

tvm.tir.stmt_functor.ir_transform(stmt, preorder, postorder, only_enable=None)

以后序 DFS 顺序递归访问和转换 ir 节点。

参数:

• stmt (tvm.tir.Stmt) – 要转换的输入。

• preorder (function) – 在递归突变之前调用的函数。如果 preorder 返回 None，则转换将继续进行递归调用。如果 preorder 返回非 None 的 tvm.tir.Stmt/Expr，则转换器将直接返回它，并且不会进行进一步的递归。

• postorder (function) – 在递归突变之后调用的函数。

• only_enable (Optional[List[str]]) – 我们仅启用的类型列表。

返回值:

result – 结果。

返回类型:

tvm.tir.Stmt

tvm.tir.stmt_functor.post_order_visit(stmt, fvisit)

以后序 DFS 顺序递归访问 ir 节点，应用 fvisit

保证每个节点只被访问一次。

参数:

fvisit (function) – 访问器函数。

tvm.tir.stmt_functor.pre_order_visit(stmt, fvisit)

对 stmt AST 进行递归先序访问，对每个节点应用 fvisit。

如果 fvisit 返回 False，则不会访问该节点的子节点。

参数:

fvisit (function of the signature Object -> bool) – 访问器函数。

tvm.tir.stmt_functor.substitute(node, vmap)

替换 vmap 指定的 var。

参数:

• node (ObjectRef) – 输入。

• vmap (Dict[tir.Var, PrimExpr]) – 变量映射。

返回值:

result – 结果。

返回类型:

tvm.tir.Stmt

tvm.tir.stmt_functor.renew_defs(func: PrimFunc)

为 TIR 重新生成定义节点，包括 VarDef、BufferDef。此 pass 的作用类似于简单的 DeepCopy，以复制具有不同 Var 和 Buffer 但行为相同的函数

参数:

func (PrimFunc) – 输入函数

返回值:

result – 新生成的函数。

返回类型:

PrimFunc