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_exec_scope_stmt_(op)

Visitor for ExecScopeStmt 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 (PrimExpr) – 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_exec_scope_stmt_(op)

Visitor implementation for ExecScopeStmt.

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 (PrimExpr) – The expression to be visited.

Returns:

result – The mutated expression.

Return type:

PrimExpr

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_exec_scope_stmt_(op)

Mutator implementation for ExecScopeStmt.

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 (PrimExpr) – 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 (PrimExpr) – The expression to be mutated.

Returns:

result – The mutated expression.

Return type:

PrimExpr

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(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.tirx.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.tirx.stmt_functor.substitute(node, vmap)

Substitute the var specified by vmap.

Parameters:

• node (ObjectRef) – The input.

• vmap (Dict[tirx.Var, PrimExpr]) – 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