(tvm) branch test_all_cases_on_unity updated: upd irmodule

[syfeng]

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

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


The following commit(s) were added to refs/heads/test_all_cases_on_unity by 
this push:
     new 7ba76b4a82 upd irmodule
7ba76b4a82 is described below

commit 7ba76b4a82fc239783317408f7c375e2f3301397
Author: Siyuan Feng <hzfen...@sjtu.edu.cn>
AuthorDate: Tue Dec 5 11:44:20 2023 +0800

    upd irmodule
---
 ci/jenkins/generated/cpu_jenkinsfile.groovy |  2 +-
 src/ir/module.cc                            | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/ci/jenkins/generated/cpu_jenkinsfile.groovy 
b/ci/jenkins/generated/cpu_jenkinsfile.groovy
index 3a5f888427..c2e9c4347d 100644
--- a/ci/jenkins/generated/cpu_jenkinsfile.groovy
+++ b/ci/jenkins/generated/cpu_jenkinsfile.groovy
@@ -573,7 +573,7 @@ def build(node_type) {
         // sh "${docker_run} ${ci_cpu} ./tests/scripts/task_golang.sh"
         // TODO(@jroesch): need to resolve CI issue will turn back on in 
follow up patch
 
-        // sh (script: "${docker_run} ${ci_cpu} ./tests/scripts/task_rust.sh", 
label: 'Rust build and test')
+        sh (script: "${docker_run} ${ci_cpu} ./tests/scripts/task_rust.sh", 
label: 'Rust build and test')
             })
           }
         }
diff --git a/src/ir/module.cc b/src/ir/module.cc
index d9d1586989..c016612c15 100644
--- a/src/ir/module.cc
+++ b/src/ir/module.cc
@@ -397,8 +397,20 @@ TVM_REGISTER_NODE_TYPE(IRModuleNode);
 
 TVM_REGISTER_GLOBAL("ir.IRModule")
     .set_body_typed([](tvm::Map<GlobalVar, BaseFunc> funcs, 
tvm::Map<GlobalTypeVar, TypeData> types,
-                       tvm::DictAttrs attrs, Map<String, Array<GlobalInfo>> 
global_infos) {
-      return IRModule(funcs, types, {}, {}, attrs, global_infos);
+                       tvm::ObjectRef attrs, Map<String, Array<GlobalInfo>> 
global_infos) {
+      auto dict_attrs = [&attrs]() {
+        if (!attrs.defined()) {
+          return DictAttrs();
+        } else if (auto* as_dict_attrs = attrs.as<tvm::DictAttrsNode>()) {
+          return GetRef<tvm::DictAttrs>(as_dict_attrs);
+        } else if (attrs.as<tvm::MapNode>()) {
+          return tvm::DictAttrs(Downcast<Map<String, ObjectRef>>(attrs));
+        } else {
+          LOG(FATAL) << "Expected attrs argument to be either DictAttrs or 
Map<String,ObjectRef>";
+        }
+      }();
+
+      return IRModule(funcs, types, {}, {}, dict_attrs, global_infos);
     });
 
 TVM_REGISTER_GLOBAL("ir.Module_Add")