24 #ifndef TVM_IR_MODULE_H_
25 #define TVM_IR_MODULE_H_
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>
38 #include <unordered_map>
39 #include <unordered_set>
92 template <
typename TObjectRef>
94 const std::string& attr_key,
95 Optional<TObjectRef> default_value = Optional<TObjectRef>(std::nullopt))
const {
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));
135 refl::ObjectDef<IRModuleNode>()
142 refl::TypeAttrDef<IRModuleNode>()
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;
240 TVM_DLL std::unordered_set<String>
Imports()
const;
268 Map<String, Array<GlobalInfo>> global_infos = {});
276 explicit IRModule(ObjectPtr<Object> n) : ObjectRef(n) {}
279 auto* ptr = get_mutable();
280 ICHECK(ptr !=
nullptr);
289 const Map<GlobalVar, BaseFunc>& global_funcs = {});
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
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.