Skip to content
This repository was archived by the owner on May 22, 2023. It is now read-only.

Commit 912d1fc

Browse files
junrushaojunrushao
committed
[TVMScript] New Parser: Part B
1 parent d1fb69d commit 912d1fc

File tree

23 files changed

+3349
-4
lines changed

23 files changed

+3349
-4
lines changed

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

+166
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
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/ir/expr.h>
23+
#include <tvm/ir/function.h>
24+
#include <tvm/node/node.h>
25+
26+
#include <vector>
27+
28+
namespace tvm {
29+
namespace script {
30+
namespace ir_builder {
31+
32+
////////////////////////////// Core Infra: Frame //////////////////////////////
33+
34+
class IRBuilderFrameNode : public runtime::Object {
35+
public:
36+
std::vector<runtime::TypedPackedFunc<void()>> callbacks;
37+
38+
void VisitAttrs(tvm::AttrVisitor* v) {
39+
// `callbacks` is not visited.
40+
}
41+
42+
static constexpr const char* _type_key = "script.ir_builder.IRBuilderFrame";
43+
TVM_DECLARE_BASE_OBJECT_INFO(IRBuilderFrameNode, runtime::Object);
44+
45+
public:
46+
virtual ~IRBuilderFrameNode() = default;
47+
virtual void EnterWithScope();
48+
virtual void ExitWithScope();
49+
50+
void AddCallback(runtime::TypedPackedFunc<void()> callback);
51+
};
52+
53+
class IRBuilderFrame : public runtime::ObjectRef {
54+
public:
55+
virtual ~IRBuilderFrame() = default;
56+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRBuilderFrame, ObjectRef, IRBuilderFrameNode);
57+
58+
protected:
59+
IRBuilderFrame() = default;
60+
61+
public:
62+
inline void EnterWithScope();
63+
inline void ExitWithScope();
64+
};
65+
66+
////////////////////////////// Core Infra: Builder //////////////////////////////
67+
///
68+
class IRBuilderNode : public runtime::Object {
69+
public:
70+
runtime::Array<IRBuilderFrame> frames;
71+
Optional<ObjectRef> result;
72+
73+
void VisitAttrs(tvm::AttrVisitor* v) {
74+
v->Visit("frames", &frames);
75+
v->Visit("result", &result);
76+
}
77+
78+
static constexpr const char* _type_key = "script.ir_builder.IRBuilder";
79+
TVM_DECLARE_FINAL_OBJECT_INFO(IRBuilderNode, runtime::Object);
80+
81+
public:
82+
template <typename TFrame>
83+
inline Optional<TFrame> FindFrame() const;
84+
template <typename TFrame>
85+
inline Optional<TFrame> GetLastFrame() const;
86+
template <typename TObjectRef>
87+
inline TObjectRef Get() const;
88+
};
89+
90+
class IRBuilder : public runtime::ObjectRef {
91+
public:
92+
IRBuilder();
93+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRBuilder, ObjectRef, IRBuilderNode);
94+
95+
public:
96+
void EnterWithScope();
97+
void ExitWithScope();
98+
static IRBuilder Current();
99+
template <class TObjectRef>
100+
inline static TObjectRef Name(String name, TObjectRef obj);
101+
};
102+
103+
////////////////////////////// Details //////////////////////////////
104+
105+
namespace details {
106+
107+
class Namer {
108+
public:
109+
using FType = NodeFunctor<void(const ObjectRef&, String)>;
110+
static FType& vtable();
111+
static void Name(ObjectRef node, String name);
112+
};
113+
114+
} // namespace details
115+
116+
template <class TObjectRef>
117+
inline TObjectRef IRBuilder::Name(String name, TObjectRef obj) {
118+
details::Namer::Name(obj, name);
119+
return Downcast<TObjectRef>(obj);
120+
}
121+
122+
inline void IRBuilderFrame::EnterWithScope() {
123+
ICHECK(data_ != nullptr);
124+
static_cast<IRBuilderFrameNode*>(data_.get())->EnterWithScope();
125+
}
126+
127+
inline void IRBuilderFrame::ExitWithScope() {
128+
ICHECK(data_ != nullptr);
129+
static_cast<IRBuilderFrameNode*>(data_.get())->ExitWithScope();
130+
data_.reset();
131+
}
132+
133+
template <typename TFrame>
134+
inline Optional<TFrame> IRBuilderNode::FindFrame() const {
135+
using TFrameNode = typename TFrame::ContainerType;
136+
for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
137+
if (const TFrameNode* p = (*it).template as<TFrameNode>()) {
138+
return GetRef<TFrame>(p);
139+
}
140+
}
141+
return NullOpt;
142+
}
143+
144+
template <typename TFrame>
145+
inline Optional<TFrame> IRBuilderNode::GetLastFrame() const {
146+
using TFrameNode = typename TFrame::ContainerType;
147+
if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
148+
return Downcast<TFrame>(frames.back());
149+
}
150+
return NullOpt;
151+
}
152+
153+
template <typename TObjectRef>
154+
inline TObjectRef IRBuilderNode::Get() const {
155+
using TObject = typename TObjectRef::ContainerType;
156+
CHECK(result.defined()) << "IndexError: No result exists in IRBuilder yet";
157+
const auto* n = result.as<TObject>();
158+
CHECK(n != nullptr) << "IndexError: IRBuilder result is not of type: " << TObject::_type_key;
159+
return GetRef<TObjectRef>(n);
160+
}
161+
162+
} // namespace ir_builder
163+
} // namespace script
164+
} // namespace tvm
165+
166+
#endif // TVM_SCRIPT_IR_BUILDER_BASE_H_

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

