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 ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt))
const {
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));
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;
182 TVM_DLL
void UpdateGlobalInfo(
const ffi::String& name,
const ffi::Array<GlobalInfo>& info);
240 TVM_DLL std::unordered_set<ffi::String>
Imports()
const;
267 ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});
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);
292 const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});
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
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.