The state of scheduling, which exposes a Replace method as the primary interface for all the scheduling primitives to manipulate the TensorIR. More...

#include <state.h>

Inheritance diagram for tvm::tir::ScheduleStateNode:

Collaboration diagram for tvm::tir::ScheduleStateNode:

Public Member Functions
void	VisitAttrs (AttrVisitor *v)

void	Replace (const tir::StmtSRef &src_sref, const Stmt &tgt_stmt, const Map< Block, Block > &block_sref_reuse)
	Replace the part of the AST, as being pointed to by `src_sref`, with a specific statement `tgt_stmt`, and maintain the sref tree accordingly. Replace will try to perform copy on write as much as possible when the ScheduleState holds the only copy to the IRModule and IR nodes. More...

void	DebugVerify () const
	Trigger the verification according to the `debug_mask` bitmask. 1) If the bitmask `kVerifySRefTree` is on, verify the correctness of the sref tree. 2) If the bitmask `kVerifyCachedFlags` is on, verify the correctness of `affine_binding`, `region_cover` and `stage_pipeline` More...

	TVM_DECLARE_FINAL_OBJECT_INFO (ScheduleStateNode, Object)

BlockInfo	GetBlockInfo (const StmtSRef &block_sref) const
	Returns the BlockInfo correpsonding to the block sref. More...

void	UpdateScopeBlockInfo (const Stmt &stmt)
	Recalculate the BlockInfo recursively under stmt. If stmt is a Block itself, we will not reset its affine binding flag unless it doesn't have block vars, since the affine flag depends on the outer scope of stmt. More...

BlockScope	GetBlockScope (const StmtSRef &scope_root) const
	Get the BlockScope correpsonding to the sref of scope root block. More...

bool	IsAffineBlockBinding (const StmtSRef &block_sref) const
	Check a cached flag indicating if the specific block has quasi-affine bindings. More...

bool	IsRegionCoveredConsumer (const StmtSRef &consumer_block_sref) const
	Check a cached flag indicating if each of the specific consumer block's read region is fully produced by its producers. More...

bool	IsStagePipeline (const StmtSRef &scope_root) const
	Check a cached flag indicating if a block scope is an equivalence of a stage pipeline. More...

Public Member Functions inherited from tvm::runtime::Object
uint32_t	type_index () const

std::string	GetTypeKey () const

size_t	GetTypeKeyHash () const

template<typename TargetType >
bool	IsInstance () const

bool	unique () const

	Object ()

	Object (const Object &other)

	Object (Object &&other)

Object &	operator= (const Object &other)

Object &	operator= (Object &&other)

Public Attributes
IRModule	mod
	The AST of the module being scheduled. More...

std::unordered_map< StmtSRef, BlockInfo, ObjectPtrHash, ObjectPtrEqual >	block_info
	Mapping from a block sref to its correpsonding BlockInfo, tracking the dependency inside the block scope, and storing necessary information flags for scheduling. More...

std::unordered_map< const StmtNode *, StmtSRef >	stmt2ref
	The reverse mapping from block/for-loop to their corresponding srefs. More...

int	debug_mask
	Do extra correctness checking after the class creation and each time after calling the Replace method. More...

bool	enable_check
	Whether to enable prequisite checks for schedule primitives. More...

Static Public Attributes
static constexpr const char *	_type_key = "tir.ScheduleState"

Static Public Attributes inherited from tvm::runtime::Object
static constexpr const char *	_type_key = "runtime.Object"

static constexpr bool	_type_final = false

static constexpr uint32_t	_type_child_slots = 0

static constexpr bool	_type_child_slots_can_overflow = true

static constexpr bool	_type_has_method_visit_attrs = true

static constexpr bool	_type_has_method_sequal_reduce = false

static constexpr bool	_type_has_method_shash_reduce = false

static constexpr uint32_t	_type_index = TypeIndex::kDynamic

