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  ffi::Map<GlobalVar, BaseFunc> functions;
63  /* \brief Additional attributes storing meta-data about the module. */
66  ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
71  ffi::Map<ffi::String, GlobalVar> global_var_map_;
72 
92  template <typename TObjectRef>
93  ffi::Optional<TObjectRef> GetAttr(
94  const std::string& attr_key,
95  ffi::Optional<TObjectRef> default_value = ffi::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  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
101  return GetAttr<TObjectRef>(attr_key, ffi::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 ffi::String& name, const ffi::Array<GlobalInfo>& info);
183 
188  TVM_DLL void Remove(const GlobalVar& var);
189 
195  TVM_DLL bool ContainGlobalVar(const ffi::String& name) const;
196 
202  TVM_DLL GlobalVar GetGlobalVar(const ffi::String& str) const;
203 
209  TVM_DLL ffi::Array<GlobalVar> GetGlobalVars() const;
210 
216  TVM_DLL BaseFunc Lookup(const GlobalVar& var) const;
217 
223  TVM_DLL BaseFunc Lookup(const ffi::String& name) const;
224 
230  TVM_DLL void Update(const IRModule& other);
231 
240  TVM_DLL std::unordered_set<ffi::String> Imports() const;
241 
243 
244  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
245 
247 
248  private:
249  friend class IRModule;
250 };
251 
256 class IRModule : public ObjectRef {
257  public:
265  TVM_DLL explicit IRModule(ffi::Map<GlobalVar, BaseFunc> functions, SourceMap map = {},
266  DictAttrs attrs = DictAttrs(),
267  ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});
268 
270  IRModule() : IRModule(ffi::Map<GlobalVar, BaseFunc>({})) {}
275  explicit IRModule(ObjectPtr<IRModuleNode> n) : ObjectRef(n) {}
279  explicit IRModule(ffi::UnsafeInit tag) : ObjectRef(tag) {}
282  auto* ptr = get_mutable();
283  ICHECK(ptr != nullptr);
284  return static_cast<IRModuleNode*>(ptr);
285  }
286 
291  TVM_DLL static IRModule FromExpr(const RelaxExpr& expr,
292  const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});
293 
300 
303 
304  // allow copy on write.
306 };
307 
308 namespace attr {
309 
310 // Following are attributes for IRModule only.
311 
317 constexpr const char* kModuleName = "mod_name";
318 
319 /*
320  * \brief All the runtime::Tensors extracted from PrimFunc tir::AllocateConst nodes. The
321  * node will record the index into this array. See also kConstNameToConstant below, which is
322  * the analog for Realy Functions.
323  *
324  * Type: ffi::Array<runtime::Tensor>
325  */
326 constexpr const char* kConstants = "constants";
327 
334 constexpr const char* kExternalMods = "external_mods";
335 
363 constexpr const char* kSystemLibPrefix = "system_lib_prefix";
364 
373 constexpr const char* kConstNameToConstant = "const_name_to_constant";
374 
375 } // namespace attr
376 } // namespace tvm
377 #endif // TVM_IR_MODULE_H_
Managed reference to BaseFuncNode.
Definition: function.h:233
Managed reference to DictAttrsNode.
Definition: attrs.h:162
bool HasNonzeroAttr(const std::string &attr_key) const
Check whether the function has an non-zero integer attr.
Definition: attrs.h:233
ffi::Optional< TObjectRef > GetAttr(const std::string &attr_key, ffi::Optional< TObjectRef > default_value=ffi::Optional< TObjectRef >(std::nullopt)) const
Get a function attribute.
Definition: attrs.h:197
Managed reference to GlobalVarNode.
Definition: expr.h:481
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
ffi::Array< GlobalVar > GetGlobalVars() const
Collect all global vars defined in this module, ordered by the global variable name.
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IRModule", IRModuleNode, Object)
static void RegisterReflection()
Definition: module.h:133
IRModuleNode()
Definition: module.h:131
void Remove(const GlobalVar &var)
Remove a function from the global environment.
bool ContainGlobalVar(const ffi::String &name) const
Check if the global_var_map_ contains a global variable.
ffi::Map< ffi::String, GlobalVar > global_var_map_
A map from string names to global variables that ensures global uniqueness.
Definition: module.h:71
void UpdateGlobalInfo(const ffi::String &name, const ffi::Array< GlobalInfo > &info)
Update an array of global infos in the global environment.
void AddUnchecked(const GlobalVar &var, const BaseFunc &func)
Add a function to the global environment.
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: module.h:244
std::unordered_set< ffi::String > Imports() const
The set of imported files.
ffi::Map< ffi::String, ffi::Array< GlobalInfo > > global_infos
Globally static object that are referred by the IR itself.
Definition: module.h:66
IRModule ShallowCopy()
Create a shallow copy of this IRModule.
BaseFunc Lookup(const GlobalVar &var) const
Look up a global function by its variable.
void Update(const IRModule &other)
Update the functions inside this environment by functions in another environment.
ffi::Optional< TObjectRef > GetAttr(const std::string &attr_key, TObjectRef default_value) const
Definition: module.h:100
BaseFunc Lookup(const ffi::String &name) const
Look up a global function by its string name.
DictAttrs GetAttrs() const
Get the metadata attributes.
Definition: module.h:108
ffi::Map< GlobalVar, BaseFunc > functions
A map from ids to all global functions.
Definition: module.h:60
ffi::Optional< TObjectRef > GetAttr(const std::string &attr_key, ffi::Optional< TObjectRef > default_value=ffi::Optional< TObjectRef >(std::nullopt)) const
Get a module attribute.
Definition: module.h:93
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
uint64_t SHash(uint64_t init_hash, ffi::TypedFunction< uint64_t(AnyView, uint64_t, bool)> hash) const
TVM_OBJECT_ENABLE_SCRIPT_PRINTER()
GlobalVar GetGlobalVar(const ffi::String &str) const
Lookup a global function by its variable.
DictAttrs attrs
Definition: module.h:64
Managed reference class to IRModuleNode.
Definition: module.h:256
IRModuleNode * operator->() const
Definition: module.h:281
static IRModule FromExpr(const RelaxExpr &expr, const ffi::Map< GlobalVar, BaseFunc > &global_funcs={})
As for FromExprInContext, but assuming expr is bound to 'main' and no imports.
IRModule(ffi::UnsafeInit tag)
constructor with UnsafeInit
Definition: module.h:279
IRModule(ObjectPtr< IRModuleNode > n)
constructor
Definition: module.h:275
TVM_DEFINE_OBJECT_REF_COW_METHOD(IRModuleNode)
IRModule()
default constructor
Definition: module.h:270
IRModule ShallowCopyIRModule(IRModule mod)
Create a shallow copy of an IRModule.
IRModule(ffi::Map< GlobalVar, BaseFunc > functions, SourceMap map={}, DictAttrs attrs=DictAttrs(), ffi::Map< ffi::String, ffi::Array< GlobalInfo >> global_infos={})
constructor
Managed reference to RelaxExprNode.
Definition: expr.h:439
Definition: source_map.h:205
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:326
constexpr const char * kConstNameToConstant
All the named runtime::Tensors accumulated during compilation by external codegen....
Definition: module.h:373
constexpr const char * kModuleName
Name of the module.
Definition: module.h:317
constexpr const char * kExternalMods
All the runtime::Modules accumulated during compilation by external codegen. These modules must be ei...
Definition: module.h:334
constexpr const char * kSystemLibPrefix
A prefix for generating C symbols system lib creation.
Definition: module.h:363
Definition: repr_printer.h:91
IRModuleFrame IRModule()
The IRModule declaration statement.
Definition: module.h:249
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:308
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.