tvm
module.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 
24 #ifndef TVM_IR_MODULE_H_
25 #define TVM_IR_MODULE_H_
26 
27 #include <tvm/ffi/container/array.h>
28 #include <tvm/ffi/container/map.h>
29 #include <tvm/ffi/reflection/registry.h>
30 #include <tvm/ffi/string.h>
31 #include <tvm/ir/expr.h>
32 #include <tvm/ir/function.h>
33 #include <tvm/ir/global_info.h>
34 #include <tvm/ir/source_map.h>
35 #include <tvm/ir/type.h>
36 
37 #include <string>
38 #include <unordered_map>
39 #include <unordered_set>
40 #include <utility>
41 #include <vector>
42 
43 namespace tvm {
44 
45 class IRModule;
46 
57 class IRModuleNode : public Object {
58  public:
60  Map<GlobalVar, BaseFunc> functions;
63  /* \brief Additional attributes storing meta-data about the module. */
66  Map<String, Array<GlobalInfo>> global_infos;
71  Map<String, GlobalVar> global_var_map_;
72 
92  template <typename TObjectRef>
93  Optional<TObjectRef> GetAttr(
94  const std::string& attr_key,
95  Optional<TObjectRef> default_value = Optional<TObjectRef>(std::nullopt)) const {
96  return attrs.GetAttr(attr_key, default_value);
97  }
98  // variant that uses TObjectRef to enable implicit conversion to default value.
99  template <typename TObjectRef>
100  Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
101  return GetAttr<TObjectRef>(attr_key, Optional<TObjectRef>(default_value));
102  }
103 
108  DictAttrs GetAttrs() const { return attrs; }
109 
129  bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); }
130 
132 
133  static void RegisterReflection() {
134  namespace refl = tvm::ffi::reflection;
135  refl::ObjectDef<IRModuleNode>()
136  .def_ro("functions", &IRModuleNode::functions)
137  .def_ro("global_var_map_", &IRModuleNode::global_var_map_)
138  .def_ro("source_map", &IRModuleNode::source_map)
139  .def_ro("attrs", &IRModuleNode::attrs)
140  .def_ro("global_infos", &IRModuleNode::global_infos);
141  // register custom structural equal and hash.
142  refl::TypeAttrDef<IRModuleNode>()
143  .def("__s_equal__", &IRModuleNode::SEqual)
144  .def("__s_hash__", &IRModuleNode::SHash);
145  }
146 
147  TVM_DLL bool SEqual(const IRModuleNode* other,
148  ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
149  TVM_DLL uint64_t SHash(uint64_t init_hash,
150  ffi::TypedFunction<uint64_t(AnyView, uint64_t, bool)> hash) const;
151 
159  TVM_DLL void Add(const GlobalVar& var, const BaseFunc& func, bool update = false);
160 
168  TVM_DLL void AddUnchecked(const GlobalVar& var, const BaseFunc& func);
169 
175  TVM_DLL void Update(const GlobalVar& var, const BaseFunc& func);
176 
182  TVM_DLL void UpdateGlobalInfo(const String& name, const Array<GlobalInfo>& info);
183 
188  TVM_DLL void Remove(const GlobalVar& var);
189 
195  TVM_DLL bool ContainGlobalVar(const String& name) const;
196 
202  TVM_DLL GlobalVar GetGlobalVar(const String& str) const;
203 
209  TVM_DLL Array<GlobalVar> GetGlobalVars() const;
210 
216  TVM_DLL BaseFunc Lookup(const GlobalVar& var) const;
217 
223  TVM_DLL BaseFunc Lookup(const String& name) const;
224 
230  TVM_DLL void Update(const IRModule& other);
231 
240  TVM_DLL std::unordered_set<String> Imports() const;
241 
243 
244  static constexpr const char* _type_key = "ir.IRModule";
245  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
246 
248 
249  private:
250  friend class IRModule;
251 };
252 
257 class IRModule : public ObjectRef {
258  public:
266  TVM_DLL explicit IRModule(Map<GlobalVar, BaseFunc> functions, SourceMap map = {},
267  DictAttrs attrs = DictAttrs(),
268  Map<String, Array<GlobalInfo>> global_infos = {});
269 
276  explicit IRModule(ObjectPtr<Object> n) : ObjectRef(n) {}
279  auto* ptr = get_mutable();
280  ICHECK(ptr != nullptr);
281  return static_cast<IRModuleNode*>(ptr);
282  }
283 
288  TVM_DLL static IRModule FromExpr(const RelaxExpr& expr,
289  const Map<GlobalVar, BaseFunc>& global_funcs = {});
290 
297 
300 
301  // allow copy on write.
303 };
304 
305 namespace attr {
306 
307 // Following are attributes for IRModule only.
308 
314 constexpr const char* kModuleName = "mod_name";
315 
316 /*
317  * \brief All the runtime::NDArrays extracted from PrimFunc tir::AllocateConst nodes. The
318  * node will record the index into this array. See also kConstNameToConstant below, which is
319  * the analog for Realy Functions.
320  *
321  * Type: Array<runtime::NDArray>
322  */
323 constexpr const char* kConstants = "constants";
324 
331 constexpr const char* kExternalMods = "external_mods";
332 
360 constexpr const char* kSystemLibPrefix = "system_lib_prefix";
361 
370 constexpr const char* kConstNameToConstant = "const_name_to_constant";
371 
372 } // namespace attr
373 } // namespace tvm
374 #endif // TVM_IR_MODULE_H_
Managed reference to BaseFuncNode.
Definition: function.h:234
Managed reference to DictAttrsNode.
Definition: attrs.h:166
bool HasNonzeroAttr(const std::string &attr_key) const
Check whether the function has an non-zero integer attr.
Definition: attrs.h:233
Optional< TObjectRef > GetAttr(const std::string &attr_key, Optional< TObjectRef > default_value=Optional< TObjectRef >(std::nullopt)) const
Get a function attribute.
Definition: attrs.h:197
Managed reference to GlobalVarNode.
Definition: expr.h:489
IRModule that holds functions and type definitions.
Definition: module.h:57
bool HasNonzeroAttr(const std::string &attr_key) const
Check whether the module has an non-zero integer attr.
Definition: module.h:129
Map< String, GlobalVar > global_var_map_
A map from string names to global variables that ensures global uniqueness.
Definition: module.h:71
static void RegisterReflection()
Definition: module.h:133
IRModuleNode()
Definition: module.h:131
void Remove(const GlobalVar &var)
Remove a function from the global environment.
Array< GlobalVar > GetGlobalVars() const
Collect all global vars defined in this module, ordered by the global variable name.
GlobalVar GetGlobalVar(const String &str) const
Lookup a global function by its variable.
void AddUnchecked(const GlobalVar &var, const BaseFunc &func)
Add a function to the global environment.
TVM_DECLARE_FINAL_OBJECT_INFO(IRModuleNode, Object)
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: module.h:245
static constexpr const char * _type_key
Definition: module.h:244
BaseFunc Lookup(const String &name) const
Look up a global function by its string name.
IRModule ShallowCopy()
Create a shallow copy of this IRModule.
Map< String, Array< GlobalInfo > > global_infos
Globally static object that are referred by the IR itself.
Definition: module.h:66
Map< GlobalVar, BaseFunc > functions
A map from ids to all global functions.
Definition: module.h:60
BaseFunc Lookup(const GlobalVar &var) const
Look up a global function by its variable.
std::unordered_set< String > Imports() const
The set of imported files.
void Update(const IRModule &other)
Update the functions inside this environment by functions in another environment.
DictAttrs GetAttrs() const
Get the metadata attributes.
Definition: module.h:108
Optional< TObjectRef > GetAttr(const std::string &attr_key, TObjectRef default_value) const
Definition: module.h:100
void Add(const GlobalVar &var, const BaseFunc &func, bool update=false)
Add a function to the global environment.
SourceMap source_map
The source map for the module.
Definition: module.h:62
void Update(const GlobalVar &var, const BaseFunc &func)
Update a function in the global environment.
bool SEqual(const IRModuleNode *other, ffi::TypedFunction< bool(AnyView, AnyView, bool, AnyView)> equal) const
bool ContainGlobalVar(const String &name) const
Check if the global_var_map_ contains a global variable.
uint64_t SHash(uint64_t init_hash, ffi::TypedFunction< uint64_t(AnyView, uint64_t, bool)> hash) const
TVM_OBJECT_ENABLE_SCRIPT_PRINTER()
Optional< TObjectRef > GetAttr(const std::string &attr_key, Optional< TObjectRef > default_value=Optional< TObjectRef >(std::nullopt)) const
Get a module attribute.
Definition: module.h:93
DictAttrs attrs
Definition: module.h:64
void UpdateGlobalInfo(const String &name, const Array< GlobalInfo > &info)
Update an array of global infos in the global environment.
Managed reference class to IRModuleNode.
Definition: module.h:257
IRModuleNode * operator->() const
Definition: module.h:278
IRModule(ObjectPtr< Object > n)
constructor
Definition: module.h:276
static IRModule FromExpr(const RelaxExpr &expr, const Map< GlobalVar, BaseFunc > &global_funcs={})
As for FromExprInContext, but assuming expr is bound to 'main' and no imports.
IRModule(Map< GlobalVar, BaseFunc > functions, SourceMap map={}, DictAttrs attrs=DictAttrs(), Map< String, Array< GlobalInfo >> global_infos={})
constructor
TVM_DEFINE_OBJECT_REF_COW_METHOD(IRModuleNode)
IRModule()
default constructor
Definition: module.h:271
IRModule ShallowCopyIRModule(IRModule mod)
Create a shallow copy of an IRModule.
Managed reference to RelaxExprNode.
Definition: expr.h:446
Definition: source_map.h:212
Base expr nodes in TVM.
Function nodes.
GlobalInfo are globally static object that are referred by the IR itself.
IR/AST nodes for the unified type system in TVM.
constexpr const char * kConstants
Definition: module.h:323
constexpr const char * kConstNameToConstant
All the named runtime::NDArrays accumulated during compilation by external codegen....
Definition: module.h:370
constexpr const char * kModuleName
Name of the module.
Definition: module.h:314
constexpr const char * kExternalMods
All the runtime::Modules accumulated during compilation by external codegen. These modules must be ei...
Definition: module.h:331
constexpr const char * kSystemLibPrefix
A prefix for generating C symbols system lib creation.
Definition: module.h:360
Definition: repr_printer.h:91
IRModuleFrame IRModule()
The IRModule declaration statement.
Definition: module.h:250
Var var(std::string name_hint, DataType t=DataType::Int(32))
Construct a new Var expression.
tvm::PrimExpr mod(const tvm::PrimExpr &a, const tvm::PrimExpr &b)
Definition: broadcast.h:306
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
PrimExpr equal(PrimExpr a, PrimExpr b, Span span=Span())
equal
A map from source names to source code.