tvm.tirx.stmt_functor

tvm.tirx.stmt_functor#

Statement functor utilities for IR transformations

class tvm.tirx.stmt_functor.StmtFunctor#

An abstract visitor over Statement, with visiting functions defined for each Stmt type.

visit_stmt(stmt)#

Apply the visitor to a statement.

Parameters:: stmt (tvm.tirx.Stmt) – The statement to be visited.
Returns:: result – The result of the visit.
Return type:: Any

visit_stmt_default_(op)#: Default visitor implementation for statements.

visit_bind_(op)#: Visitor for Bind nodes.

visit_attr_(op)#: Visitor for AttrStmt nodes.

visit_if_then_else_(op)#: Visitor for IfThenElse nodes.

visit_for_(op)#: Visitor for For nodes.

visit_while_(op)#: Visitor for While nodes.

visit_break_(op)#: Visitor for Break nodes.

visit_continue_(op)#: Visitor for Continue nodes.

visit_allocate_(op)#: Visitor for Allocate nodes.

visit_allocate_const_(op)#: Visitor for AllocateConst nodes.

visit_decl_buffer_(op)#: Visitor for DeclBuffer nodes.

visit_buffer_store_(op)#: Visitor for BufferStore nodes.

visit_buffer_realize_(op)#: Visitor for BufferRealize nodes.

visit_assert_(op)#: Visitor for AssertStmt nodes.

visit_producer_store_(op)#: Visitor for ProducerStore nodes.

visit_producer_realize_(op)#: Visitor for ProducerRealize nodes.

visit_prefetch_(op)#: Visitor for Prefetch nodes.

visit_seqstmt_(op)#: Visitor for SeqStmt nodes.

visit_evaluate_(op)#: Visitor for Evaluate nodes.

visit_block_(op)#: Visitor for Block nodes.

visit_block_realize_(op)#: Visitor for BlockRealize nodes.

visit_scope_id_def_stmt_(op)#: Visitor for ScopeIdDefStmt nodes.

visit_op_call_(op)#: Visitor for TilePrimitiveCall nodes.

visit_buffer_region_(op)#: Visitor for BufferRegion nodes.

visit_alloc_buffer_(op)#: Visitor for AllocBuffer nodes.

class tvm.tirx.stmt_functor.StmtVisitor#

A visitor over Stmt.

This is a visitor that recursively traverses a statement. Subclasses can override the visit methods to customize the behavior.

visit_expr(expr)#

Visit expressions that occur in a statement.

This method can be overridden to implement expression traversal in a statement visitor.

Parameters:: expr (Expr) – The expression to be visited.

visit_bind_(op)#: Visitor implementation for Bind.

visit_attr_(op)#: Visitor implementation for AttrStmt.

visit_if_then_else_(op)#: Visitor implementation for IfThenElse.

visit_for_(op)#: Visitor implementation for For.

visit_while_(op)#: Visitor implementation for While.

visit_break_(op)#: Visitor implementation for Break.

visit_continue_(op)#: Visitor implementation for Continue.

visit_allocate_(op)#: Visitor implementation for Allocate.

visit_allocate_const_(op)#: Visitor implementation for AllocateConst.

visit_decl_buffer_(op)#: Visitor implementation for DeclBuffer.

visit_buffer_store_(op)#: Visitor implementation for BufferStore.

visit_assert_(op)#: Visitor implementation for AssertStmt.

visit_seqstmt_(op)#: Visitor implementation for SeqStmt.

visit_evaluate_(op)#: Visitor implementation for Evaluate.

visit_block_(op)#: Visitor implementation for Block.

visit_block_realize_(op)#: Visitor implementation for BlockRealize.

visit_scope_id_def_stmt_(op)#

Visitor implementation for ScopeIdDefStmt.

Mirrors the C++ visitor: walk extents and preferred_extents via visit_expr; there is no body to recurse into (the def vars themselves are leaves the visitor doesn’t otherwise inspect).

visit_op_call_(op)#: Visitor implementation for TilePrimitiveCall.

visit_buffer_region_(op)#: Visitor implementation for BufferRegion.

visit_alloc_buffer_(op)#: Visitor implementation for AllocBuffer.

class tvm.tirx.stmt_functor.StmtMutator#

A mutator over Stmt.

This is a mutator that recursively transforms a statement. Subclasses can override the visit methods to customize the behavior.

visit_expr(expr)#

