25 #ifndef TVM_TIR_BLOCK_SCOPE_H_
26 #define TVM_TIR_BLOCK_SCOPE_H_
33 #include <unordered_map>
75 static constexpr
const char*
_type_key =
"tir.StmtSRef";
81 this->parent =
nullptr;
92 template <
typename StmtType>
94 if (
stmt !=
nullptr &&
stmt->IsInstance<StmtType>()) {
95 return static_cast<const StmtType*
>(
stmt);
157 bool include_loops =
true) {
159 for (
const auto& kv :
mod->functions) {
160 const BaseFunc& base_func = kv.second;
161 if (
auto opt = base_func.as<
PrimFunc>()) {
162 auto func = opt.value();
163 creator.VisitStmt(func->body);
166 return std::move(creator.stmt2ref_);
170 explicit SRefTreeCreator(
bool include_loops) : include_loops_(include_loops) {}
176 void PushSRef(
const StmtNode* stmt);
179 void PopAndRecordSRef();
181 void VisitStmt_(
const ForNode* loop)
final;
183 void VisitStmt_(
const BlockRealizeNode* realize)
final;
185 void VisitStmt_(
const SeqStmtNode* seq_stmt)
final;
189 std::unordered_map<const StmtNode*, StmtSRef> stmt2ref_;
191 std::vector<StmtSRef> srefs_;
224 refl::ObjectDef<DependencyNode>()
230 static constexpr
const char*
_type_key =
"tir.Dependency";
265 std::unordered_map<StmtSRef, Array<Dependency>, ObjectPtrHash, ObjectPtrEqual>
src2deps;
267 std::unordered_map<StmtSRef, Array<Dependency>, ObjectPtrHash, ObjectPtrEqual>
dst2deps;
269 std::unordered_map<Buffer, Array<StmtSRef>, ObjectPtrHash, ObjectPtrEqual>
buffer_writers;
275 static constexpr
const char*
_type_key =
"tir.BlockScope";
308 TVM_DLL
explicit BlockScope(
const Array<StmtSRef>& child_block_srefs);
Managed reference to BaseFuncNode.
Definition: function.h:234
Managed reference class to IRModuleNode.
Definition: module.h:257
An object with 1-to-1 correspondence with each block reference in the sref tree. This data structure ...
Definition: block_scope.h:258
Array< Dependency > GetDepsByDst(const StmtSRef &dst) const
Get all dependencies whose dst equals dst
static void RegisterReflection()
Definition: block_scope.h:271
Array< Dependency > GetDepsBySrc(const StmtSRef &src) const
Get all dependencies whose src equals src
std::unordered_map< StmtSRef, Array< Dependency >, ObjectPtrHash, ObjectPtrEqual > dst2deps
Lookup table for the dst of dependencies.
Definition: block_scope.h:267
std::unordered_map< StmtSRef, Array< Dependency >, ObjectPtrHash, ObjectPtrEqual > src2deps
Lookup table for the src of dependencies.
Definition: block_scope.h:265
std::unordered_map< Buffer, Array< StmtSRef >, ObjectPtrHash, ObjectPtrEqual > buffer_writers
The mapping from the buffer to the blocks who write it.
Definition: block_scope.h:269
TVM_DECLARE_FINAL_OBJECT_INFO(BlockScopeNode, Object)
static constexpr const char * _type_key
Definition: block_scope.h:275
Managed reference to BlockScopeNode.
Definition: block_scope.h:298
BlockScope()
The constructor creating an empty block scope with on dependency information.
BlockScope(const Array< StmtSRef > &child_block_srefs)
Create the object with the specific leaf blocks, and compute the dependency information between the l...
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockScope, ObjectRef, BlockScopeNode)
A tuple (src, dst, kind) representing certain types of dependency. For example, (A,...
Definition: block_scope.h:213
static void RegisterReflection()
Definition: block_scope.h:222
StmtSRef dst
The destination of the dependency relation.
Definition: block_scope.h:218
StmtSRef src
The source of the dependency relation.
Definition: block_scope.h:216
static constexpr const char * _type_key
Definition: block_scope.h:230
DepKind kind
The dependency kind.
Definition: block_scope.h:220
TVM_DECLARE_FINAL_OBJECT_INFO(DependencyNode, Object)
Managed reference to DependencyNode.
Definition: block_scope.h:238
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Dependency, ObjectRef, DependencyNode)
Dependency(StmtSRef src, StmtSRef dst, DepKind kind)
Constructor.
Managed reference to PrimFuncNode.
Definition: function.h:131
Definition: block_scope.h:149
static std::unordered_map< const StmtNode *, StmtSRef > Create(IRModule mod, bool include_loops=true)
StmtSRef Tree Creator.
Definition: block_scope.h:156
Base node of all statements.
Definition: stmt.h:38
An object that refers to schedulable elements (block/for-loop) in TensorIR, aka "sref".
Definition: block_scope.h:54
void Reset()
Reset the object inplace to the invalid state.
Definition: block_scope.h:79
int64_t seq_index
If the statement the sref points to is an element of a SeqStmt in the AST, then seq_index is set to i...
Definition: block_scope.h:68
TVM_DECLARE_FINAL_OBJECT_INFO(StmtSRefNode, Object)
static void RegisterReflection()
Definition: block_scope.h:70
StmtSRefNode * parent
The parent sref.
Definition: block_scope.h:63
const StmtNode * stmt
The block or for stmt the object refers to.
Definition: block_scope.h:61
static constexpr const char * _type_key
Definition: block_scope.h:75
const StmtType * StmtAs() const
Get the referenced statement with proper type checking. It serves the same purpose as ObjectRef::as,...
Definition: block_scope.h:93
Managed reference to StmtSRefNode.
Definition: block_scope.h:106
static StmtSRef InlineMark()
StmtSRef(const StmtNode *stmt, StmtSRefNode *parent, int64_t seq_index)
The constructor.
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(StmtSRef, ObjectRef, StmtSRefNode)
static StmtSRef RootMark()
StmtSRefNode * get() const
Definition: block_scope.h:118
StmtVisitor.
Definition: stmt_functor.h:135
IRModule that holds the functions and type definitions.
Definition: repr_printer.h:91
DepKind
Type of dependency. Right now we have 4 types of dependencies 1) Read-after-write (kRAW) 2) Write-aft...
Definition: block_scope.h:201
@ kOpaque
IterVar is opaque,.
Definition: var.h:227
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
Functors for tir stmts utility functions to call common functors.