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 namespace tvm {
30 namespace script {
31 namespace ir_builder {
32 namespace relax {
33 
36  public:
37  static void RegisterReflection() {
38  namespace refl = tvm::ffi::reflection;
39  refl::ObjectDef<RelaxFrameNode>();
40  }
41 
42  static constexpr const char* _type_key = "script.ir_builder.relax.RelaxFrame";
44 };
45 
46 class RelaxFrame : public IRBuilderFrame {
47  public:
49 
50  protected:
51  RelaxFrame() = default;
52 };
53 
58  public:
60  Array<tvm::relax::BindingBlock> binding_blocks;
62  Optional<tvm::relax::Expr> output;
63 
64  static void RegisterReflection() {
65  namespace refl = tvm::ffi::reflection;
66  refl::ObjectDef<SeqExprFrameNode>()
67  .def_ro("binding_blocks", &SeqExprFrameNode::binding_blocks)
68  .def_ro("output", &SeqExprFrameNode::output);
69  }
70 
71  static constexpr const char* _type_key = "script.ir_builder.relax.SeqExprFrame";
73 
74  public:
75  void EnterWithScope() override;
76  void ExitWithScope() override;
77 };
78 
79 class SeqExprFrame : public RelaxFrame {
80  public:
82 };
83 
86  public:
92  Optional<String> name;
94  Array<tvm::relax::Var> params;
104  Optional<tvm::relax::StructInfo> ret_struct_info;
106  Optional<Bool> is_pure;
108  Optional<Bool> is_private;
110  Map<String, Any> attrs;
113 
114  static void RegisterReflection() {
115  namespace refl = tvm::ffi::reflection;
116  refl::ObjectDef<FunctionFrameNode>()
117  .def_ro("name", &FunctionFrameNode::name)
118  .def_ro("params", &FunctionFrameNode::params)
119  .def_ro("ret_struct_info", &FunctionFrameNode::ret_struct_info)
120  .def_ro("is_pure", &FunctionFrameNode::is_pure)
121  .def_ro("attrs", &FunctionFrameNode::attrs)
122  .def_ro("binding_blocks", &FunctionFrameNode::binding_blocks)
123  .def_ro("output", &FunctionFrameNode::output);
124  // `block_builder` is not registered as it's not visited.
125  }
126 
127  static constexpr const char* _type_key = "script.ir_builder.relax.FunctionFrame";
129 
130  public:
131  void EnterWithScope() final;
132  void ExitWithScope() final;
133 };
134 
135 class FunctionFrame : public SeqExprFrame {
136  public:
138 };
139 
142  public:
146  Array<tvm::relax::Var> emitted_vars;
157  Array<tvm::relax::Var> output_vars;
158 
159  static void RegisterReflection() {
160  namespace refl = tvm::ffi::reflection;
161  refl::ObjectDef<BlockFrameNode>()
162  .def_ro("is_dataflow", &BlockFrameNode::is_dataflow)
163  .def_ro("emitted_vars", &BlockFrameNode::emitted_vars)
164  .def_ro("output_vars", &BlockFrameNode::output_vars);
165  // `block_ended` is not registered as it's not visited.
166  }
167 
168  static constexpr const char* _type_key = "script.ir_builder.relax.BlockFrame";
170 
171  public:
172  void EnterWithScope() final;
173  void ExitWithScope() final;
174 };
175 
176 class BlockFrame : public RelaxFrame {
177  public:
179 };
180 
186 class IfFrameNode : public RelaxFrameNode {
187  public:
191  Optional<tvm::relax::Expr> then_expr;
193  Optional<tvm::relax::Expr> else_expr;
197  String var_name;
198 
199  static void RegisterReflection() {
200  namespace refl = tvm::ffi::reflection;
201  refl::ObjectDef<IfFrameNode>()
202  .def_ro("condition", &IfFrameNode::condition)
203  .def_ro("then_expr", &IfFrameNode::then_expr)
204  .def_ro("else_expr", &IfFrameNode::else_expr)
205  .def_ro("var", &IfFrameNode::var)
206  .def_ro("var_name", &IfFrameNode::var_name);
207  }
208 
209  static constexpr const char* _type_key = "script.ir_builder.relax.IfFrame";
211 
212  public:
217  void EnterWithScope() final;
222  void ExitWithScope() final;
223 };
224 
230 class IfFrame : public RelaxFrame {
231  public:
233 };
234 
241  public:
242  static void RegisterReflection() {
243  namespace refl = tvm::ffi::reflection;
244  refl::ObjectDef<ThenFrameNode>();
245  }
246 
247  static constexpr const char* _type_key = "script.ir_builder.relax.ThenFrame";
249 
250  public:
255  void EnterWithScope() final;
260  void ExitWithScope() final;
261 };
262 
268 class ThenFrame : public SeqExprFrame {
269  public:
271 };
272 
279  public:
280  static void RegisterReflection() {
281  namespace refl = tvm::ffi::reflection;
282  refl::ObjectDef<ElseFrameNode>();
283  }
284 
285  static constexpr const char* _type_key = "script.ir_builder.relax.ElseFrame";
287 
288  public:
293  void EnterWithScope() final;
298  void ExitWithScope() final;
299 };
300 
306 class ElseFrame : public SeqExprFrame {
307  public:
309 };
310 
311 } // namespace relax
312 } // namespace ir_builder
313 } // namespace script
314 } // namespace tvm
315 
316 #endif // TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
The utility for constructing Relax binding blocks.
Managed reference to RelaxExprNode.
Definition: expr.h:446
Definition: block_builder.h:264
Definition: expr.h:387
Managed reference to an IRBuilderFrameNode.
Definition: base.h:103
The ir_builder frame for relax binding blocks.
Definition: frame.h:141
TVM_DECLARE_FINAL_OBJECT_INFO(BlockFrameNode, RelaxFrameNode)
void EnterWithScope() final
The method called when entering RAII scope.
static void RegisterReflection()
Definition: frame.h:159
bool is_dataflow
The flag that indicates whether the block is a dataflow block.
Definition: frame.h:144
Array< tvm::relax::Var > output_vars
The output vars of the dataflow block.
Definition: frame.h:157
Array< tvm::relax::Var > emitted_vars
The variables emitted in this block.
Definition: frame.h:146
bool block_ended
A boolean indicating if the dataflow block is ended of construction. If it is true,...
Definition: frame.h:152
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockFrame, RelaxFrame, BlockFrameNode)
A frame that represents else.
Definition: frame.h:278
static void RegisterReflection()
Definition: frame.h:280
TVM_DECLARE_FINAL_OBJECT_INFO(ElseFrameNode, SeqExprFrameNode)
void EnterWithScope() final
The method called when entering RAII scope.
Managed reference to ElseFrameNode.
Definition: frame.h:306
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ElseFrame, SeqExprFrame, ElseFrameNode)
The ir_builder frame for the relax function.
Definition: frame.h:85
tvm::relax::BlockBuilder block_builder
The block builder to create Relax function.
Definition: frame.h:112
Optional< String > name
The function name.
Definition: frame.h:92
void EnterWithScope() final
The method called when entering RAII scope.
static void RegisterReflection()
Definition: frame.h:114
void ExitWithScope() final
The method called when exiting RAII scope.
Optional< Bool > is_private
Whether the function is annotated as private.
Definition: frame.h:108
Map< String, Any > attrs
The function attributes.
Definition: frame.h:110
static constexpr const char * _type_key
Definition: frame.h:127
Optional< tvm::relax::StructInfo > ret_struct_info
The function return struct info.
Definition: frame.h:104
Optional< Bool > is_pure
Whether the function is annotated as pure.
Definition: frame.h:106
TVM_DECLARE_FINAL_OBJECT_INFO(FunctionFrameNode, SeqExprFrameNode)
Array< tvm::relax::Var > params
The function params.
Definition: frame.h:94
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionFrame, SeqExprFrame, FunctionFrameNode)
A frame that represents if statement.
Definition: frame.h:186
void EnterWithScope() final
The method called when entering RAII scope.
TVM_DECLARE_FINAL_OBJECT_INFO(IfFrameNode, RelaxFrameNode)
static void RegisterReflection()
Definition: frame.h:199
tvm::relax::Var var
The Binding var.
Definition: frame.h:195
tvm::relax::Expr condition
The condition of the if statement.
Definition: frame.h:189
String var_name
The binding var name.
Definition: frame.h:197
Optional< tvm::relax::Expr > then_expr
The Bindings in the true branch.
Definition: frame.h:191
Optional< tvm::relax::Expr > else_expr
The Bindings in the false branch.
Definition: frame.h:193
Managed reference to IfFrameNode.
Definition: frame.h:230
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IfFrame, RelaxFrame, IfFrameNode)
The base ir_builder frame for the relax dialect.
Definition: frame.h:35
static constexpr const char * _type_key
Definition: frame.h:42
static void RegisterReflection()
Definition: frame.h:37
TVM_DECLARE_BASE_OBJECT_INFO(RelaxFrameNode, IRBuilderFrameNode)
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RelaxFrame, IRBuilderFrame, RelaxFrameNode)
The base ir_builder frame for frames with SeqExpr i.e. Functions, If branches.
Definition: frame.h:57
TVM_DECLARE_BASE_OBJECT_INFO(SeqExprFrameNode, RelaxFrameNode)
void EnterWithScope() override
The method called when entering RAII scope.
static void RegisterReflection()
Definition: frame.h:64
void ExitWithScope() override
The method called when exiting RAII scope.
Array< tvm::relax::BindingBlock > binding_blocks
The binding blocks inside the frame.
Definition: frame.h:60
Optional< tvm::relax::Expr > output
The frame output expr. std::nullopt when undefined.
Definition: frame.h:62
static constexpr const char * _type_key
Definition: frame.h:71
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(SeqExprFrame, RelaxFrame, SeqExprFrameNode)
A frame that represents then.
Definition: frame.h:240
void EnterWithScope() final
The method called when entering RAII scope.
static void RegisterReflection()
Definition: frame.h:242
TVM_DECLARE_FINAL_OBJECT_INFO(ThenFrameNode, SeqExprFrameNode)
Managed reference to ThenFrameNode.
Definition: frame.h:268
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ThenFrame, SeqExprFrame, ThenFrameNode)
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37