Skip to content

Commit fbc2ec2

Browse files
junrushaoYuchenJinHzfengsyMasterJH5574
committed
Recover dropped commits
[TVMScript] B4: If branch support (tlc-pack#263) B8: Local Function Support (tlc-pack#258) [TVMScript] B3: Type annotation checks (tlc-pack#256) [TVMScript][Parser] B1: Dataflow block (tlc-pack#252) [TVMScript] B2: match shape support (tlc-pack#251) [TVMScript] B6/B7: Symbolic shape and var shadowing (tlc-pack#245) [TVMScript] B5: Support relax op (tlc-pack#244) [TVMScript] B0: Call_tir support (tlc-pack#243) enhance parser error reporting (tlc-pack#242) [TVMScript] A1: Relax Parser infra (tlc-pack#240) update ci image versions. (tlc-pack#241) [TVMScript] B2-4: TIR IRBuilder (tlc-pack#239) [TVMScript] A0: Relax IRBuilder infra (tlc-pack#235) [TVMScript] B5-6: TIR IRBuilder (tlc-pack#231) [TVMScript] B1: IRBuilder (tlc-pack#228) [TVMScript] New Parser: Part C (tlc-pack#218) [TVMScript] New Parser: Part A (tlc-pack#221) [TVMScript] New Parser: Part B (tlc-pack#217) Not recovered: [Pass] Separate ApplyHistoryBest from tuning passes (tlc-pack#226) [Bugfix] Couple of bug fixes to run TVM-gen code together with BYOC (tlc-pack#249) co-authored-by: Yuchen Jin <[email protected]> co-authored-by: Siyuan Feng <[email protected]> co-authored-by: Ruihang Lai <[email protected]>
1 parent eb8aa20 commit fbc2ec2

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

62 files changed

+4635
-353
lines changed

‎include/tvm/relax/transform.h

-10
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,6 @@
2525
#define TVM_RELAX_TRANSFORM_H_
2626

2727
#include <tvm/ir/transform.h>
28-
#include <tvm/meta_schedule/apply_history_best.h>
29-
#include <tvm/meta_schedule/database.h>
3028
#include <tvm/relax/expr.h>
3129

3230
namespace tvm {
@@ -119,14 +117,6 @@ TVM_DLL Pass CanonicalizeBindings();
119117
*/
120118
TVM_DLL Pass Normalize();
121119

122-
/*!
123-
* \brief Apply the best schedule from tuning database.
124-
*
125-
* \return The Pass.
126-
*/
127-
TVM_DLL Pass MetaScheduleApplyHistoryBest(const tvm::meta_schedule::Database& database,
128-
Target target);
129-
130120
/*!
131121
* \brief Bind params of function of the module to constant tensors.
132122
*

‎include/tvm/script/ir_builder/ir/frame.h

+8-3
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,17 @@ namespace ir {
3838
*/
3939
class IRModuleFrameNode : public IRBuilderFrameNode {
4040
public:
41-
Array<GlobalVar> global_vars;
42-
Array<BaseFunc> functions;
41+
/*! \brief A map from string names to global variables that ensures global uniqueness. */
42+
Map<String, GlobalVar> global_var_map;
43+
/*!
44+
* \brief A map from GlobalVar to all global functions.
45+
* \note Only defined functions are in the map, while declared functions are not included.
46+
*/
47+
Map<GlobalVar, BaseFunc> functions;
4348

4449
void VisitAttrs(tvm::AttrVisitor* v) {
4550
IRBuilderFrameNode::VisitAttrs(v);
46-
v->Visit("global_vars", &global_vars);
51+
v->Visit("global_vars", &global_var_map);
4752
v->Visit("functions", &functions);
4853
}
4954

‎include/tvm/script/ir_builder/ir/ir.h

+15
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,21 @@ namespace ir {
3737
*/
3838
TVM_DLL IRModuleFrame IRModule();
3939

40+
/*!
41+
* \brief Declare a Function without given the specific function implementation.
42+
* \note It is usually used in cross-function call. And we can specify the function by `DefFunction`
43+
* \param func_name The function unique name.
44+
* \return The corresponding GlobalVar.
45+
*/
46+
TVM_DLL GlobalVar DeclFunction(const String& func_name);
47+
48+
/*!
49+
* \brief Define the function which is declared before.
50+
* \param func_name The function unique name.
51+
* \param func The given function implementation
52+
*/
53+
TVM_DLL void DefFunction(const String& func_name, const BaseFunc& func);
54+
4055
} // namespace ir
4156
} // namespace ir_builder
4257
} // namespace script

‎include/tvm/script/ir_builder/relax/frame.h

+283
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
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/relax/block_builder.h>
23+
#include <tvm/relax/expr.h>
24+
#include <tvm/script/ir_builder/base.h>
25+
#include <tvm/script/ir_builder/ir/frame.h>
26+
#include <tvm/script/ir_builder/ir/ir.h>
27+
28+
namespace tvm {
29+
namespace script {
30+
namespace ir_builder {
31+
namespace relax {
32+
33+
/*! \brief The base ir_builder frame for the relax dialect. */
34+
class RelaxFrameNode : public IRBuilderFrameNode {
35+
public:
36+
void VisitAttrs(tvm::AttrVisitor* v) { IRBuilderFrameNode::VisitAttrs(v); }
37+
38+
static constexpr const char* _type_key = "script.ir_builder.relax.RelaxFrame";
39+
TVM_DECLARE_BASE_OBJECT_INFO(RelaxFrameNode, IRBuilderFrameNode);
40+
};
41+
42+
class RelaxFrame : public IRBuilderFrame {
43+
public:
44+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RelaxFrame, IRBuilderFrame, RelaxFrameNode);
45+
46+
protected:
47+
RelaxFrame() = default;
48+
};
49+
50+
/*! \brief The base ir_builder frame for frames with SeqExpr
51+
i.e. Functions, If branches
52+
*/
53+
class SeqExprFrameNode : public RelaxFrameNode {
54+
public:
55+
/*! \brief The binding blocks inside the frame. */
56+
Array<tvm::relax::BindingBlock> binding_blocks;
57+
/*! \brief The frame output expr. `NullOpt` when undefined. */
58+
Optional<tvm::relax::Expr> output;
59+
60+
void VisitAttrs(tvm::AttrVisitor* v) {
61+
RelaxFrameNode::VisitAttrs(v);
62+
v->Visit("binding_blocks", &binding_blocks);
63+
v->Visit("output", &output);
64+
}
65+
66+
static constexpr const char* _type_key = "script.ir_builder.relax.SeqExprFrame";
67+
TVM_DECLARE_BASE_OBJECT_INFO(SeqExprFrameNode, RelaxFrameNode);
68+
69+
public:
70+
void ExitWithScope() override;
71+
};
72+
73+
class SeqExprFrame : public RelaxFrame {
74+
public:
75+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(SeqExprFrame, RelaxFrame, SeqExprFrameNode);
76+
};
77+
78+
/*! \brief The ir_builder frame for the relax function. */
79+
class FunctionFrameNode : public SeqExprFrameNode {
80+
public:
81+
/*!
82+
* \brief The function name.
83+
* \note The name will not be specified in constructor, so it is "Optional",
84+
* However, we must specify the name by `R.func_name` before exit this frame.
85+
*/
86+
Optional<String> name;
87+
/*! \brief The function params. */
88+
Array<tvm::relax::Var> params;
89+
/*!
90+
* \brief The function return type.
91+
* \note Usually the function return type can be deduced by the function body.
92+
* But we can use this field to specify a more "accurate" return type.
93+
* i.e. If the `ret_type` is None, try to use the deduced type from body
94+
* If the `ret_type` is not None, check the deduced type is a base type of the given one.
95+
*/
96+
Optional<Type> ret_type;
97+
/*! \brief The function attributes. */
98+
Map<String, ObjectRef> attrs;
99+
/*! \brief The block builder to create Relax function. */
100+
tvm::relax::BlockBuilder block_builder;
101+
102+
void VisitAttrs(tvm::AttrVisitor* v) {
103+
SeqExprFrameNode::VisitAttrs(v);
104+
v->Visit("name", &name);
105+
v->Visit("params", &params);
106+
v->Visit("ret_type", &ret_type);
107+
v->Visit("attrs", &attrs);
108+
v->Visit("binding_blocks", &binding_blocks);
109+
v->Visit("output", &output);
110+
// `block_builder` is not visited.
111+
}
112+
113+
static constexpr const char* _type_key = "script.ir_builder.relax.FunctionFrame";
114+
TVM_DECLARE_FINAL_OBJECT_INFO(FunctionFrameNode, SeqExprFrameNode);
115+
116+
public:
117+
void ExitWithScope() final;
118+
};
119+
120+
class FunctionFrame : public SeqExprFrame {
121+
public:
122+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionFrame, SeqExprFrame, FunctionFrameNode);
123+
};
124+
125+
/*! \brief The ir_builder frame for relax binding blocks. */
126+
class BlockFrameNode : public RelaxFrameNode {
127+
public:
128+
/*! \brief The flag that indicates whether the block is a dataflow block. */
129+
bool is_dataflow;
130+
/*! \brief The variables emitted in this block. */
131+
Array<tvm::relax::Var> emitted_vars;
132+
/*!
133+
* \brief (Only used for a dataflow block.) A boolean indicating if the dataflow block is ended of
134+
* construction. If it is true, any new binding trying to be emitted into this block will cause an
135+
* error.
136+
*/
137+
bool block_ended;
138+
139+
void VisitAttrs(tvm::AttrVisitor* v) {
140+
RelaxFrameNode::VisitAttrs(v);
141+
v->Visit("is_dataflow", &is_dataflow);
142+
v->Visit("emitted_vars", &emitted_vars);
143+
v->Visit("block_ended", &block_ended);
144+
}
145+
146+
static constexpr const char* _type_key = "script.ir_builder.relax.BlockFrame";
147+
TVM_DECLARE_FINAL_OBJECT_INFO(BlockFrameNode, RelaxFrameNode);
148+
149+
public:
150+
void EnterWithScope() final;
151+
void ExitWithScope() final;
152+
};
153+
154+
class BlockFrame : public RelaxFrame {
155+
public:
156+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockFrame, RelaxFrame, BlockFrameNode);
157+
};
158+
159+
/*!
160+
* \brief A frame that represents if statement.
161+
*
162+
* \sa IfFrame
163+
*/
164+
class IfFrameNode : public RelaxFrameNode {
165+
public:
166+
/*! \brief The condition of the if statement. */
167+
tvm::relax::Expr condition;
168+
/*! \brief The Bindings in the true branch. */
169+
Optional<tvm::relax::Expr> then_expr;
170+
/*! \brief The Bindings in the false branch. */
171+
Optional<tvm::relax::Expr> else_expr;
172+
/*! \brief The Binding var. */
173+
tvm::relax::Var var;
174+
/*! \brief The binding var name. */
175+
String var_name;
176+
177+
void VisitAttrs(tvm::AttrVisitor* v) {
178+
RelaxFrameNode::VisitAttrs(v);
179+
v->Visit("condition", &condition);
180+
v->Visit("then_expr", &then_expr);
181+
v->Visit("else_expr", &else_expr);
182+
v->Visit("var", &var);
183+
v->Visit("var_name", &var_name);
184+
}
185+
186+
static constexpr const char* _type_key = "script.ir_builder.relax.IfFrame";
187+
TVM_DECLARE_FINAL_OBJECT_INFO(IfFrameNode, RelaxFrameNode);
188+
189+
public:
190+
/*!
191+
* \brief The method called when entering RAII scope.
192+
* \sa tvm::support::With
193+
*/
194+
void EnterWithScope() final;
195+
/*!
196+
* \brief The method called when exiting RAII scope.
197+
* \sa tvm::support::With
198+
*/
199+
void ExitWithScope() final;
200+
};
201+
202+
/*!
203+
* \brief Managed reference to IfFrameNode.
204+
*
205+
* \sa IfFrameNode
206+
*/
207+
class IfFrame : public RelaxFrame {
208+
public:
209+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IfFrame, RelaxFrame, IfFrameNode);
210+
};
211+
212+
/*!
213+
* \brief A frame that represents then.
214+
*
215+
* \sa ThenFrame
216+
*/
217+
class ThenFrameNode : public SeqExprFrameNode {
218+
public:
219+
static constexpr const char* _type_key = "script.ir_builder.relax.ThenFrame";
220+
TVM_DECLARE_FINAL_OBJECT_INFO(ThenFrameNode, SeqExprFrameNode);
221+
222+
public:
223+
/*!
224+
* \brief The method called when entering RAII scope.
225+
* \sa tvm::support::With
226+
*/
227+
void EnterWithScope() final;
228+
/*!
229+
* \brief The method called when exiting RAII scope.
230+
* \sa tvm::support::With
231+
*/
232+
void ExitWithScope() final;
233+
};
234+
235+
/*!
236+
* \brief Managed reference to ThenFrameNode.
237+
*
238+
* \sa ThenFrameNode
239+
*/
240+
class ThenFrame : public SeqExprFrame {
241+
public:
242+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ThenFrame, SeqExprFrame, ThenFrameNode);
243+
};
244+
245+
/*!
246+
* \brief A frame that represents else.
247+
*
248+
* \sa ElseFrame
249+
*/
250+
class ElseFrameNode : public SeqExprFrameNode {
251+
public:
252+
static constexpr const char* _type_key = "script.ir_builder.relax.ElseFrame";
253+
TVM_DECLARE_FINAL_OBJECT_INFO(ElseFrameNode, SeqExprFrameNode);
254+
255+
public:
256+
/*!
257+
* \brief The method called when entering RAII scope.
258+
* \sa tvm::support::With
259+
*/
260+
void EnterWithScope() final;
261+
/*!
262+
* \brief The method called when exiting RAII scope.
263+
* \sa tvm::support::With
264+
*/
265+
void ExitWithScope() final;
266+
};
267+
268+
/*!
269+
* \brief Managed reference to ElseFrameNode.
270+
*
271+
* \sa ElseFrameNode
272+
*/
273+
class ElseFrame : public SeqExprFrame {
274+
public:
275+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ElseFrame, SeqExprFrame, ElseFrameNode);
276+
};
277+
278+
} // namespace relax
279+
} // namespace ir_builder
280+
} // namespace script
281+
} // namespace tvm
282+
283+
#endif // TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_

0 commit comments

Comments
 (0)