tvm.tir.stmt_functor#

Statement functor utilities for IR transformations

tvm.tir.stmt_functor.ir_transform(stmt, preorder, postorder, only_enable=None)[源代码]#

Recursively visit and transform ir nodes in post DFS order.

Parameters#

stmttvm.tir.Stmt

The input to be transformed.

preorder: function

The function called in before recursive mutation If preorder returns None, then the transform will proceed to recursive call. If preorder returns a not None tvm.tir.Stmt/Expr, the transformer will simply return it and won't do further recursion.

postorderfunction

The function called after recursive mutation.

only_enableOptional[List[str]]

List of types that we only enable.

Returns#

resulttvm.tir.Stmt

The result.

tvm.tir.stmt_functor.post_order_visit(stmt, fvisit)[源代码]#
Recursively visit the ir in post DFS order node, apply fvisit

Each node is guaranteed to be visited only once.

Parameters#

fvisit: function

The visitor function.

tvm.tir.stmt_functor.pre_order_visit(stmt, fvisit)[源代码]#
Recursive pre-order visit on stmt AST, applying fvisit on each node.

If fvisit returns False, it won't visit the children of the node.

Parameters#

fvisit: function of the signature Object -> bool

The visitor function.

tvm.tir.stmt_functor.renew_defs(func)[源代码]#

Re-generate the definition nodes for a TIR, including VarDef, BufferDef. This pass works as a simple DeepCopy to duplicate a function with different Vars and Buffers but the same behavior

Parameters#

func: PrimFunc

The input function

Returns#

resultPrimFunc

The new generated func.

参数:

func (PrimFunc)

tvm.tir.stmt_functor.substitute(node, vmap)[源代码]#

Substitute the var specified by vmap.

Parameters#

node: ObjectRef

The input.

vmapDict[Var, PrimExpr]

The variable mapping.

Returns#

resulttvm.tir.Stmt

The result.