From f6d80371a90931867245f60dd23f63412ab59363 Mon Sep 17 00:00:00 2001
From: Zhi Chen <chzhi@amazon.com>
Date: Mon, 27 Apr 2020 23:06:02 +0000
Subject: [PATCH 1/2] add constantnode to atomic

---
 src/relay/transforms/pass_util.h              |  6 +++++-
 tests/python/relay/test_pass_fold_constant.py | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/relay/transforms/pass_util.h b/src/relay/transforms/pass_util.h
index 56b064573841..08e3dee85f86 100644
--- a/src/relay/transforms/pass_util.h
+++ b/src/relay/transforms/pass_util.h
@@ -97,7 +97,11 @@ inline Expr TransformF(const std::function<Expr(const Expr&)>& func, const Expr&
  *   if so, the compute cost of the expression is bounded so it can be copy without graph mode.
  */
 inline bool IsAtomic(const Expr& e) {
-  return e.as<VarNode>() || e.as<OpNode>() || e.as<ConstructorNode>() || e.as<GlobalVarNode>();
+  return e.as<VarNode>() ||
+         e.as<OpNode>() ||
+         e.as<ConstructorNode>() ||
+         e.as<GlobalVarNode>() ||
+         e.as<ConstantNode>();
 }
 
 template<typename ConditionObjectPtr>
diff --git a/tests/python/relay/test_pass_fold_constant.py b/tests/python/relay/test_pass_fold_constant.py
index b212b26c99a7..a981667219cd 100644
--- a/tests/python/relay/test_pass_fold_constant.py
+++ b/tests/python/relay/test_pass_fold_constant.py
@@ -32,6 +32,25 @@ def run_opt_pass(expr, opt_pass):
     return entry if isinstance(expr, relay.Function) else entry.body
 
 
+def test_concatenate_const():
+    def before():
+        data = tvm.nd.array(np.array([1.0, 2.0, 3.0]))
+        const = relay.const(data)
+        concat = relay.op.concatenate([const, const], axis=0)
+        func = relay.Function([], concat)
+        return func
+
+    def expected():
+        data = tvm.nd.array(np.array([1.0, 2.0, 3.0, 1.0, 2.0, 3.0]))
+        const = relay.const(data)
+        func = relay.Function([], const)
+        return func
+
+    zz = run_opt_pass(before(), transform.FoldConstant())
+    zexpected = run_opt_pass(expected(), transform.InferType())
+    assert tvm.ir.structural_equal(zz, zexpected)
+
+
 def test_fold_const():
     c_data = np.array([1, 2, 3]).astype("float32")
     t = relay.TensorType([1, 2, 3], "float32")

From 28d41ea0b85d7b76e1722baa75d71b42d3847d42 Mon Sep 17 00:00:00 2001
From: Zhi Chen <chzhi@amazon.com>
Date: Wed, 29 Apr 2020 23:21:36 +0000
Subject: [PATCH 2/2] Add ToANormalForm to FoldConstant

---
 src/relay/transforms/fold_constant.cc | 1 +
 src/relay/transforms/pass_util.h      | 6 +-----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/relay/transforms/fold_constant.cc b/src/relay/transforms/fold_constant.cc
index a52f42054c3e..fab184cfb5a6 100644
--- a/src/relay/transforms/fold_constant.cc
+++ b/src/relay/transforms/fold_constant.cc
@@ -203,6 +203,7 @@ class ConstantFolder : public ExprMutator {
   // Constant evaluate a expression.
   Expr ConstEvaluate(Expr expr) {
     std::vector<transform::Pass> passes = {transform::FuseOps(0),
+                                           transform::ToANormalForm(),
                                            transform::InferType()};
     Function func;
     if (expr.as<FunctionNode>()) {
diff --git a/src/relay/transforms/pass_util.h b/src/relay/transforms/pass_util.h
index 08e3dee85f86..56b064573841 100644
--- a/src/relay/transforms/pass_util.h
+++ b/src/relay/transforms/pass_util.h
@@ -97,11 +97,7 @@ inline Expr TransformF(const std::function<Expr(const Expr&)>& func, const Expr&
  *   if so, the compute cost of the expression is bounded so it can be copy without graph mode.
  */
 inline bool IsAtomic(const Expr& e) {
-  return e.as<VarNode>() ||
-         e.as<OpNode>() ||
-         e.as<ConstructorNode>() ||
-         e.as<GlobalVarNode>() ||
-         e.as<ConstantNode>();
+  return e.as<VarNode>() || e.as<OpNode>() || e.as<ConstructorNode>() || e.as<GlobalVarNode>();
 }
 
 template<typename ConditionObjectPtr>