(tvm) branch unity updated: [Unity] Speed up NormalizeGlobalVar (#16219)

[syfeng]

This is an automated email from the ASF dual-hosted git repository.

syfeng pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new f18d186559 [Unity] Speed up NormalizeGlobalVar (#16219)
f18d186559 is described below

commit f18d18655923c2afc0adb8f52a46e6e7e670c2cb
Author: Siyuan Feng <hzfen...@sjtu.edu.cn>
AuthorDate: Mon Dec 11 10:39:33 2023 +0800

    [Unity] Speed up NormalizeGlobalVar (#16219)
    
    The current implementation of NormalizeGlobalVar couples the
    `AddPublicFunction` and checking if the module needs renaming, which
    cause the modules that don't need renaming still need to run
    unnecessary code, such as `AddFunction` to the new module.
    
    This patch separates the two steps, so that the modules that don't need
    renaming can skip the `AddPublicFunction` step.
---
 src/relax/transform/normalize.cc | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/relax/transform/normalize.cc b/src/relax/transform/normalize.cc
index 6b45a4c8e9..0939674e81 100644
--- a/src/relax/transform/normalize.cc
+++ b/src/relax/transform/normalize.cc
@@ -183,12 +183,13 @@ class GlobalVarNormalizer : private ExprMutator {
   using ExprMutator::VisitExpr_;
 
   IRModule RenameModule() {
-    // Step 1. Add public functions (functions with global_symbol attributes)
-    auto name_changes = AddPublicFunctions();
-    if (!name_changes) {
+    if (!NeedRename()) {
       return module_;
     }
 
+    // Step 1. Add public functions (functions with global_symbol attributes)
+    AddPublicFunctions();
+
     // Step 2. Rename private functions
     AddPrivateFunctions();
 
@@ -210,12 +211,19 @@ class GlobalVarNormalizer : private ExprMutator {
     return module_;
   }
 
-  /**
-   * \brief Add public functions to the builder, and update the name supplier.
-   * \return true if any name changes are made.
-   */
-  bool AddPublicFunctions() {
-    bool name_changes = false;
+  /*! \brief Check if any function needs to be renamed. */
+  bool NeedRename() {
+    for (const auto& [gvar, func] : module_->functions) {
+      auto global_symbol = func->GetAttr<String>("global_symbol");
+      if (global_symbol && global_symbol.value() != gvar->name_hint) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /*! \brief Add public functions to the builder, and update the name 
supplier. */
+  void AddPublicFunctions() {
     for (const auto& [gvar, func] : module_->functions) {
       auto global_symbol = func->GetAttr<String>("global_symbol");
       if (!global_symbol) {
@@ -228,15 +236,10 @@ class GlobalVarNormalizer : private ExprMutator {
       name_supply_->ReserveName(global_symbol_value);
       auto new_gvar = builder_->AddFunction(func, global_symbol_value);
       gvar_map_.Set(gvar, new_gvar);
-
-      if (global_symbol.value() != gvar->name_hint) {
-        name_changes = true;
-      }
     }
-    return name_changes;
   }
 
-  /**
+  /*!
    * \brief Add private functions to the builder with names provided by name 
supplier. Renaming may
    * happen if the name of any function conflicts with the name of a public 
function.
    */