[incubator-tvm] branch master updated: upstream (#6436)

[jroesch]

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

jroesch pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git


The following commit(s) were added to refs/heads/master by this push:
     new 9b0e37a  upstream (#6436)
9b0e37a is described below

commit 9b0e37a42de0f4cb0efcaa907bfa54d61f1c3faf
Author: 雾雨魔理沙 <lol...@marisa.moe>
AuthorDate: Fri Sep 11 17:00:50 2020 -0700

    upstream (#6436)
---
 src/relay/transforms/fold_constant.cc | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/relay/transforms/fold_constant.cc 
b/src/relay/transforms/fold_constant.cc
index 0a1f173..a3f2f69 100644
--- a/src/relay/transforms/fold_constant.cc
+++ b/src/relay/transforms/fold_constant.cc
@@ -105,7 +105,23 @@ class ConstantFolder : public ExprMutator {
     }
   }
 
+  bool inside_primitive = false;
+  Expr VisitExpr_(const FunctionNode* op) final {
+    if (op->HasNonzeroAttr(attr::kPrimitive)) {
+      CHECK_EQ(inside_primitive, false);
+      inside_primitive = true;
+      auto ret = ExprMutator::VisitExpr_(op);
+      inside_primitive = false;
+      return ret;
+    } else {
+      return ExprMutator::VisitExpr_(op);
+    }
+  }
+
   Expr VisitExpr_(const CallNode* call) final {
+    if (inside_primitive) {
+      return GetRef<Expr>(call);
+    }
     static auto op_stateful = Op::GetAttrMap<TOpIsStateful>("TOpIsStateful");
 
     std::unordered_set<std::string> skip_list{"zeros_like", "ones_like", 
"full_like", "full"};