1 | |
2 | /* |
3 | * Licensed to the Apache Software Foundation (ASF) under one |
4 | * or more contributor license agreements. See the NOTICE file |
5 | * distributed with this work for additional information |
6 | * regarding copyright ownership. The ASF licenses this file |
7 | * to you under the Apache License, Version 2.0 (the |
8 | * "License"); you may not use this file except in compliance |
9 | * with the License. You may obtain a copy of the License at |
10 | * |
11 | * http://www.apache.org/licenses/LICENSE-2.0 |
12 | * |
13 | * Unless required by applicable law or agreed to in writing, |
14 | * software distributed under the License is distributed on an |
15 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
16 | * KIND, either express or implied. See the License for the |
17 | * specific language governing permissions and limitations |
18 | * under the License. |
19 | */ |
20 | |
21 | /*! |
22 | * \file codegen_source_base.h |
23 | * \brief Common utilities to source code in text form. |
24 | */ |
25 | #ifndef TVM_TARGET_SOURCE_CODEGEN_SOURCE_BASE_H_ |
26 | #define TVM_TARGET_SOURCE_CODEGEN_SOURCE_BASE_H_ |
27 | |
28 | #include <tvm/ir/name_supply.h> |
29 | #include <tvm/runtime/metadata.h> |
30 | #include <tvm/target/codegen.h> |
31 | #include <tvm/tir/expr.h> |
32 | #include <tvm/tir/op.h> |
33 | |
34 | #include <functional> |
35 | #include <string> |
36 | #include <unordered_map> |
37 | #include <vector> |
38 | |
39 | #include "../../runtime/meta_data.h" |
40 | |
41 | namespace tvm { |
42 | namespace codegen { |
43 | |
44 | /*! |
45 | * \brief A base class to generate source code. |
46 | * Contains helper utilities to generate nest and ssa form. |
47 | */ |
48 | class CodeGenSourceBase { |
49 | public: |
50 | virtual ~CodeGenSourceBase() = default; |
51 | /*! |
52 | * \brief Register constant value appeared in expresion tree |
53 | * This avoid generated a ssa id for each appearance of the value |
54 | * \param value The constant value. |
55 | */ |
56 | void MarkConst(std::string value); |
57 | /*! |
58 | * Print Type representation of type type. |
59 | * \param t The type representation. |
60 | * \param os The stream to print the ctype into |
61 | */ |
62 | virtual void PrintType(DataType type, std::ostream& os); // NOLINT(*) |
63 | /*! |
64 | * Print Type representation of type type. |
65 | * \param type The type representation. |
66 | * \param os The stream to print the ctype into |
67 | */ |
68 | virtual void PrintType(const Type& type, std::ostream& os); // NOLINT(*) |
69 | |
70 | protected: |
71 | /*! \brief entry in ssa assign map */ |
72 | struct SSAEntry { |
73 | /*! \brief The value id */ |
74 | std::string vid; |
75 | /*! \brief The scope id, used to check if this entry is invalid. */ |
76 | int scope_id; |
77 | }; |
78 | /*! \brief Clear the states that might relates to function generation */ |
79 | void ClearFuncState(); |
80 | /*! \brief print the current indented value */ |
81 | void PrintIndent(); |
82 | /*! |
83 | * \brief Allocate a variable name for a newly defined var. |
84 | * \param v The variable. |
85 | * \return the variable name. |
86 | */ |
87 | std::string AllocVarID(const tir::VarNode* v); |
88 | /*! |
89 | * \brief Get a variable name. |
90 | * \param v The variable. |
91 | * \return the variable name. |
92 | */ |
93 | std::string GetVarID(const tir::VarNode* v) const; |
94 | /*! |
95 | * \brief Get the SSA ID corresponds to src |
96 | * If necessary, generate new assignment |
97 | * \param src The source expression |
98 | * \param t The type of the expression. |
99 | */ |
100 | std::string SSAGetID(std::string src, DataType t); |
101 | /*! |
102 | * \brief mark the beginning of a new scope |
103 | * \return The scope id. |
104 | */ |
105 | int BeginScope(); |
106 | /*! |
107 | * \brief mark the end of an old scope. |
108 | * \param scope_id The scope id to be ended. |
109 | */ |
110 | void EndScope(int scope_id); |
111 | /*! |
112 | * \brief Print assignment of src to the id in ssa entry. |
113 | * \param target id of target variable. |
114 | * \param src The source expression. |
115 | * \param t The type of target. |
116 | */ |
117 | virtual void PrintSSAAssign(const std::string& target, const std::string& src, DataType t) = 0; |
118 | |
119 | /*! \brief the declaration stream */ |
120 | std::ostringstream decl_stream; |
121 | /*! \brief the stream to be printed */ |
122 | std::ostringstream stream; |
123 | /*! \brief the forward declaration stream */ |
124 | std::ostringstream fwd_decl_stream; |
125 | /*! \brief name of each variable */ |
126 | std::unordered_map<const tir::VarNode*, std::string> var_idmap_; |
127 | /*! \brief NameSupply for allocation */ |
128 | NameSupply name_supply_ = NameSupply("" ); |
129 | |
130 | private: |
131 | /*! \brief assignment map of ssa */ |
132 | std::unordered_map<std::string, SSAEntry> ssa_assign_map_; |
133 | /*! \brief array to check whether we are inside certain scope */ |
134 | std::vector<bool> scope_mark_; |
135 | /*! \brief The current indentation value */ |
136 | int indent_{0}; |
137 | }; |
138 | |
139 | /*! |
140 | * \brief Create a source module for viewing. |
141 | * \param code The code to be viewed. |
142 | * \param fmt The code. format. |
143 | */ |
144 | runtime::Module SourceModuleCreate(std::string code, std::string fmt); |
145 | |
146 | /*! |
147 | * \brief Create a C source module for viewing and compiling GCC code. |
148 | * \param code The code to be viewed. |
149 | * \param fmt The code format. |
150 | * \param func_names The name of functions inside the runtime module. |
151 | * \param const_vars. The constant variables that the c source module needs. |
152 | * \return The created module. |
153 | */ |
154 | runtime::Module CSourceModuleCreate(const String& code, const String& fmt, |
155 | const Array<String>& func_names, |
156 | const Array<String>& const_vars = {}); |
157 | |
158 | /*! |
159 | * \brief Wrap the submodules in a metadata module. |
160 | * \param params The variable to constant mapping that is collected by the host |
161 | * module. |
162 | * \param target_module The main TIR-lowered internal runtime module |
163 | * \param modules All the external modules that needs to be imported inside the metadata module(s). |
164 | * \param target The target that all the modules are compiled for |
165 | * \param metadata Metadata which should be exported to the runtime. |
166 | * \return The wrapped module. |
167 | */ |
168 | runtime::Module CreateMetadataModule( |
169 | const std::unordered_map<std::string, runtime::NDArray>& params, runtime::Module target_module, |
170 | const Array<runtime::Module>& ext_modules, Target target, runtime::metadata::Metadata metadata); |
171 | |
172 | /*! |
173 | * \brief Create a source module for viewing and limited saving for device. |
174 | * \param data The code data to be viewed. |
175 | * \param fmt The code. format. |
176 | * \param fmap The map function information map of each function. |
177 | * \param type_key The type_key of the runtime module of this source code |
178 | * \param fget_source a closure to replace default get source behavior. |
179 | */ |
180 | runtime::Module DeviceSourceModuleCreate( |
181 | std::string data, std::string fmt, std::unordered_map<std::string, runtime::FunctionInfo> fmap, |
182 | std::string type_key, std::function<std::string(const std::string&)> fget_source = nullptr); |
183 | |
184 | /*! |
185 | * \brief Wrap the submodules that are to be wrapped in a c-source metadata module for C runtime. |
186 | * \param modules The modules to be wrapped. |
187 | * \param target the target the modules are compiled for. |
188 | * \param metadata the metadata needed for code generation. |
189 | * \return The wrapped module. |
190 | */ |
191 | runtime::Module CreateCSourceCrtMetadataModule(const Array<runtime::Module>& modules, Target target, |
192 | runtime::metadata::Metadata metadata); |
193 | |
194 | } // namespace codegen |
195 | } // namespace tvm |
196 | #endif // TVM_TARGET_SOURCE_CODEGEN_SOURCE_BASE_H_ |
197 | |