tvm::s_tir::ScheduleStateNode Class Reference

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::s_tir::ScheduleStateNode:

Collaboration diagram for tvm::s_tir::ScheduleStateNode:

Public Member Functions
void	Replace (const tir::StmtSRef &src_sref, const Stmt &tgt_stmt, const ffi::Map< SBlock, SBlock > &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_FFI_DECLARE_OBJECT_INFO_FINAL ("s_tir.ScheduleState", ScheduleStateNode, Object)
SBlockInfo	GetSBlockInfo (const StmtSRef &block_sref) const
	Returns the SBlockInfo correpsonding to the block sref. More...
void	UpdateScopeSBlockInfo (const Stmt &stmt)
	Recalculate the SBlockInfo recursively under stmt. If stmt is a SBlock 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...
SBlockScope	GetSBlockScope (const StmtSRef &scope_root) const
	Get the SBlockScope 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...

Static Public Member Functions
static void	RegisterReflection ()

Public Attributes
IRModule	mod
	The AST of the module being scheduled. More...
std::unordered_map< StmtSRef, SBlockInfo, ObjectPtrHash, ObjectPtrEqual >	block_info
	Mapping from a block sref to its correpsonding SBlockInfo, 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 bool	_type_mutable = true

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::s_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

GetSBlockInfo()

SBlockInfo tvm::s_tir::ScheduleStateNode::GetSBlockInfo

(

const StmtSRef &

block_sref

)

const

Returns the SBlockInfo correpsonding to the block sref.

GetSBlockScope()

SBlockScope tvm::s_tir::ScheduleStateNode::GetSBlockScope ( const StmtSRef &  scope_root ) const

inline

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

Parameters

scope_root

The block sref to be retrieved

Returns

The corresponding SBlockScope

IsAffineBlockBinding()

bool tvm::s_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::s_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::s_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 SBlockScope

RegisterReflection()

static void tvm::s_tir::ScheduleStateNode::RegisterReflection ( )

inlinestatic

Replace()

void tvm::s_tir::ScheduleStateNode::Replace	(	const tir::StmtSRef &	src_sref,
		const Stmt &	tgt_stmt,
		const ffi::Map< SBlock, SBlock > &	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) SBlock -> SBlock 2) Loop -> Loop 3) Loop -> SBlockRealize

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_FFI_DECLARE_OBJECT_INFO_FINAL()

tvm::s_tir::ScheduleStateNode::TVM_FFI_DECLARE_OBJECT_INFO_FINAL	(	"s_tir.ScheduleState"	,
		ScheduleStateNode	,
		Object
	)

UpdateScopeSBlockInfo()

void tvm::s_tir::ScheduleStateNode::UpdateScopeSBlockInfo

(

const Stmt &

stmt

)

Recalculate the SBlockInfo recursively under stmt. If stmt is a SBlock 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.

Member Data Documentation

_type_mutable

constexpr const bool tvm::s_tir::ScheduleStateNode::_type_mutable = true

staticconstexpr

block_info

std::unordered_map<StmtSRef, SBlockInfo, ObjectPtrHash, ObjectPtrEqual> tvm::s_tir::ScheduleStateNode::block_info

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

debug_mask

int tvm::s_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::s_tir::ScheduleStateNode::enable_check

Whether to enable prequisite checks for schedule primitives.

mod

IRModule tvm::s_tir::ScheduleStateNode::mod

The AST of the module being scheduled.

stmt2ref

std::unordered_map<const StmtNode*, StmtSRef> tvm::s_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/s_tir/schedule/state.h