Visit and mutate expressions that occur in a statement.

This method can be overridden to implement expression mutation in a statement mutator.

Parameters:: expr (Expr) – The expression to be visited.
Returns:: result – The mutated expression.
Return type:: Expr

visit_bind_(op)#: Mutator implementation for Bind.

visit_attr_(op)#: Mutator implementation for AttrStmt.

visit_if_then_else_(op)#: Mutator implementation for IfThenElse.

visit_for_(op)#: Mutator implementation for For.

visit_while_(op)#: Mutator implementation for While.

visit_break_(op)#: Mutator implementation for Break.

visit_continue_(op)#: Mutator implementation for Continue.

visit_allocate_(op)#: Mutator implementation for Allocate.

visit_allocate_const_(op)#: Mutator implementation for AllocateConst.

visit_decl_buffer_(op)#: Mutator implementation for DeclBuffer.

visit_buffer_store_(op)#: Mutator implementation for BufferStore.

visit_buffer_realize_(op)#: Mutator implementation for BufferRealize.

visit_assert_(op)#: Mutator implementation for AssertStmt.

visit_producer_store_(op)#: Mutator implementation for ProducerStore.

visit_producer_realize_(op)#: Mutator implementation for ProducerRealize.

visit_prefetch_(op)#: Mutator implementation for Prefetch.

visit_seqstmt_(op)#: Mutator implementation for SeqStmt.

visit_evaluate_(op)#: Mutator implementation for Evaluate.

visit_block_(op)#: Mutator implementation for Block.

visit_block_realize_(op)#: Mutator implementation for BlockRealize.

visit_scope_id_def_stmt_(op)#

Mutator implementation for ScopeIdDefStmt.

Mirrors the C++ mutator: rewrite extents and preferred_extents via visit_expr. Deferred-extent defs (extents is None) and unchanged extents pass through.

visit_op_call_(op)#: Mutator implementation for TilePrimitiveCall.

visit_buffer_region_(op)#: Mutator implementation for BufferRegion.

visit_alloc_buffer_(op)#: Mutator implementation for AllocBuffer.

class tvm.tirx.stmt_functor.StmtExprVisitor#

A visitor over both statements and expressions.

This class inherits from both StmtVisitor and ExprVisitor to recursively visit both statements and expressions.

visit_expr(expr)#

Visit an expression used in a statement.

Parameters:: expr (Expr) – The expression to be visited.

class tvm.tirx.stmt_functor.StmtExprMutator#

A mutator over both statements and expressions.

This class inherits from both StmtMutator and ExprMutator to recursively transform both statements and expressions.

visit_expr(expr)#

Mutate an expression used in a statement.

Parameters:: expr (Expr) – The expression to be mutated.
Returns:: result – The mutated expression.
Return type:: Expr

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

Recursively visit and transform ir nodes in post DFS order.

Parameters:

stmt (tvm.tirx.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.tirx.Stmt/Expr, the transformer will simply return it and won’t do further recursion.
postorder (function) – The function called after recursive mutation.
only_enable (Optional[List[str]]) – List of types that we only enable.

Returns:

result – The result.

Return type:

tvm.tirx.Stmt

tvm.tirx.stmt_functor.post_order_visit(node, fvisit)#

Recursively visit a statement or expression in post DFS order, applying fvisit.: Each node is guaranteed to be visited only once.

Parameters:

node (tvm.tirx.Stmt or tvm.ir.Expr) – The statement or expression to visit.
fvisit (function) – The visitor function.

tvm.tirx.stmt_functor.pre_order_visit(node, fvisit)#

Recursively visit a statement or expression in pre-order, applying fvisit.: If fvisit returns False, it won’t visit the children of the node.

Parameters:

node (tvm.tirx.Stmt or tvm.ir.Expr) – The statement or expression to visit.
fvisit (function of the signature Object -> bool) – The visitor function.

tvm.tirx.stmt_functor.substitute(node, vmap)#

Substitute the var specified by vmap.

Parameters:

node (ObjectRef) – The input.
vmap (Dict[tirx.Var, Expr]) – The variable mapping.

Returns:

result – The result.

Return type:

tvm.tirx.Stmt

tvm.tirx.stmt_functor.renew_defs(func: PrimFunc)#

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:: result – The new generated func.
Return type:: PrimFunc

tvm.tirx.stmt_functor

Contents

tvm.tirx.stmt_functor#