[tvm] branch unity updated: [Unity][Fix] Fix bug in MergeCompositeFunctions (#14117)

[ruihangl]

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

ruihangl 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 268c06528c [Unity][Fix] Fix bug in MergeCompositeFunctions (#14117)
268c06528c is described below

commit 268c06528cfc1a8ec1d99431a520a22e075c89fc
Author: Yixin Dong <ubosp...@gmail.com>
AuthorDate: Fri Feb 24 23:27:25 2023 +0800

    [Unity][Fix] Fix bug in MergeCompositeFunctions (#14117)
    
    Currently `MergeCompositeFunctions` will modify the map while iterating 
over it, and that makes 
tests/python/relax/test_transform_merge_composite_functions.py does not pass. 
This PR fixes this bug.
---
 src/relax/transform/merge_composite_functions.cc | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/relax/transform/merge_composite_functions.cc 
b/src/relax/transform/merge_composite_functions.cc
index db73392b02..609dd173f2 100644
--- a/src/relax/transform/merge_composite_functions.cc
+++ b/src/relax/transform/merge_composite_functions.cc
@@ -324,13 +324,17 @@ IRModule MergeCompositeFunctions(IRModule mod) {
   auto new_mod = MakeGroupedFunctions(mod, group_map);
 
   CompositeInliner inliner(mod);
+  std::vector<std::pair<GlobalVar, BaseFunc>> to_update;
   for (const auto& [gvar, func] : new_mod->functions) {
     if (func->GetAttr<String>(attr::kCodegen)) {
       auto new_func = inliner.Run(Downcast<Function>(func));
       new_func = WithAttr(new_func, tvm::attr::kGlobalSymbol, gvar->name_hint);
-      new_mod->Update(gvar, new_func);
+      to_update.emplace_back(gvar, new_func);
     }
   }
+  for (const auto& [gvar, func] : to_update) {
+    new_mod->Update(gvar, func);
+  }
   // TODO(@tvm-team): Implicit pass dependency. Revisit when we have a better 
way to handle this.
   return RemoveUnusedFunctions(new_mod, {"main"});
 }