Construct GlobalVarSupply from IRModule

Author: Florin-Gabriel Blanaru

include/tvm/driver/driver_api.h

@@ -29,6 +29,7 @@
 #ifndef TVM_DRIVER_DRIVER_API_H_
 #define TVM_DRIVER_DRIVER_API_H_
 
+#include <tvm/ir/global_var_supply.h>
 #include <tvm/ir/module.h>
 #include <tvm/ir/transform.h>
 #include <tvm/runtime/packed_func.h>

include/tvm/ir/global_var_supply.h

@@ -24,6 +24,7 @@
 #include <unordered_map>
 
 #include "tvm/ir/expr.h"
+#include "tvm/ir/module.h"
 #include "tvm/ir/name_supply.h"
 
 namespace tvm {
@@ -57,13 +58,12 @@ class GlobalVarSupplyNode : public Object {
 
 class GlobalVarSupply : public ObjectRef {
  public:
-  TVM_DLL explicit GlobalVarSupply(
-      const NameSupply& name_supply = NameSupply::NameSupplyWithPrefix(""),
-      std::unordered_map<std::string, GlobalVar> name_to_var_map = {});
+  TVM_DLL explicit GlobalVarSupply(const NameSupply& name_supply = NameSupply(),
+                                   std::unordered_map<std::string, GlobalVar> name_to_var_map = {});
 
-  TVM_DLL static GlobalVarSupply GlobalVarSupplyFromNameSupply(const NameSupply& name_supply);
+  TVM_DLL explicit GlobalVarSupply(const Array<IRModule>& modules);
 
-  TVM_DLL static GlobalVarSupply EmptySupply();
+  TVM_DLL explicit GlobalVarSupply(const IRModule module);
 
   explicit GlobalVarSupply(ObjectPtr<Object> n) : ObjectRef(n) {}
   /*! \return mutable pointers to the node. */

include/tvm/ir/module.h

@@ -27,7 +27,6 @@
 #include <tvm/ir/adt.h>
 #include <tvm/ir/expr.h>
 #include <tvm/ir/function.h>
-#include <tvm/ir/global_var_supply.h>
 #include <tvm/ir/type.h>
 #include <tvm/parser/source_map.h>
 #include <tvm/runtime/container/array.h>

include/tvm/ir/name_supply.h

@@ -64,15 +64,11 @@ class NameSupplyNode : public Object {
 
 class NameSupply : public ObjectRef {
  public:
-  TVM_DLL NameSupply();
+  TVM_DLL explicit NameSupply();
 
   TVM_DLL explicit NameSupply(const String& prefix,
                               std::unordered_map<std::string, int> name_map = {});
 
-  TVM_DLL static NameSupply NameSupplyWithPrefix(const String& prefix = "");
-
-  TVM_DLL static NameSupply EmptySupply();
-
   explicit NameSupply(ObjectPtr<Object> n) : ObjectRef(n) {}
   /*! \return mutable pointers to the node. */
   NameSupplyNode* operator->() const {

python/tvm/ir/supply.py

@@ -16,7 +16,7 @@
 # under the License.
 """Suppliers that are used to guarantee uniqueness of names and GlobalVars."""
 import tvm
-from tvm import Object
+from tvm import Object, IRModule
 from . import _ffi_api
 
 
@@ -38,6 +38,12 @@ def fresh_name(self, name, add_prefix=True):
     def reserve_name(self, name, add_prefix=True):
         return _ffi_api.NameSupply_ReserveName(self, name, add_prefix)
 
+    def contains_name(self, name, add_prefix=True):
+        return _ffi_api.NameSupply_ContainsName(self, name, add_prefix)
+
+    def clear(self):
+        return _ffi_api.NameSupply_Clear(self)
+
 
 @tvm._ffi.register_object("GlobalVarSupply")
 class GlobalVarSupply(Object):
@@ -48,15 +54,24 @@ class GlobalVarSupply(Object):
 
     Parameters
     ----------
-    name_supply: The NameSupply to be used by this GlobalVarSupply.
+    value: Union[List[IRModule], IRModule, NameSupply]
+        The IRModules used to build this GlobalVarSupply or a NameSupply.
     """
 
-    def __init__(self, name_supply=None):
-        name_supply = name_supply if name_supply is not None else NameSupply("")
-        self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply, name_supply)
+    def __init__(self, value=None):
+        if value is None:
+            name_supply = NameSupply("")
+            self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply_NameSupply, name_supply)
+        elif isinstance(value, (list, tvm.container.Array)):
+            self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply_IRModules, value)
+        elif isinstance(value, IRModule):
+            self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply_IRModule, value)
 
     def fresh_global(self, name, add_prefix=True):
         return _ffi_api.GlobalVarSupply_FreshGlobal(self, name, add_prefix)
 
     def unique_global_for(self, name, add_prefix=True):
         return _ffi_api.GlobalVarSupply_UniqueGlobalFor(self, name, add_prefix)
+
+    def reserve_global(self, global_var, allow_conflict=False):
+        return _ffi_api.GlobalVarSupply_ReserveGlobalVar(self, global_var, allow_conflict)

src/auto_scheduler/feature.cc

@@ -27,6 +27,7 @@
 #include <tvm/auto_scheduler/measure.h>
 #include <tvm/auto_scheduler/measure_record.h>
 #include <tvm/driver/driver_api.h>
+#include <tvm/ir/global_var_supply.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/support/parallel_for.h>
 #include <tvm/te/operation.h>
@@ -1371,8 +1372,7 @@ void GetPerStoreFeaturesWorkerFunc(const SearchTask& task, const State& state, i
     auto pass_ctx = tvm::transform::PassContext::Current();
 
     auto mod = ScheduleToModule(sch, Array<ObjectRef>{tensors.begin(), tensors.end()}, name,
-                                std::unordered_map<te::Tensor, te::Buffer>(),
-                                GlobalVarSupply::EmptySupply());
+                                std::unordered_map<te::Tensor, te::Buffer>(), GlobalVarSupply());
 
     bool disable_vectorize =
         pass_ctx->GetConfig<Bool>("tir.disable_vectorize", Bool(false)).value();

src/contrib/hybrid/codegen_hybrid.h

@@ -24,6 +24,7 @@
 #ifndef TVM_CONTRIB_HYBRID_CODEGEN_HYBRID_H_
 #define TVM_CONTRIB_HYBRID_CODEGEN_HYBRID_H_
 
+#include <tvm/ir/name_supply.h>
 #include <tvm/target/codegen.h>
 #include <tvm/te/operation.h>
 #include <tvm/te/schedule.h>
@@ -146,7 +147,7 @@ class CodeGenHybrid : public ExprFunctor<void(const PrimExpr&, std::ostream&)>,
   /*! \brief Print the current indent spaces. */
   inline void PrintIndent();
   /*! \brief NameSupply for allocated ids.  */
-  NameSupply ids_allocated = NameSupply::EmptySupply();
+  NameSupply ids_allocated = NameSupply();
   /*!
    * \brief Keys are either (tensors, value_index) or (variables, 0).
    *        Values are the corresponding IDs.*/

src/driver/driver_api.cc

@@ -303,8 +303,7 @@ TVM_REGISTER_GLOBAL("driver.schedule_to_module")
           c_binds.insert({kv.first, kv.second});
         }
       }
-      IRModule mod =
-          ScheduleToModule(std::move(sch), args, name, c_binds, GlobalVarSupply::EmptySupply());
+      IRModule mod = ScheduleToModule(std::move(sch), args, name, c_binds, GlobalVarSupply());
       return mod;
     });
 
@@ -367,8 +366,7 @@ TVM_REGISTER_GLOBAL("driver.lower_schedule")
           c_binds.insert({kv.first, kv.second});
         }
       }
-      return LowerSchedule(std::move(sch), args, name, c_binds, GlobalVarSupply::EmptySupply(),
-                           simple_mode);
+      return LowerSchedule(std::move(sch), args, name, c_binds, GlobalVarSupply(), simple_mode);
     });
 
 /**

src/ir/global_var_supply.cc

@@ -26,23 +26,32 @@
 #include "tvm/ir/expr.h"
 
 namespace tvm {
-
 GlobalVarSupply::GlobalVarSupply(const NameSupply& name_supply,
                                  std::unordered_map<std::string, GlobalVar> name_to_var_map) {
   auto n = make_object<GlobalVarSupplyNode>(name_supply);
   n->name_to_var_map_ = std::move(name_to_var_map);
   data_ = std::move(n);
 }
 
-GlobalVarSupply GlobalVarSupply::GlobalVarSupplyFromNameSupply(const NameSupply& name_supply) {
-  auto global_var_supply = GlobalVarSupply(name_supply);
-  return global_var_supply;
+std::string GetModuleName(const IRModule& module) {
+  return module->GetAttr<String>(tvm::attr::kModuleName).value_or("tvmgen_default");
 }
 
-GlobalVarSupply GlobalVarSupply::EmptySupply() {
-  return GlobalVarSupplyFromNameSupply(NameSupply::NameSupplyWithPrefix(""));
+GlobalVarSupply::GlobalVarSupply(const Array<IRModule>& modules) : GlobalVarSupply() {
+  if (!modules.empty()) {
+    IRModule first_mod = modules.front();
+    this->operator->()->name_supply_->prefix_ = GetModuleName(first_mod);
+  }
+  for (auto& mod : modules) {
+    for (auto kv : mod->functions) {
+      this->operator->()->ReserveGlobalVar(kv.first);
+    }
+  }
 }
 
+GlobalVarSupply::GlobalVarSupply(const IRModule module)
+    : GlobalVarSupply(Array<IRModule>{module}) {}
+
 void GlobalVarSupplyNode::ReserveGlobalVar(const GlobalVar& var, bool allow_conflict) {
   name_supply_->ReserveName(var->name_hint, false);
   if (!allow_conflict) {
@@ -79,8 +88,15 @@ GlobalVar GlobalVarSupplyNode::FreshGlobal(String name, bool add_prefix) {
 
 TVM_REGISTER_NODE_TYPE(GlobalVarSupplyNode);
 
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply").set_body_typed([](NameSupply name_supply) {
-  return GlobalVarSupply(name_supply);
+TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_NameSupply")
+    .set_body_typed([](const NameSupply& name_supply) { return GlobalVarSupply(name_supply); });
+
+TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModule").set_body_typed([](IRModule mod) {
+  return GlobalVarSupply(std::move(mod));
+});
+
+TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModules").set_body_typed([](const Array<IRModule>& mods) {
+  return GlobalVarSupply(mods);
 });
 
 TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_FreshGlobal")
@@ -89,4 +105,7 @@ TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_FreshGlobal")
 TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_UniqueGlobalFor")
     .set_body_method<GlobalVarSupply>(&GlobalVarSupplyNode::UniqueGlobalFor);
 
+TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_ReserveGlobalVar")
+    .set_body_method<GlobalVarSupply>(&GlobalVarSupplyNode::ReserveGlobalVar);
+
 }  // namespace tvm

src/ir/module.cc

@@ -40,7 +40,7 @@
 #include <sstream>
 #include <unordered_set>
 
-#include "../relay/backend/supply_provider.h"
+#include "tvm/ir/global_var_supply.h"
 
 namespace tvm {
 
@@ -386,7 +386,7 @@ std::pair<IRModule, GlobalVar> IRModule::FromExprInContext(
   }
 
   GlobalVar main_gv;
-  auto global_var_supply = tvm::BuildGlobalVarSupply(mod);
+  auto global_var_supply = GlobalVarSupply(mod);
   if (gv_name.empty()) {
     // Bind function to 'main' (though rename if would clash with existing 'main').
     main_gv = global_var_supply->FreshGlobal("main", false);