tvm
frame.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_RELAX_FRAME_H_
20 #define TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
21 
22 #include <tvm/ffi/reflection/registry.h>
24 #include <tvm/relax/expr.h>
28 
29 #include <utility>
30 
31 namespace tvm {
32 namespace script {
33 namespace ir_builder {
34 namespace relax {
35 
38  public:
39  static void RegisterReflection() {
40  namespace refl = tvm::ffi::reflection;
41  refl::ObjectDef<RelaxFrameNode>();
42  }
43  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.RelaxFrame", RelaxFrameNode,
45 };
46 
47 class RelaxFrame : public IRBuilderFrame {
48  public:
49  explicit RelaxFrame(ObjectPtr<RelaxFrameNode> data) : IRBuilderFrame(ffi::UnsafeInit{}) {
50  TVM_FFI_ICHECK(data != nullptr);
51  data_ = std::move(data);
52  }
54 
55  protected:
56  RelaxFrame() = default;
57 };
58 
63  public:
65  ffi::Array<tvm::relax::BindingBlock> binding_blocks;
67  ffi::Optional<tvm::relax::Expr> output;
68 
69  static void RegisterReflection() {
70  namespace refl = tvm::ffi::reflection;
71  refl::ObjectDef<SeqExprFrameNode>()
72  .def_ro("binding_blocks", &SeqExprFrameNode::binding_blocks)
73  .def_ro("output", &SeqExprFrameNode::output);
74  }
75  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.SeqExprFrame", SeqExprFrameNode,
77 
78  public:
79  void EnterWithScope() override;
80  void ExitWithScope() override;
81 };
82 
83 class SeqExprFrame : public RelaxFrame {
84  public:
85  explicit SeqExprFrame(ObjectPtr<SeqExprFrameNode> data) : RelaxFrame(data) {
86  TVM_FFI_ICHECK(data != nullptr);
87  }
89 };
90 
93  public:
99  ffi::Optional<ffi::String> name;
101  ffi::Array<tvm::relax::Var> params;
111  ffi::Optional<tvm::relax::StructInfo> ret_struct_info;
113  ffi::Optional<Bool> is_pure;
115  ffi::Optional<Bool> is_private;
117  ffi::Map<ffi::String, Any> attrs;
120 
121  static void RegisterReflection() {
122  namespace refl = tvm::ffi::reflection;
123  refl::ObjectDef<FunctionFrameNode>()
124  .def_ro("name", &FunctionFrameNode::name)
125  .def_ro("params", &FunctionFrameNode::params)
126  .def_ro("ret_struct_info", &FunctionFrameNode::ret_struct_info)
127  .def_ro("is_pure", &FunctionFrameNode::is_pure)
128  .def_ro("attrs", &FunctionFrameNode::attrs)
129  .def_ro("binding_blocks", &FunctionFrameNode::binding_blocks)
130  .def_ro("output", &FunctionFrameNode::output);
131  // `block_builder` is not registered as it's not visited.
132  }
133  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.FunctionFrame", FunctionFrameNode,
135 
136  public:
137  void EnterWithScope() final;
138  void ExitWithScope() final;
139 };
140 
141 class FunctionFrame : public SeqExprFrame {
142  public:
143  explicit FunctionFrame(ObjectPtr<FunctionFrameNode> data) : SeqExprFrame(data) {
144  TVM_FFI_ICHECK(data != nullptr);
145  }
147 };
148 
151  public:
155  ffi::Array<tvm::relax::Var> emitted_vars;
166  ffi::Array<tvm::relax::Var> output_vars;
167 
168  static void RegisterReflection() {
169  namespace refl = tvm::ffi::reflection;
170  refl::ObjectDef<BlockFrameNode>()
171  .def_ro("is_dataflow", &BlockFrameNode::is_dataflow)
172  .def_ro("emitted_vars", &BlockFrameNode::emitted_vars)
173  .def_ro("output_vars", &BlockFrameNode::output_vars);
174  // `block_ended` is not registered as it's not visited.
175  }
176  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.BlockFrame", BlockFrameNode,
178 
179  public:
180  void EnterWithScope() final;
181  void ExitWithScope() final;
182 };
183 
184 class BlockFrame : public RelaxFrame {
185  public:
186  explicit BlockFrame(ObjectPtr<BlockFrameNode> data) : RelaxFrame(data) {
187  TVM_FFI_ICHECK(data != nullptr);
188  }
190 };
191 
197 class IfFrameNode : public RelaxFrameNode {
198  public:
202  ffi::Optional<tvm::relax::Expr> then_expr;
204  ffi::Optional<tvm::relax::Expr> else_expr;
208  ffi::String var_name;
209 
210  static void RegisterReflection() {
211  namespace refl = tvm::ffi::reflection;
212  refl::ObjectDef<IfFrameNode>()
213  .def_ro("condition", &IfFrameNode::condition)
214  .def_ro("then_expr", &IfFrameNode::then_expr)
215  .def_ro("else_expr", &IfFrameNode::else_expr)
216  .def_ro("var", &IfFrameNode::var)
217  .def_ro("var_name", &IfFrameNode::var_name);
218  }
219  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.IfFrame", IfFrameNode, RelaxFrameNode);
220 
221  public:
226  void EnterWithScope() final;
231  void ExitWithScope() final;
232 };
233 
239 class IfFrame : public RelaxFrame {
240  public:
241  explicit IfFrame(ObjectPtr<IfFrameNode> data) : RelaxFrame(data) {
242  TVM_FFI_ICHECK(data != nullptr);
243  }
245 };
246 
253  public:
254  static void RegisterReflection() {
255  namespace refl = tvm::ffi::reflection;
256  refl::ObjectDef<ThenFrameNode>();
257  }
258  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ThenFrame", ThenFrameNode,
260 
261  public:
266  void EnterWithScope() final;
271  void ExitWithScope() final;
272 };
273 
279 class ThenFrame : public SeqExprFrame {
280  public:
281  explicit ThenFrame(ObjectPtr<ThenFrameNode> data) : SeqExprFrame(data) {
282  TVM_FFI_ICHECK(data != nullptr);
283  }
285 };
286 
293  public:
294  static void RegisterReflection() {
295  namespace refl = tvm::ffi::reflection;
296  refl::ObjectDef<ElseFrameNode>();
297  }
298  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ElseFrame", ElseFrameNode,
300 
301  public:
306  void EnterWithScope() final;
311  void ExitWithScope() final;
312 };
313 
319 class ElseFrame : public SeqExprFrame {
320  public:
321  explicit ElseFrame(ObjectPtr<ElseFrameNode> data) : SeqExprFrame(data) {
322  TVM_FFI_ICHECK(data != nullptr);
323  }
325 };
326 
327 } // namespace relax
328 } // namespace ir_builder
329 } // namespace script
330 } // namespace tvm
331 
332 #endif // TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
The utility for constructing Relax binding blocks.
Managed reference to RelaxExprNode.
Definition: expr.h:439
Definition: block_builder.h:264
Definition: expr.h:377
Managed reference to an IRBuilderFrameNode.
Definition: base.h:104
The ir_builder frame for relax binding blocks.
Definition: frame.h:150
void EnterWithScope() final
The method called when entering RAII scope.
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.BlockFrame", BlockFrameNode, RelaxFrameNode)
ffi::Array< tvm::relax::Var > emitted_vars
The variables emitted in this block.
Definition: frame.h:155
ffi::Array< tvm::relax::Var > output_vars
The output vars of the dataflow block.
Definition: frame.h:166
static void RegisterReflection()
Definition: frame.h:168
bool is_dataflow
The flag that indicates whether the block is a dataflow block.
Definition: frame.h:153
bool block_ended
A boolean indicating if the dataflow block is ended of construction. If it is true,...
Definition: frame.h:161
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockFrame, RelaxFrame, BlockFrameNode)
BlockFrame(ObjectPtr< BlockFrameNode > data)
Definition: frame.h:186
A frame that represents else.
Definition: frame.h:292
static void RegisterReflection()
Definition: frame.h:294
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ElseFrame", ElseFrameNode, SeqExprFrameNode)
void EnterWithScope() final
The method called when entering RAII scope.
Managed reference to ElseFrameNode.
Definition: frame.h:319
ElseFrame(ObjectPtr< ElseFrameNode > data)
Definition: frame.h:321
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, SeqExprFrame, ElseFrameNode)
The ir_builder frame for the relax function.
Definition: frame.h:92
tvm::relax::BlockBuilder block_builder
The block builder to create Relax function.
Definition: frame.h:119
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.FunctionFrame", FunctionFrameNode, SeqExprFrameNode)
void EnterWithScope() final
The method called when entering RAII scope.
ffi::Optional< Bool > is_pure
Whether the function is annotated as pure.
Definition: frame.h:113
static void RegisterReflection()
Definition: frame.h:121
void ExitWithScope() final
The method called when exiting RAII scope.
ffi::Array< tvm::relax::Var > params
The function params.
Definition: frame.h:101
ffi::Optional< ffi::String > name
The function name.
Definition: frame.h:99
ffi::Map< ffi::String, Any > attrs
The function attributes.
Definition: frame.h:117
ffi::Optional< Bool > is_private
Whether the function is annotated as private.
Definition: frame.h:115
ffi::Optional< tvm::relax::StructInfo > ret_struct_info
The function return struct info.
Definition: frame.h:111
FunctionFrame(ObjectPtr< FunctionFrameNode > data)
Definition: frame.h:143
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionFrame, SeqExprFrame, FunctionFrameNode)
A frame that represents if statement.
Definition: frame.h:197
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.IfFrame", IfFrameNode, RelaxFrameNode)
ffi::Optional< tvm::relax::Expr > then_expr
The Bindings in the true branch.
Definition: frame.h:202
void EnterWithScope() final
The method called when entering RAII scope.
static void RegisterReflection()
Definition: frame.h:210
tvm::relax::Var var
The Binding var.
Definition: frame.h:206
tvm::relax::Expr condition
The condition of the if statement.
Definition: frame.h:200
ffi::String var_name
The binding var name.
Definition: frame.h:208
ffi::Optional< tvm::relax::Expr > else_expr
The Bindings in the false branch.
Definition: frame.h:204
Managed reference to IfFrameNode.
Definition: frame.h:239
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfFrame, RelaxFrame, IfFrameNode)
IfFrame(ObjectPtr< IfFrameNode > data)
Definition: frame.h:241
The base ir_builder frame for the relax dialect.
Definition: frame.h:37
static void RegisterReflection()
Definition: frame.h:39
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.RelaxFrame", RelaxFrameNode, IRBuilderFrameNode)
RelaxFrame(ObjectPtr< RelaxFrameNode > data)
Definition: frame.h:49
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RelaxFrame, IRBuilderFrame, RelaxFrameNode)
The base ir_builder frame for frames with SeqExpr i.e. Functions, If branches.
Definition: frame.h:62
void EnterWithScope() override
The method called when entering RAII scope.
ffi::Optional< tvm::relax::Expr > output
The frame output expr. std::nullopt when undefined.
Definition: frame.h:67
static void RegisterReflection()
Definition: frame.h:69
void ExitWithScope() override
The method called when exiting RAII scope.
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.SeqExprFrame", SeqExprFrameNode, RelaxFrameNode)
ffi::Array< tvm::relax::BindingBlock > binding_blocks
The binding blocks inside the frame.
Definition: frame.h:65
SeqExprFrame(ObjectPtr< SeqExprFrameNode > data)
Definition: frame.h:85
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExprFrame, RelaxFrame, SeqExprFrameNode)
A frame that represents then.
Definition: frame.h:252
void EnterWithScope() final
The method called when entering RAII scope.
static void RegisterReflection()
Definition: frame.h:254
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ThenFrame", ThenFrameNode, SeqExprFrameNode)
Managed reference to ThenFrameNode.
Definition: frame.h:279
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThenFrame, SeqExprFrame, ThenFrameNode)
ThenFrame(ObjectPtr< ThenFrameNode > data)
Definition: frame.h:281
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37