tvm::tirx::StmtMutator Class Reference

StmtMutator that mutates the statements. More...

#include <stmt_functor.h>

Inheritance diagram for tvm::tirx::StmtMutator:

Collaboration diagram for tvm::tirx::StmtMutator:

Public Member Functions
Stmt	operator() (Stmt stmt)
	Mutate stmt.

Protected Member Functions
template<typename TNode >
ffi::ObjectPtr< TNode >	CopyOnWrite (const TNode *node)
	Perform copy on write on node.
Stmt	VisitStmt (const Stmt &stmt) override
	Internal mutator that everyone calls.
virtual Expr	VisitExpr (const Expr &e)
	Visitor to Exprs, can be overriden to do recursive changes to Exprs.
PrimExpr	VisitPrimExpr (const PrimExpr &e)
	Mutate a primitive expression and verify that it remains primitive.
virtual BufferVar	VisitBufferDef (const BufferVar &buffer, bool alloc_data)
	Visit buffer at definition site. Visits shape/strides/elem_offset via VisitExpr. If any field changes, creates a new buffer and records it in buffer_remap_.
virtual BufferVar	VisitBufferUse (const BufferVar &buffer)
	Visit buffer at use site (BufferStore, BufferLoad, SBlock reads/writes). By default, returns the remapped buffer from buffer_remap_ if exists, otherwise returns the original buffer. BufferVar fields are visited at their definition site.
Stmt	VisitStmt_ (const BindNode *op) override
Stmt	VisitStmt_ (const AttrStmtNode *op) override
Stmt	VisitStmt_ (const IfThenElseNode *op) override
Stmt	VisitStmt_ (const ForNode *op) override
Stmt	VisitStmt_ (const WhileNode *op) override
Stmt	VisitStmt_ (const ReturnNode *op) override
Stmt	VisitStmt_ (const BreakNode *op) override
Stmt	VisitStmt_ (const ContinueNode *op) override
Stmt	VisitStmt_ (const AllocBufferNode *op) override
Stmt	VisitStmt_ (const DeclBufferNode *op) override
Stmt	VisitStmt_ (const BufferStoreNode *op) override
Stmt	VisitStmt_ (const AssertStmtNode *op) override
Stmt	VisitStmt_ (const SeqStmtNode *op) override
Stmt	VisitStmt_ (const EvaluateNode *op) override
Stmt	VisitStmt_ (const SBlockNode *op) override
Stmt	VisitStmt_ (const SBlockRealizeNode *op) override
Stmt	VisitStmt_ (const ScopeIdDefStmtNode *op) override
Stmt	VisitStmt_ (const tirx::TilePrimitiveCallNode *op) override
Stmt	VisitSeqStmt_ (const SeqStmtNode *op, bool flatten_before_visit, std::function< Stmt(const Stmt &)> fmutate=nullptr)
	Alternative advance method for SeqStmtNode.

Protected Attributes
ffi::Map< BufferVar, BufferVar >	buffer_remap_
	Map from old buffer to new buffer, populated by VisitBufferDef.
bool	allow_copy_on_write_ {false}
	Internal state to indicate whether copy on write is enabled. COW is enabled iff all the parents of the node are unique.

Detailed Description

StmtMutator that mutates the statements.

Member Function Documentation

CopyOnWrite()

template<typename TNode >

ffi::ObjectPtr< TNode > tvm::tirx::StmtMutator::CopyOnWrite ( const TNode *  node )

inlineprotected

Perform copy on write on node.

If CopyOnWrite is allowed, directly return a strong reference to the node container. Otherwise, return a copy of the node.

Returns

The result object pointer.

operator()()

Stmt tvm::tirx::StmtMutator::operator() ( Stmt  stmt )

inline

Mutate stmt.

Parameters

stmt	The input statement to be mutated.

Returns

The result of the call

Note

It is important that stmt is passed by value. so copy on write can be triggered correctly. do mutator(std::move(stmt)) or when copy elison is triggered.

VisitBufferDef()

virtual BufferVar tvm::tirx::StmtMutator::VisitBufferDef ( const BufferVar &  buffer, bool  alloc_data  )

protectedvirtual

Visit buffer at definition site. Visits shape/strides/elem_offset via VisitExpr. If any field changes, creates a new buffer and records it in buffer_remap_.

