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 bool _type_mutable = true;
77  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.IRBuilderFrame", IRBuilderFrameNode,
78  runtime::Object);
79 
80  public:
82  virtual ~IRBuilderFrameNode() = default;
87  virtual void EnterWithScope();
92  virtual void ExitWithScope();
97  void AddCallback(ffi::TypedFunction<void()> callback);
98 };
99 
104 class IRBuilderFrame : public runtime::ObjectRef {
105  public:
107 
108  protected:
110  IRBuilderFrame() = default;
111  explicit IRBuilderFrame(ObjectPtr<IRBuilderFrameNode> data) : ObjectRef(data) {}
112 
113  public:
118  inline void EnterWithScope() {
119  ICHECK(data_ != nullptr);
120  static_cast<IRBuilderFrameNode*>(data_.get())->EnterWithScope();
121  }
126  inline void ExitWithScope() {
127  ICHECK(data_ != nullptr);
128  static_cast<IRBuilderFrameNode*>(data_.get())->ExitWithScope();
129  data_.reset();
130  }
131 };
132 
134 
159 class IRBuilderNode : public runtime::Object {
160  public:
162  ffi::Array<IRBuilderFrame> frames;
164  ffi::Optional<ObjectRef> result;
165 
166  static void RegisterReflection() {
167  namespace refl = tvm::ffi::reflection;
168  refl::ObjectDef<IRBuilderNode>()
169  .def_ro("frames", &IRBuilderNode::frames)
170  .def_ro("result", &IRBuilderNode::result);
171  }
172 
173  static constexpr const bool _type_mutable = true;
174  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.IRBuilder", IRBuilderNode, runtime::Object);
175 
176  public:
182  template <typename TFrame>
183  inline ffi::Optional<TFrame> FindFrame() const;
190  template <typename TFrame>
191  inline ffi::Optional<TFrame> GetLastFrame() const;
197  template <typename TObjectRef>
198  inline TObjectRef Get() const;
199 };
200 
205 class IRBuilder : public runtime::ObjectRef {
206  public:
210 
211  public:
244  static IRBuilder Current();
246  static bool IsInScope();
253  template <class TObjectRef>
254  inline static TObjectRef Name(ffi::String name, TObjectRef obj);
255 };
256 
258 
259 namespace details {
260 
261 class Namer {
262  public:
263  using FType = NodeFunctor<void(const ObjectRef&, ffi::String)>;
264  static FType& vtable();
265  static void Name(ObjectRef node, ffi::String name);
266 };
267 
268 } // namespace details
269 
270 template <class TObjectRef>
271 inline TObjectRef IRBuilder::Name(ffi::String name, TObjectRef obj) {
272  details::Namer::Name(obj, name);
273  return Downcast<TObjectRef>(obj);
274 }
275 
276 template <typename TFrame>
277 inline ffi::Optional<TFrame> IRBuilderNode::FindFrame() const {
278  using TFrameNode = typename TFrame::ContainerType;
279  for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
280  if (const TFrameNode* p = (*it).template as<TFrameNode>()) {
281  return ffi::GetRef<TFrame>(p);
282  }
283  }
284  return std::nullopt;
285 }
286 
287 template <typename TFrame>
288 inline ffi::Optional<TFrame> IRBuilderNode::GetLastFrame() const {
289  using TFrameNode = typename TFrame::ContainerType;
290  if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
291  return Downcast<TFrame>(frames.back());
292  }
293  return std::nullopt;
294 }
295 
296 template <typename TObjectRef>
297 inline TObjectRef IRBuilderNode::Get() const {
298  using TObject = typename TObjectRef::ContainerType;
299  CHECK(result.defined()) << "IndexError: No result exists in IRBuilder yet";
300  const auto* n = result.as<TObject>();
301  CHECK(n != nullptr) << "TypeError: IRBuilder result is not of type: " << TObject::_type_key;
302  return ffi::GetRef<TObjectRef>(n);
303 }
304 
305 } // namespace ir_builder
306 } // namespace script
307 } // namespace tvm
308 
309 #endif // TVM_SCRIPT_IR_BUILDER_BASE_H_
A dynamically dispatched functor on the type of the first argument.
Definition: functor.h:65
static constexpr const bool _type_mutable
Definition: base.h:76
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.IRBuilderFrame", IRBuilderFrameNode, runtime::Object)
std::vector< ffi::TypedFunction< void()> > callbacks
A list of callbacks used when exiting the frame.
Definition: base.h:68
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:104
void EnterWithScope()
Redirected to IRBuilderFrameNode::EnterWithScope.
Definition: base.h:118
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRBuilderFrame, ObjectRef, IRBuilderFrameNode)
IRBuilderFrame(ObjectPtr< IRBuilderFrameNode > data)
Definition: base.h:111
IRBuilderFrame()=default
Disallow direct construction of this object.
void ExitWithScope()
Redirected to IRBuilderFrameNode::ExitWithScope.
Definition: base.h:126
A dialect-agnostic IRBuilder that constructs any IR of TVM. An idiomatic use of this class is to put ...
Definition: base.h:159
ffi::Optional< TFrame > FindFrame() const
Find a frame of the given type in the stack this->frames from top to bottom.
Definition: base.h:277
ffi::Optional< TFrame > GetLastFrame() const
Get the frame on top of the stack this->frames if its type is TFrame.
Definition: base.h:288
TObjectRef Get() const
Get the IR being constructed.
Definition: base.h:297
ffi::Optional< ObjectRef > result
The outcome of IR construction.
Definition: base.h:164
static constexpr const bool _type_mutable
Definition: base.h:173
ffi::Array< IRBuilderFrame > frames
A stack of context frames in the IRBuilder.
Definition: base.h:162
static void RegisterReflection()
Definition: base.h:166
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.IRBuilder", IRBuilderNode, runtime::Object)
Managed reference to an IRBuilderNode.
Definition: base.h:205
static TObjectRef Name(ffi::String name, TObjectRef obj)
Give a string name to the obj
Definition: base.h:271
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.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRBuilder, ObjectRef, IRBuilderNode)
IRBuilder()
Creates an IRBuilder.
static void Name(ObjectRef node, ffi::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.