Additional Inherited Members
Public Types inherited from tvm::runtime::Object
typedef void(*	FDeleter) (Object *self)
	Object deleter. More...

using	RefCounterType = std::atomic< int32_t >

Static Public Member Functions inherited from tvm::runtime::Object
static std::string	TypeIndex2Key (uint32_t tindex)
	Get the type key of the corresponding index from runtime. More...

static size_t	TypeIndex2KeyHash (uint32_t tindex)
	Get the type key hash of the corresponding index from runtime. More...

static uint32_t	TypeKey2Index (const std::string &key)
	Get the type index of the corresponding key from runtime. More...

static uint32_t	_GetOrAllocRuntimeTypeIndex ()

static uint32_t	RuntimeTypeIndex ()

Protected Member Functions inherited from tvm::runtime::Object
void	IncRef ()
	developer function, increases reference counter. More...

void	DecRef ()
	developer function, decrease reference counter. More...

Static Protected Member Functions inherited from tvm::runtime::Object
static uint32_t	GetOrAllocRuntimeTypeIndex (const std::string &key, uint32_t static_tindex, uint32_t parent_tindex, uint32_t type_child_slots, bool type_child_slots_can_overflow)
	Get the type index using type key. More...

Protected Attributes inherited from tvm::runtime::Object
uint32_t	type_index_ {0}
	Type index(tag) that indicates the type of the object. More...

RefCounterType	ref_counter_ {0}
	The internal reference counter. More...

FDeleter	deleter_ = nullptr
	deleter of this object to enable customized allocation. If the deleter is nullptr, no deletion will be performed. The creator of the object must always set the deleter field properly. More...

Detailed Description

The state of scheduling, which exposes a Replace method as the primary interface for all the scheduling primitives to manipulate the TensorIR.

The data structure contains the following information 1) The AST being scheduled (mod) 2) The sref tree of schedulable statements (indicated by the srefs) 3) The dependency information of each block scope (block_info) 4) A reverse mapping from the AST nodes to that in the sref tree (stmt2ref) 5) A debug flag, if set, extra checking is enabled (debug_mask) 6) A check flag, if set, enable prequisite check for schedule primitives (enable_check)

Member Function Documentation

◆ DebugVerify()

void tvm::tir::ScheduleStateNode::DebugVerify ( ) const

Trigger the verification according to the debug_mask bitmask. 1) If the bitmask kVerifySRefTree is on, verify the correctness of the sref tree. 2) If the bitmask kVerifyCachedFlags is on, verify the correctness of affine_binding, region_cover and stage_pipeline

◆ GetBlockInfo()

BlockInfo tvm::tir::ScheduleStateNode::GetBlockInfo ( const StmtSRef & block_sref ) const

Returns the BlockInfo correpsonding to the block sref.

◆ GetBlockScope()

BlockScope tvm::tir::ScheduleStateNode::GetBlockScope ( const StmtSRef & scope_root ) const

inline

Get the BlockScope correpsonding to the sref of scope root block.

Parameters

scope_root The block sref to be retrieved

Returns: The corresponding BlockScope

◆ IsAffineBlockBinding()

bool tvm::tir::ScheduleStateNode::IsAffineBlockBinding ( const StmtSRef & block_sref ) const

inline

Check a cached flag indicating if the specific block has quasi-affine bindings.

Parameters

block_sref The block sref to be checked

Returns: A boolean flag indicating if the block has quasi-affine bindings

◆ IsRegionCoveredConsumer()

bool tvm::tir::ScheduleStateNode::IsRegionCoveredConsumer ( const StmtSRef & consumer_block_sref ) const

inline

Check a cached flag indicating if each of the specific consumer block's read region is fully produced by its producers.

Parameters

consumer_block_sref The specific consumer block

Returns: A boolean flag indicating if the block has quasi-affine bindings

◆ IsStagePipeline()

