25 #ifndef TVM_TIR_BLOCK_SCOPE_H_
26 #define TVM_TIR_BLOCK_SCOPE_H_
33 #include <unordered_map>
81 this->parent =
nullptr;
92 template <
typename StmtType>
94 if (
stmt !=
nullptr &&
stmt->IsInstance<StmtType>()) {
95 return static_cast<const StmtType*
>(
stmt);
154 bool include_loops =
true) {
156 for (
const auto& kv :
mod->functions) {
157 const BaseFunc& base_func = kv.second;
158 if (
auto opt = base_func.as<
PrimFunc>()) {
159 auto func = opt.value();
160 creator.VisitStmt(func->body);
163 return std::move(creator.stmt2ref_);
167 explicit SRefTreeCreator(
bool include_loops) : include_loops_(include_loops) {}
173 void PushSRef(
const StmtNode* stmt);
176 void PopAndRecordSRef();
178 void VisitStmt_(
const ForNode* loop)
final;
180 void VisitStmt_(
const BlockRealizeNode* realize)
final;
182 void VisitStmt_(
const SeqStmtNode* seq_stmt)
final;
186 std::unordered_map<const StmtNode*, StmtSRef> stmt2ref_;
188 std::vector<StmtSRef> srefs_;
221 refl::ObjectDef<DependencyNode>()
260 std::unordered_map<StmtSRef, ffi::Array<Dependency>, ObjectPtrHash, ObjectPtrEqual>
src2deps;
262 std::unordered_map<StmtSRef, ffi::Array<Dependency>, ObjectPtrHash, ObjectPtrEqual>
dst2deps;
264 std::unordered_map<Buffer, ffi::Array<StmtSRef>, ObjectPtrHash, ObjectPtrEqual>
buffer_writers;
297 explicit BlockScope(ObjectPtr<BlockScopeNode> data) : ObjectRef(data) {
298 TVM_FFI_ICHECK(data !=
nullptr);
308 TVM_DLL
explicit BlockScope(
const ffi::Array<StmtSRef>& child_block_srefs);
Managed reference to BaseFuncNode.
Definition: function.h:233
Managed reference class to IRModuleNode.
Definition: module.h:256
An object with 1-to-1 correspondence with each block reference in the sref tree. This data structure ...
Definition: block_scope.h:253
ffi::Array< Dependency > GetDepsByDst(const StmtSRef &dst) const
Get all dependencies whose dst equals dst
std::unordered_map< StmtSRef, ffi::Array< Dependency >, ObjectPtrHash, ObjectPtrEqual > src2deps
Lookup table for the src of dependencies.
Definition: block_scope.h:260
std::unordered_map< Buffer, ffi::Array< StmtSRef >, ObjectPtrHash, ObjectPtrEqual > buffer_writers
The mapping from the buffer to the blocks who write it.
Definition: block_scope.h:264
std::unordered_map< StmtSRef, ffi::Array< Dependency >, ObjectPtrHash, ObjectPtrEqual > dst2deps
Lookup table for the dst of dependencies.
Definition: block_scope.h:262
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.BlockScope", BlockScopeNode, Object)
static void RegisterReflection()
Definition: block_scope.h:266
ffi::Array< Dependency > GetDepsBySrc(const StmtSRef &src) const
Get all dependencies whose src equals src
Managed reference to BlockScopeNode.
Definition: block_scope.h:291
BlockScope()
The constructor creating an empty block scope with on dependency information.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockScope, ObjectRef, BlockScopeNode)
BlockScope(ObjectPtr< BlockScopeNode > data)
Constructor from ObjectPtr<BlockScopeNode>.
Definition: block_scope.h:297
BlockScope(const ffi::Array< StmtSRef > &child_block_srefs)
Create the object with the specific leaf blocks, and compute the dependency information between the l...
A tuple (src, dst, kind) representing certain types of dependency. For example, (A,...
Definition: block_scope.h:210
static void RegisterReflection()
Definition: block_scope.h:219
StmtSRef dst
The destination of the dependency relation.
Definition: block_scope.h:215
StmtSRef src
The source of the dependency relation.
Definition: block_scope.h:213
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Dependency", DependencyNode, Object)
DepKind kind
The dependency kind.
Definition: block_scope.h:217
Managed reference to DependencyNode.
Definition: block_scope.h:233
Dependency(StmtSRef src, StmtSRef dst, DepKind kind)
Constructor.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Dependency, ObjectRef, DependencyNode)
Managed reference to PrimFuncNode.
Definition: function.h:129
Definition: block_scope.h:146
static std::unordered_map< const StmtNode *, StmtSRef > Create(IRModule mod, bool include_loops=true)
StmtSRef Tree Creator.
Definition: block_scope.h:153
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
static constexpr const bool _type_mutable
Definition: block_scope.h:75
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.StmtSRef", 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
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.
static StmtSRef RootMark()
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StmtSRef, ObjectRef, StmtSRefNode)
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:198
@ kOpaque
IterVar is opaque,.
Definition: var.h:227
tvm::PrimExpr mod(const tvm::PrimExpr &a, const tvm::PrimExpr &b)
Definition: broadcast.h:308
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Functors for tir stmts utility functions to call common functors.