tvm
state.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
23 #ifndef TVM_TIR_SCHEDULE_STATE_H_
24 #define TVM_TIR_SCHEDULE_STATE_H_
25 
26 #include <tvm/ffi/reflection/registry.h>
27 #include <tvm/ir/module.h>
28 #include <tvm/tir/block_scope.h>
29 #include <tvm/tir/function.h>
30 
31 #include <unordered_map>
32 #include <utility>
33 
34 namespace tvm {
35 namespace tir {
36 
44 struct BlockInfo {
46  BlockScope scope{nullptr};
47  // The properties below are information about the current block realization under its parent scope
49  bool affine_binding{false};
54  bool region_cover{false};
64  bool stage_pipeline{false};
65 
66  BlockInfo() = default;
67 
68  explicit BlockInfo(BlockScope scope, bool affine_binding = false, bool region_cover = false,
69  bool stage_pipeline = false)
70  : scope(std::move(scope)), //
74 };
75 
80 enum ScheduleDebugMask : uint32_t {
85 };
86 
99 class ScheduleStateNode : public Object {
100  public:
108  std::unordered_map<StmtSRef, BlockInfo, ObjectPtrHash, ObjectPtrEqual> block_info;
110  std::unordered_map<const StmtNode*, StmtSRef> stmt2ref;
121 
122  static void RegisterReflection() {
123  namespace refl = tvm::ffi::reflection;
124  refl::ObjectDef<ScheduleStateNode>()
125  .def_ro("mod", &ScheduleStateNode::mod)
126  .def_ro("debug_mask", &ScheduleStateNode::debug_mask)
127  .def_ro("enable_check", &ScheduleStateNode::enable_check);
128  }
129 
149  TVM_DLL void Replace(const tir::StmtSRef& src_sref, const Stmt& tgt_stmt,
150  const Map<Block, Block>& block_sref_reuse);
157  TVM_DLL void DebugVerify() const;
158 
159  static constexpr const char* _type_key = "tir.ScheduleState";
161 
162  /******** Property of blocks ********/
164  TVM_DLL BlockInfo GetBlockInfo(const StmtSRef& block_sref) const;
170  TVM_DLL void UpdateScopeBlockInfo(const Stmt& stmt);
176  BlockScope GetBlockScope(const StmtSRef& scope_root) const {
177  return GetBlockInfo(scope_root).scope;
178  }
184  bool IsAffineBlockBinding(const StmtSRef& block_sref) const {
185  return GetBlockInfo(block_sref).affine_binding;
186  }
193  bool IsRegionCoveredConsumer(const StmtSRef& consumer_block_sref) const {
194  return GetBlockInfo(consumer_block_sref).region_cover;
195  }
201  bool IsStagePipeline(const StmtSRef& scope_root) const {
202  return GetBlockInfo(scope_root).stage_pipeline;
203  }
204 };
205 
210 class ScheduleState : public ObjectRef {
211  public:
219  TVM_DLL explicit ScheduleState(IRModule mod, int debug_mask = 0, bool enable_check = true);
220 
222  ScheduleStateNode* get() const { return static_cast<ScheduleStateNode*>(data_.get()); }
223 
225 };
226 
227 } // namespace tir
228 } // namespace tvm
229 
230 #endif // TVM_TIR_SCHEDULE_STATE_H_
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
TIR Function.