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
41namespace tvm {
42namespace codegen {
43
44/*!
45 * \brief A base class to generate source code.
46 * Contains helper utilities to generate nest and ssa form.
47 */
48class 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 */
144runtime::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 */
154runtime::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 */
168runtime::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 */
180runtime::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 */
191runtime::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