(tvm) branch main updated: [Relax] Fix fuseOps via pattern (#17160)

[tqchen]

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

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


The following commit(s) were added to refs/heads/main by this push:
     new 70c53082e6 [Relax] Fix fuseOps via pattern  (#17160)
70c53082e6 is described below

commit 70c53082e6715516aefefcdca6262e195f36a0de
Author: Siyuan Feng <hzfen...@sjtu.edu.cn>
AuthorDate: Wed Jul 17 02:34:19 2024 +0800

    [Relax] Fix fuseOps via pattern  (#17160)
    
    fix fuseops via pattern
---
 src/relax/transform/fuse_ops.cc | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index 45d70fc3e2..2be7ad41f3 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -1222,7 +1222,12 @@ class CompositeFunctionAnnotator : public ExprMutator {
   IRModule Run() {
     auto mod = builder_->GetContextIRModule();
     for (const auto& gv : mod->GetGlobalVars()) {
-      const auto& base_func = mod->Lookup(gv);
+      auto it = mod->functions.find(gv);
+      // Note that the fusion pass may have already removed the function.
+      if (it == mod->functions.end()) {
+        continue;
+      }
+      const auto& base_func = (*it).second;
       if (const auto* func = base_func.as<FunctionNode>()) {
         if (func->GetAttr<String>(attr::kComposite).defined() ||
             func->GetAttr<String>(attr::kCodegen).defined()) {
@@ -1399,7 +1404,7 @@ Pass FuseOps(int fuse_opt_level) {
       };
   return CreateModulePass(/*pass_function=*/pass_func,  //
                           /*opt_level=*/0,              //
-                          /*pass_name=*/"FuseOps",      //
+                          /*name=*/"FuseOps",           //
                           /*required=*/{});
 }
 
@@ -1412,9 +1417,9 @@ Pass FuseOpsByPattern(const tvm::Array<FusionPattern>& 
patterns, bool bind_const
         return relax::FuseOpsByPattern(patterns, m, bind_constants, 
annotate_codegen,
                                        entry_function_names);
       };
-  return CreateModulePass(/*pass_function=*/pass_func,       //
-                          /*opt_level=*/0,                   //
-                          /*pass_name=*/"FuseOpsByPattern",  //
+  return CreateModulePass(/*pass_function=*/pass_func,  //
+                          /*opt_level=*/0,              //
+                          /*name=*/"FuseOpsByPattern",  //
                           /*required=*/{});
 }