tvm
base.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  */
19 #ifndef TVM_SCRIPT_IR_BUILDER_BASE_H_
20 #define TVM_SCRIPT_IR_BUILDER_BASE_H_
21 
22 #include <tvm/ffi/reflection/registry.h>
23 #include <tvm/ir/expr.h>
24 #include <tvm/ir/function.h>
25 #include <tvm/node/node.h>
26 
27 #include <vector>
28 
29 namespace tvm {
30 namespace script {
31 namespace ir_builder {
32 
34 
65 class IRBuilderFrameNode : public runtime::Object {
66  public:
68  std::vector<ffi::TypedFunction<void()>> callbacks;
69 
70  static void RegisterReflection() {
71  namespace refl = tvm::ffi::reflection;
72  refl::ObjectDef<IRBuilderFrameNode>();
73  // `callbacks` is not registered as it's not visited.
74  }
75 
76  static constexpr const char* _type_key = "script.ir_builder.IRBuilderFrame";
78 
79  public:
81  virtual ~IRBuilderFrameNode() = default;
86  virtual void EnterWithScope();
91  virtual void ExitWithScope();
96  void AddCallback(ffi::TypedFunction<void()> callback);
97 };
98 
103 class IRBuilderFrame : public runtime::ObjectRef {
104  public:
106 
107  protected:
109  IRBuilderFrame() = default;
110 
111  public:
116  inline void EnterWithScope() {
117  ICHECK(data_ != nullptr);
118  static_cast<IRBuilderFrameNode*>(data_.get())->EnterWithScope();
119  }
124  inline void ExitWithScope() {
125  ICHECK(data_ != nullptr);
126  static_cast<IRBuilderFrameNode*>(data_.get())->ExitWithScope();
127  data_.reset();
128  }
129 };
130 
132 
157 class IRBuilderNode : public runtime::Object {
158  public:
160  Array<IRBuilderFrame> frames;
162  Optional<ObjectRef> result;
163 
164  static void RegisterReflection() {
165  namespace refl = tvm::ffi::reflection;
166  refl::ObjectDef<IRBuilderNode>()
167  .def_ro("frames", &IRBuilderNode::frames)
168  .def_ro("result", &IRBuilderNode::result);
169  }
170 
171  static constexpr const char* _type_key = "script.ir_builder.IRBuilder";
173 
174  public:
180  template <typename TFrame>
181  inline Optional<TFrame> FindFrame() const;
188  template <typename TFrame>
189  inline Optional<TFrame> GetLastFrame() const;
195  template <typename TObjectRef>
196  inline TObjectRef Get() const;
197 };
198 
203 class IRBuilder : public runtime::ObjectRef {
204  public:
208 
209  public:
242  static IRBuilder Current();
244  static bool IsInScope();
251  template <class TObjectRef>
252  inline static TObjectRef Name(String name, TObjectRef obj);
253 };
254 
256 
257 namespace details {
258 
259 class Namer {
260  public:
261  using FType = NodeFunctor<void(const ObjectRef&, String)>;
262  static FType& vtable();
263  static void Name(ObjectRef node, String name);
264 };
265 
266 } // namespace details
267 
268 template <class TObjectRef>
269 inline TObjectRef IRBuilder::Name(String name, TObjectRef obj) {
270  details::Namer::Name(obj, name);
271  return Downcast<TObjectRef>(obj);
272 }
273 
274 template <typename TFrame>
275 inline Optional<TFrame> IRBuilderNode::FindFrame() const {
276  using TFrameNode = typename TFrame::ContainerType;
277  for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
278  if (const TFrameNode* p = (*it).template as<TFrameNode>()) {
279  return GetRef<TFrame>(p);
280  }
281  }
282  return std::nullopt;
283 }
284 
285 template <typename TFrame>
286 inline Optional<TFrame> IRBuilderNode::GetLastFrame() const {
287  using TFrameNode = typename TFrame::ContainerType;
288  if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
289  return Downcast<TFrame>(frames.back());
290  }
291  return std::nullopt;
292 }
293 
294 template <typename TObjectRef>
295 inline TObjectRef IRBuilderNode::Get() const {
296  using TObject = typename TObjectRef::ContainerType;
297  CHECK(result.defined()) << "IndexError: No result exists in IRBuilder yet";
298  const auto* n = result.as<TObject>();
299  CHECK(n != nullptr) << "TypeError: IRBuilder result is not of type: " << TObject::_type_key;
300  return GetRef<TObjectRef>(n);
301 }
302 
303 } // namespace ir_builder
304 } // namespace script
305 } // namespace tvm
306 
307 #endif // TVM_SCRIPT_IR_BUILDER_BASE_H_
A dynamically dispatched functor on the type of the first argument.
Definition: functor.h:65
TVM_DECLARE_BASE_OBJECT_INFO(IRBuilderFrameNode, runtime::Object)
std::vector< ffi::TypedFunction< void()> > callbacks
A list of callbacks used when exiting the frame.
Definition: base.h:68
static constexpr const char * _type_key
Definition: base.h:76
virtual ~IRBuilderFrameNode()=default
Default destructor.
virtual void ExitWithScope()
The method called when exiting RAII scope.
virtual void EnterWithScope()
The method called when entering RAII scope.
void AddCallback(ffi::TypedFunction< void()> callback)
Add a callback method invoked when exiting the RAII scope.
static void RegisterReflection()
Definition: base.h:70
Managed reference to an IRBuilderFrameNode.
Definition: base.h:103
void EnterWithScope()
Redirected to IRBuilderFrameNode::EnterWithScope.
Definition: base.h:116
IRBuilderFrame()=default
Disallow direct construction of this object.
void ExitWithScope()
Redirected to IRBuilderFrameNode::ExitWithScope.
Definition: base.h:124
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRBuilderFrame, ObjectRef, IRBuilderFrameNode)
A dialect-agnostic IRBuilder that constructs any IR of TVM. An idiomatic use of this class is to put ...
Definition: base.h:157
TObjectRef Get() const
Get the IR being constructed.
Definition: base.h:295
Array< IRBuilderFrame > frames
A stack of context frames in the IRBuilder.
Definition: base.h:160
Optional< TFrame > GetLastFrame() const
Get the frame on top of the stack this->frames if its type is TFrame.
Definition: base.h:286
static constexpr const char * _type_key
Definition: base.h:171
TVM_DECLARE_FINAL_OBJECT_INFO(IRBuilderNode, runtime::Object)
Optional< TFrame > FindFrame() const
Find a frame of the given type in the stack this->frames from top to bottom.
Definition: base.h:275
Optional< ObjectRef > result
The outcome of IR construction.
Definition: base.h:162
static void RegisterReflection()
Definition: base.h:164
Managed reference to an IRBuilderNode.
Definition: base.h:203
void ExitWithScope()
Exit the RAII scope.
static IRBuilder Current()
Get the current IRBuilder in the current thread-local scope.
void EnterWithScope()
Puts the current IRBuilder into a thread-local scope, which can be retrieved using IRBuilder::Current...
static bool IsInScope()
See if the current thread-local scope has an IRBuilder.
static TObjectRef Name(String name, TObjectRef obj)
Give a string name to the obj
Definition: base.h:269
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRBuilder, ObjectRef, IRBuilderNode)
IRBuilder()
Creates an IRBuilder.
static void Name(ObjectRef node, String name)
Base expr nodes in TVM.
Function nodes.
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Definitions and helper macros for IR/AST nodes.