23 #ifndef TVM_TIR_SCHEDULE_STATE_H_
24 #define TVM_TIR_SCHEDULE_STATE_H_
26 #include <tvm/ffi/reflection/registry.h>
31 #include <unordered_map>
108 std::unordered_map<StmtSRef, BlockInfo, ObjectPtrHash, ObjectPtrEqual>
block_info;
110 std::unordered_map<const StmtNode*, StmtSRef>
stmt2ref;
124 refl::ObjectDef<ScheduleStateNode>()
150 const Map<Block, Block>& block_sref_reuse);
159 static constexpr
const char*
_type_key =
"tir.ScheduleState";
Definition of two pillar data structure for TensorIR scheduling: StmtSRef, BlockScope.
Managed reference class to IRModuleNode.
Definition: module.h:257
Managed reference to BlockScopeNode.
Definition: block_scope.h:298
The state of scheduling, which exposes a Replace method as the primary interface for all the scheduli...
Definition: state.h:99
TVM_DECLARE_FINAL_OBJECT_INFO(ScheduleStateNode, Object)
bool IsAffineBlockBinding(const StmtSRef &block_sref) const
Check a cached flag indicating if the specific block has quasi-affine bindings.
Definition: state.h:184
std::unordered_map< const StmtNode *, StmtSRef > stmt2ref
The reverse mapping from block/for-loop to their corresponding srefs.
Definition: state.h:110
int debug_mask
Do extra correctness checking after the class creation and each time after calling the Replace method...
Definition: state.h:116
static constexpr const char * _type_key
Definition: state.h:159
BlockInfo GetBlockInfo(const StmtSRef &block_sref) const
Returns the BlockInfo correpsonding to the block sref.
void DebugVerify() const
Trigger the verification according to the debug_mask bitmask. 1) If the bitmask kVerifySRefTree is on...
bool enable_check
Whether to enable prequisite checks for schedule primitives.
Definition: state.h:120
IRModule mod
The AST of the module being scheduled.
Definition: state.h:102
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...
Definition: state.h:193
std::unordered_map< StmtSRef, BlockInfo, ObjectPtrHash, ObjectPtrEqual > block_info
Mapping from a block sref to its correpsonding BlockInfo, tracking the dependency inside the block sc...
Definition: state.h:108
void UpdateScopeBlockInfo(const Stmt &stmt)
Recalculate the BlockInfo recursively under stmt. If stmt is a Block itself, we will not reset its af...
static void RegisterReflection()
Definition: state.h:122
bool IsStagePipeline(const StmtSRef &scope_root) const
Check a cached flag indicating if a block scope is an equivalence of a stage pipeline.
Definition: state.h:201
BlockScope GetBlockScope(const StmtSRef &scope_root) const
Get the BlockScope correpsonding to the sref of scope root block.
Definition: state.h:176
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,...
Managed reference to ScheduleStateNode.
Definition: state.h:210
ScheduleState(IRModule mod, int debug_mask=0, bool enable_check=true)
Construct a schedule state from an IRModule.
ScheduleStateNode * get() const
Definition: state.h:222
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ScheduleState, ObjectRef, ScheduleStateNode)
Managed reference to StmtSRefNode.
Definition: block_scope.h:106
Container of all statements.
Definition: stmt.h:64
IRModule that holds the functions and type definitions.
Definition: repr_printer.h:91
ScheduleDebugMask
The bitmask of the debug flag in the ScheduleStateNode.
Definition: state.h:80
@ kVerifySRefTree
Verify the correctness of the sref tree.
Definition: state.h:82
@ kVerifyCachedFlags
Verify the correctness of affine_binding, region_cover and stage_pipeline.
Definition: state.h:84
tvm::PrimExpr mod(const tvm::PrimExpr &a, const tvm::PrimExpr &b)
Definition: broadcast.h:306
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
The information about a TensorIR block, it contains two categories of information 1) Info on the bloc...
Definition: state.h:44
BlockScope scope
Property of a block scope rooted at the block, storing dependencies in the scope.
Definition: state.h:46
BlockInfo(BlockScope scope, bool affine_binding=false, bool region_cover=false, bool stage_pipeline=false)
Definition: state.h:68
bool stage_pipeline
This property indicates that the block scope (rooted at its corresponding block) is equivalent to of ...
Definition: state.h:64
bool affine_binding
Property of a block, indicating the block realization binding is quasi-affine.
Definition: state.h:49
bool region_cover
Property of a block, indicating each of the block's read regions is fully produced by its producers.
Definition: state.h:54