bool tvm::tir::ScheduleStateNode::IsStagePipeline ( const StmtSRef & scope_root ) const

inline

Check a cached flag indicating if a block scope is an equivalence of a stage pipeline.

Parameters

scope_root The block sref to be retrieved

Returns: The corresponding BlockScope

◆ Replace()

void tvm::tir::ScheduleStateNode::Replace	(	const tir::StmtSRef &	src_sref,
		const Stmt &	tgt_stmt,
		const Map< Block, Block > &	block_sref_reuse
	)

Replace the part of the AST, as being pointed to by src_sref, with a specific statement tgt_stmt, and maintain the sref tree accordingly. Replace will try to perform copy on write as much as possible when the ScheduleState holds the only copy to the IRModule and IR nodes.

Only 3 types of replacements are allowed: from src_sref->stmt to tgt_stmt. 1) Block -> Block 2) Loop -> Loop 3) Loop -> BlockRealize

Parameters

src_sref	The sref to the statement to be replaced
tgt_stmt	The statement to be replaced in
block_sref_reuse	Maps an old block (to be replaced in the subtree under `src_sref->stmt`) to a new block (replaced to, in the subtree under `tgt_stmt`), and enforces reuse of srefs between them (rather than create new srefs) i.e. after being replaced, the sref that points to the old block will point to the new one

Note: The reuse of loop srefs are detected automatically according to the reuse of loop vars.

◆ TVM_DECLARE_FINAL_OBJECT_INFO()

tvm::tir::ScheduleStateNode::TVM_DECLARE_FINAL_OBJECT_INFO	(	ScheduleStateNode	,
		Object
	)

◆ UpdateScopeBlockInfo()

void tvm::tir::ScheduleStateNode::UpdateScopeBlockInfo ( const Stmt & stmt )

Recalculate the BlockInfo recursively under stmt. If stmt is a Block itself, we will not reset its affine binding flag unless it doesn't have block vars, since the affine flag depends on the outer scope of stmt.

◆ VisitAttrs()

void tvm::tir::ScheduleStateNode::VisitAttrs ( AttrVisitor * v )

inline

Member Data Documentation

◆ _type_key

constexpr const char* tvm::tir::ScheduleStateNode::_type_key = "tir.ScheduleState"

staticconstexpr

◆ block_info

std::unordered_map<StmtSRef, BlockInfo, ObjectPtrHash, ObjectPtrEqual> tvm::tir::ScheduleStateNode::block_info

Mapping from a block sref to its correpsonding BlockInfo, tracking the dependency inside the block scope, and storing necessary information flags for scheduling.

◆ debug_mask

int tvm::tir::ScheduleStateNode::debug_mask

Do extra correctness checking after the class creation and each time after calling the Replace method.

See also: ScheduleDebugMask

◆ enable_check

bool tvm::tir::ScheduleStateNode::enable_check

Whether to enable prequisite checks for schedule primitives.

◆ mod

IRModule tvm::tir::ScheduleStateNode::mod

The AST of the module being scheduled.

◆ stmt2ref

std::unordered_map<const StmtNode*, StmtSRef> tvm::tir::ScheduleStateNode::stmt2ref

The reverse mapping from block/for-loop to their corresponding srefs.

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

include/tvm/tir/schedule/state.h

Public Member Functions

Public Attributes

Static Public Attributes

Additional Inherited Members

Detailed Description

Member Function Documentation

◆ DebugVerify()

◆ GetBlockInfo()

◆ GetBlockScope()

◆ IsAffineBlockBinding()

◆ IsRegionCoveredConsumer()

◆ IsStagePipeline()

◆ Replace()

◆ TVM_DECLARE_FINAL_OBJECT_INFO()

◆ UpdateScopeBlockInfo()

◆ VisitAttrs()

Member Data Documentation

◆ _type_key

◆ block_info

◆ debug_mask

◆ enable_check

◆ mod

◆ stmt2ref