+61
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
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_IR_FRAME_H_
20+
#define TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
21+
22+
#include <tvm/ir/expr.h>
23+
#include <tvm/ir/function.h>
24+
#include <tvm/node/node.h>
25+
#include <tvm/script/ir_builder/base.h>
26+
27+
#include <vector>
28+
29+
namespace tvm {
30+
namespace script {
31+
namespace ir_builder {
32+
33+
class IRModuleFrameNode : public IRBuilderFrameNode {
34+
public:
35+
Array<GlobalVar> global_vars;
36+
Array<BaseFunc> functions;
37+
38+
void VisitAttrs(tvm::AttrVisitor* v) {
39+
IRBuilderFrameNode::VisitAttrs(v);
40+
v->Visit("global_vars", &global_vars);
41+
v->Visit("functions", &functions);
42+
}
43+
44+
static constexpr const char* _type_key = "script.ir_builder.IRModuleFrame";
45+
TVM_DECLARE_FINAL_OBJECT_INFO(IRModuleFrameNode, IRBuilderFrameNode);
46+
47+
public:
48+
void ExitWithScope() final;
49+
};
50+
51+
class IRModuleFrame : public IRBuilderFrame {
52+
public:
53+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRModuleFrame, IRBuilderFrame,
54+
IRModuleFrameNode);
55+
};
56+
57+
} // namespace ir_builder
58+
} // namespace script
59+
} // namespace tvm
60+
61+
#endif // TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_

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

+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
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_IR_IR_H_
20+
#define TVM_SCRIPT_IR_BUILDER_IR_IR_H_
21+
22+
#include <tvm/ir/expr.h>
23+
#include <tvm/ir/function.h>
24+
#include <tvm/node/node.h>
25+
#include <tvm/script/ir_builder/ir/frame.h>
26+
27+
#include <vector>
28+
29+
namespace tvm {
30+
namespace script {
31+
namespace ir_builder {
32+
33+
TVM_DLL IRModuleFrame IRModule();
34+
35+
}
36+
} // namespace script
37+
} // namespace tvm
38+
39+
#endif // TVM_IR_IR_BUILDER_IR_IR_H_

0 commit comments

Comments
 (0)