Parameters

buffer	The buffer being defined.
alloc_data	If true, the buffer's data pointer is a new allocation (AllocBuffer); if false, data references an existing variable (DeclBuffer).

Returns

The (possibly new) buffer.

VisitBufferUse()

virtual BufferVar tvm::tirx::StmtMutator::VisitBufferUse ( const BufferVar &  buffer )

protectedvirtual

Visit buffer at use site (BufferStore, BufferLoad, SBlock reads/writes). By default, returns the remapped buffer from buffer_remap_ if exists, otherwise returns the original buffer. BufferVar fields are visited at their definition site.

Returns

The (possibly remapped) buffer.

VisitExpr()

virtual Expr tvm::tirx::StmtMutator::VisitExpr ( const Expr &  e )

inlineprotectedvirtual

Visitor to Exprs, can be overriden to do recursive changes to Exprs.

Note

A common pattern is to call ExprMutator here, or have a class sub-class both StmtMutator and ExprMutator and redirect Mutate to ExprMutator::Mutate(Expr)

Reimplemented in tvm::tirx::StmtExprMutator.

VisitPrimExpr()

PrimExpr tvm::tirx::StmtMutator::VisitPrimExpr ( const PrimExpr &  e )

inlineprotected

Mutate a primitive expression and verify that it remains primitive.

VisitSeqStmt_()

Stmt tvm::tirx::StmtMutator::VisitSeqStmt_ ( const SeqStmtNode *  op, bool  flatten_before_visit, std::function< Stmt(const Stmt &)>  fmutate = nullptr  )

protected

Alternative advance method for SeqStmtNode.

This function can be called when a child class override VisitStmt_(const SeqStmtNode*) to introduce the special behavior to visit

Parameters

op	The sequence.
flatten_before_visit	Whether to flatten the sequence before visit.
fmutate	The mutate function, can be nullptr, which defaults to Visit.

Returns

The mutated result.

VisitStmt()

Stmt tvm::tirx::StmtMutator::VisitStmt ( const Stmt &  stmt )

inlineoverrideprotected

Internal mutator that everyone calls.

Note

To override mutate's behavior, override VisitExpr instead.

Parameters

stmt	The input stmt.

Returns

The mutated results.

VisitStmt_() [1/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const AllocBufferNode *  op )

overrideprotected

VisitStmt_() [2/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const AssertStmtNode *  op )

overrideprotected

VisitStmt_() [3/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const AttrStmtNode *  op )

overrideprotected

VisitStmt_() [4/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const BindNode *  op )

overrideprotected

VisitStmt_() [5/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const BreakNode *  op )

overrideprotected

VisitStmt_() [6/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const BufferStoreNode *  op )

overrideprotected

VisitStmt_() [7/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const ContinueNode *  op )

overrideprotected

VisitStmt_() [8/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const DeclBufferNode *  op )

overrideprotected

VisitStmt_() [9/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const EvaluateNode *  op )

overrideprotected

VisitStmt_() [10/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const ForNode *  op )

overrideprotected

VisitStmt_() [11/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const IfThenElseNode *  op )

overrideprotected

VisitStmt_() [12/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const ReturnNode *  op )

overrideprotected

VisitStmt_() [13/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const SBlockNode *  op )

overrideprotected

VisitStmt_() [14/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const SBlockRealizeNode *  op )

overrideprotected

VisitStmt_() [15/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const ScopeIdDefStmtNode *  op )

overrideprotected

VisitStmt_() [16/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const SeqStmtNode *  op )

overrideprotected

VisitStmt_() [17/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const tirx::TilePrimitiveCallNode *  op )

overrideprotected

VisitStmt_() [18/18]

Stmt tvm::tirx::StmtMutator::VisitStmt_ ( const WhileNode *  op )

overrideprotected

Member Data Documentation

allow_copy_on_write_

bool tvm::tirx::StmtMutator::allow_copy_on_write_ {false}

protected

Internal state to indicate whether copy on write is enabled. COW is enabled iff all the parents of the node are unique.

buffer_remap_

ffi::Map<BufferVar, BufferVar> tvm::tirx::StmtMutator::buffer_remap_

protected

Map from old buffer to new buffer, populated by VisitBufferDef.

---

The documentation for this class was generated from the following file:

• include/tvm/tirx/stmt_functor.h