Program Listing for File module.h

↰ Return to documentation for file (tvm/ir/module.h)

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef TVM_IR_MODULE_H_
#define TVM_IR_MODULE_H_

#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/ir/global_info.h>
#include <tvm/ir/location.h>
#include <tvm/ir/type.h>

#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

namespace tvm {

class IRModule;

class IRModuleObj : public ffi::Object {
 public:
  ffi::Map<GlobalVar, BaseFunc> functions;
  /* \brief Additional attributes storing meta-data about the module. */
  DictAttrs attrs;
  ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
  ffi::Map<ffi::String, GlobalVar> global_var_map_;

  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetAttr(
      const std::string& attr_key,
      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
    return attrs.GetAttr(attr_key, default_value);
  }
  template <typename T>
  ffi::Optional<T> GetAttr(ffi::reflection::AttrKey<T> key,
                           ffi::Optional<typename ffi::reflection::AttrKey<T>::ValueType>
                               default_value = std::nullopt) const {
    return attrs.GetAttr(key, std::move(default_value));
  }
  template <typename T>
  ffi::Optional<T> GetAttr(ffi::reflection::AttrKey<T> key,
                           typename ffi::reflection::AttrKey<T>::ValueType default_value) const {
    return GetAttr(key, ffi::Optional<T>(std::move(default_value)));
  }

  // variant that uses TObjectRef to enable implicit conversion to default value.
  template <typename TObjectRef>
  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
    return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
  }

  DictAttrs GetAttrs() const { return attrs; }

  bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); }

  IRModuleObj() = default;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IRModuleObj>()
        .def_ro("functions", &IRModuleObj::functions)
        .def_ro("global_var_map_", &IRModuleObj::global_var_map_)
        .def_ro("attrs", &IRModuleObj::attrs)
        .def_ro("global_infos", &IRModuleObj::global_infos);
    // register custom structural equal and hash.
    refl::TypeAttrDef<IRModuleObj>()
        .def(tvm::ffi::type_attr::kSEqual.name, &IRModuleObj::SEqual)
        .def(tvm::ffi::type_attr::kSHash.name, &IRModuleObj::SHash);
  }

  TVM_TIRX_DLL bool SEqual(const IRModuleObj* other,
                           ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
  TVM_TIRX_DLL int64_t SHash(int64_t init_hash,
                             ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;

  void Add(const GlobalVar& var, const BaseFunc& func, bool update = false) {
    BaseFunc checked_func = func;
    AddUnchecked(var, checked_func);
  }

  TVM_TIRX_DLL void AddUnchecked(const GlobalVar& var, const BaseFunc& func);

  void Update(const GlobalVar& var, const BaseFunc& func) { this->Add(var, func, true); }

  void UpdateGlobalInfo(const ffi::String& name, const ffi::Array<GlobalInfo>& info) {
    this->global_infos.Set(name, info);
  }

  void Remove(const GlobalVar& var) {
    auto functions_node = this->functions.CopyOnWrite();
    functions_node->erase(var);
    auto gvar_node = global_var_map_.CopyOnWrite();
    gvar_node->erase(var->name_hint);
  }

  bool ContainGlobalVar(const ffi::String& name) const {
    return global_var_map_.find(name) != global_var_map_.end();
  }

  TVM_TIRX_DLL GlobalVar GetGlobalVar(const ffi::String& str) const;

  TVM_TIRX_DLL ffi::Array<GlobalVar> GetGlobalVars() const;

  BaseFunc Lookup(const GlobalVar& var) const {
    auto it = functions.find(var);
    TVM_FFI_ICHECK(it != functions.end()) << "There is no definition of " << var;
    return (*it).second;
  }

  BaseFunc Lookup(const ffi::String& name) const {
    GlobalVar id = this->GetGlobalVar(name);
    return this->Lookup(id);
  }

  void Update(const IRModule& other);

  IRModule ShallowCopy();

  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IRModule", IRModuleObj, ffi::Object);

 private:
  friend class IRModule;
};

class IRModule : public ffi::ObjectRef {
 public:
  static constexpr bool _type_is_nullable = false;
  TVM_TIRX_DLL explicit IRModule(ffi::Map<GlobalVar, BaseFunc> functions,
                                 DictAttrs attrs = DictAttrs(),
                                 ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});

  IRModule() : IRModule(ffi::Map<GlobalVar, BaseFunc>({})) {}
  explicit IRModule(ffi::ObjectPtr<IRModuleObj> n) : ffi::ObjectRef(std::move(n)) {
    TVM_FFI_ICHECK(data_ != nullptr);
  }
  explicit IRModule(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
  IRModuleObj* operator->() const {
    auto* ptr = get_mutable();
    TVM_FFI_ICHECK(ptr != nullptr);
    return static_cast<IRModuleObj*>(ptr);
  }

  TVM_TIRX_DLL static IRModule FromExpr(const Expr& expr,
                                        const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});

  using ContainerType = IRModuleObj;

  // allow copy on write.
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IRModuleObj);
};

inline void IRModuleObj::Update(const IRModule& mod) {
  for (auto pair : mod->functions) {
    this->AddUnchecked(pair.first, pair.second);
  }
}

inline IRModule IRModuleObj::ShallowCopy() {
  return IRModule(this->functions, this->attrs, this->global_infos);
}

namespace attr {

// Following are attributes for IRModule only.

inline constexpr ffi::reflection::AttrKey<ffi::Array<ffi::Module>> kExternalMods{"external_mods"};

inline constexpr ffi::reflection::AttrKey<ffi::String> kSystemLibPrefix{"system_lib_prefix"};

}  // namespace attr
}  // namespace tvm
#endif  // TVM_IR_MODULE_H_