From d56abe33025a4089f4159062c4d0da92ac5b4f72 Mon Sep 17 00:00:00 2001
From: jiahongyu <jiahongyu@baidu.com>
Date: Fri, 3 Nov 2023 12:45:25 +0000
Subject: [PATCH] MapExpr supports concat op, WIP MapExpr2Ir

---
 paddle/cinn/adt/CMakeLists.txt               | 10 ++--
 paddle/cinn/adt/equation_function.cc         | 23 ++++++--
 paddle/cinn/adt/equation_function.h          | 20 +++++--
 paddle/cinn/adt/equation_solver.cc           | 16 ++++--
 paddle/cinn/adt/equation_value.h             | 14 +++++
 paddle/cinn/adt/equation_value_match_trait.h | 41 +++++++-------
 paddle/cinn/adt/m_ir.cc                      | 10 +++-
 paddle/cinn/adt/naive_op_equation_context.h  | 36 +++++++++---
 paddle/cinn/adt/op_equation_context.h        | 13 ++++-
 paddle/cinn/adt/print_equations.cc           | 13 ++++-
 paddle/cinn/adt/print_value.cc               | 10 +++-
 paddle/cinn/hlir/op/transform.cc             | 44 +++++++++++++++
 test/cinn/adt/test_naive_concat.py           | 59 ++++++++++++++++++++
 13 files changed, 252 insertions(+), 57 deletions(-)
 create mode 100644 test/cinn/adt/test_naive_concat.py

diff --git a/paddle/cinn/adt/CMakeLists.txt b/paddle/cinn/adt/CMakeLists.txt
index 32987eafd8f64..35ccd477151b9 100644
--- a/paddle/cinn/adt/CMakeLists.txt
+++ b/paddle/cinn/adt/CMakeLists.txt
@@ -4,6 +4,8 @@ gather_srcs(
   cinnapi_src
   SRCS
   anchor_sd_equation_context.cc
+  dim_expr_simplifier.cc
+  dim_expr.cc
   equation_function.cc
   equation_solver.cc
   equation_value.cc
@@ -17,6 +19,7 @@ gather_srcs(
   naive_bidirection_equation_generator.cc
   naive_op_equation_context.cc
   partition_op_stmts.cc
+  print_dim_expr.cc
   print_equations.cc
   print_map_expr.cc
   print_schedule_descriptor.cc
@@ -26,12 +29,9 @@ gather_srcs(
   schedule_descriptor.cc
   schedule_dim.cc
   schedule_mesh.cc
-  dim_expr.cc
-  dim_expr_simplifier.cc
-  symbolic_dim_infer_util.cc
   simplify_value.cc
-  write_broadcast_disabled_bidirection_equation_generator.cc
-  print_dim_expr.cc)
+  symbolic_dim_infer_util.cc
+  write_broadcast_disabled_bidirection_equation_generator.cc)
 
 cinn_cc_test(equation_value_match_trait_test SRCS
              equation_value_match_trait_test.cc DEPS gtest glog)
diff --git a/paddle/cinn/adt/equation_function.cc b/paddle/cinn/adt/equation_function.cc
index f69ae290e4c9c..2fdc9a8809ce2 100644
--- a/paddle/cinn/adt/equation_function.cc
+++ b/paddle/cinn/adt/equation_function.cc
@@ -33,9 +33,8 @@ CollectInputAndOutputVariables(const Function& function) {
             out_variables.emplace(Variable{out_index.value()});
             in_variables.emplace(Variable{in_index.value()});
           },
-          [&](const IndexDot<List<DimExpr>,
-                             tOut<Index>,
-                             tIn<List<Iterator>>>& dot) {
+          [&](const IndexDot<List<DimExpr>, tOut<Index>, tIn<List<Iterator>>>&
+                  dot) {
             const auto& [dims, out_index, in_iterators] = dot.tuple();
             out_variables.emplace(Variable{out_index.value()});
             for (const auto& iterator : *in_iterators.value()) {
@@ -49,9 +48,17 @@ CollectInputAndOutputVariables(const Function& function) {
             out_variables.emplace(Variable{out_iterator.value()});
             in_variables.emplace(Variable{in_iterator.value()});
           },
-          [&](const IndexUnDot<List<DimExpr>,
-                               tOut<List<Iterator>>,
-                               tIn<Index>>& undot) {
+          [&](const SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>&
+                  sub_function) {
+            {
+              const auto& [dim, out_iterator, in_iterator] =
+                  sub_function.tuple();
+              out_variables.emplace(Variable{out_iterator.value()});
+              in_variables.emplace(Variable{in_iterator.value()});
+            }
+          },
+          [&](const IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>&
+                  undot) {
             const auto& [dims, out_iterators, in_index] = undot.tuple();
             for (const auto& iterator : *out_iterators.value()) {
               out_variables.emplace(Variable{iterator});
@@ -113,6 +120,8 @@ std::string GetFunctionTypeName(const Function& function) {
                                               tIn<Iterator>>& broadcast) {
                return "GetBroadcastedIterator";
              },
+             [&](const SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>&
+                     sub_function) { return "SubFunction"; },
              [&](const IndexUnDot<List<DimExpr>,
                                   tOut<List<Iterator>>,
                                   tIn<Index>>& undot) { return "IndexUnDot"; },
@@ -142,6 +151,8 @@ const void* GetFunctionDataPtr(const Function& function) {
                                               tOut<Iterator>,
                                               tIn<Iterator>>& broadcast)
                  -> const void* { return &broadcast.tuple(); },
+             [&](const SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>& sub)
+                 -> const void* { return &sub.tuple(); },
              [&](const IndexUnDot<List<DimExpr>,
                                   tOut<List<Iterator>>,
                                   tIn<Index>>& undot) -> const void* {
diff --git a/paddle/cinn/adt/equation_function.h b/paddle/cinn/adt/equation_function.h
index c864ed0ede6f9..b12c41b09d997 100644
--- a/paddle/cinn/adt/equation_function.h
+++ b/paddle/cinn/adt/equation_function.h
@@ -106,19 +106,27 @@ struct GetBroadcastedIterator<DimExpr, tOut<Iterator>, tIn<Iterator>>
   using Tuple<DimExpr, tOut<Iterator>, tIn<Iterator>>::Tuple;
 };
 
+template <typename DimT, typename OutT, typename InT>
+struct SubFunction;
+
+template <>
+struct SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>
+    : public Tuple<DimExpr, tOut<Iterator>, tIn<Iterator>> {
+  using Tuple<DimExpr, tOut<Iterator>, tIn<Iterator>>::Tuple;
+};
+
 // clang-format off
 DEFINE_ADT_UNION(Equation,
                  Identity<tOut<Iterator>, tIn<Iterator>>,
                  Identity<tOut<Index>, tIn<Index>>,
                  GetBroadcastedIterator<DimExpr,
                                         tOut<Iterator>, tIn<Iterator>>,
-                 IndexDot<List<DimExpr>, tOut<Index>,
-                          tIn<List<Iterator>>>,
-                 IndexUnDot<List<DimExpr>,
-                            tOut<List<Iterator>>, tIn<Index>>,
+                 SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>,
+                 IndexDot<List<DimExpr>, tOut<Index>, tIn<List<Iterator>>>,
+                 IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>,
                  InMsg2OutMsg<tOut<FakeOpPlaceHolder>,
-                                    tOut<OpArgIndexes<std::optional<Index>>>,
-                                    tIn<OpArgIndexes<Index>>>,
+                                   tOut<OpArgIndexes<std::optional<Index>>>,
+                                   tIn<OpArgIndexes<Index>>>,
                  ConstantFunction<tOut<Iterator>, tIn<Index>>);
 // clang-format on
 
diff --git a/paddle/cinn/adt/equation_solver.cc b/paddle/cinn/adt/equation_solver.cc
index 35f3ae533281e..adad52e091445 100644
--- a/paddle/cinn/adt/equation_solver.cc
+++ b/paddle/cinn/adt/equation_solver.cc
@@ -85,8 +85,16 @@ std::unordered_map<Variable, Value> InferValuesImpl(
 }
 
 std::unordered_map<Variable, Value> InferValuesImpl(
-    const IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>&
-        undot,
+    const SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>& sub,
+    IndexExprInferContext* ctx) {
+  const auto& [dim, out_iterator, in_iterator] = sub.tuple();
+  SubValue<Value, DimExpr> sub_iterator{ctx->GetValue(in_iterator.value()),
+                                        dim};
+  return {{out_iterator.value(), sub_iterator}};
+}
+
+std::unordered_map<Variable, Value> InferValuesImpl(
+    const IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>& undot,
     IndexExprInferContext* ctx) {
   const auto& [dims, out_iters, in_index] = undot.tuple();
 
@@ -94,8 +102,8 @@ std::unordered_map<Variable, Value> InferValuesImpl(
   for (const auto& dim : *dims) {
     dim_constants->emplace_back(dim);
   }
-  IndexUnDotValue<Value, List<DimExpr>> index_undot{ctx->GetValue(in_index.value()),
-                                               dim_constants};
+  IndexUnDotValue<Value, List<DimExpr>> index_undot{
+      ctx->GetValue(in_index.value()), dim_constants};
 
   std::unordered_map<Variable, Value> ret{};
   for (std::size_t idx = 0; idx < out_iters.value()->size(); ++idx) {
diff --git a/paddle/cinn/adt/equation_value.h b/paddle/cinn/adt/equation_value.h
index a876ffef1bf6d..4cca35644ccc9 100644
--- a/paddle/cinn/adt/equation_value.h
+++ b/paddle/cinn/adt/equation_value.h
@@ -68,6 +68,13 @@ struct BroadcastedIterator final : public Tuple<ValueT, ConstantT> {
   const ValueT& GetArg0() const { return std::get<0>(this->tuple()); }
 };
 
+template <typename ValueT, typename ConstantT>
+struct SubValue final : public Tuple<ValueT, ConstantT> {
+  using Tuple<ValueT, ConstantT>::Tuple;
+
+  const ValueT& GetArg0() const { return std::get<0>(this->tuple()); }
+};
+
 DEFINE_ADT_UNION(Value,
                  Undefined,
                  Ok,
@@ -76,6 +83,7 @@ DEFINE_ADT_UNION(Value,
                  List<Value>,
                  IndexDotValue<Value, List<DimExpr>>,
                  IndexUnDotValue<Value, List<DimExpr>>,
+                 SubValue<Value, DimExpr>,
                  ListGetItem<Value, DimExpr>,
                  BroadcastedIterator<Value, DimExpr>,
                  PtrGetItem<Value>);
@@ -89,6 +97,8 @@ using ListGetItem_Value_DimExpr = ListGetItem<Value, DimExpr>;
 OVERLOAD_OPERATOR_EQ_NE(ListGetItem_Value_DimExpr, TupleEqual);
 using BroadcastedIterator_Value_DimExpr = BroadcastedIterator<Value, DimExpr>;
 OVERLOAD_OPERATOR_EQ_NE(BroadcastedIterator_Value_DimExpr, TupleEqual);
+using SubValue_Value_DimExpr = SubValue<Value, DimExpr>;
+OVERLOAD_OPERATOR_EQ_NE(SubValue_Value_DimExpr, TupleEqual);
 OVERLOAD_OPERATOR_EQ_NE(PtrGetItem<Value>, TupleEqual);
 
 inline std::size_t GetHashValue(const Value& value);
@@ -135,6 +145,10 @@ inline std::size_t GetHashValueImpl(
   const auto& [v, c] = value.tuple();
   return hash_combine(GetHashValue(v), GetHashValue(c));
 }
+inline std::size_t GetHashValueImpl(const SubValue<Value, DimExpr>& value) {
+  const auto& [v, c] = value.tuple();
+  return hash_combine(GetHashValue(v), GetHashValue(c));
+}
 inline std::size_t GetHashValueImpl(const PtrGetItem<Value>& value) {
   const auto& [pointer, c] = value.tuple();
   return hash_combine(pointer.value().unique_id(), GetHashValue(c));
diff --git a/paddle/cinn/adt/equation_value_match_trait.h b/paddle/cinn/adt/equation_value_match_trait.h
index 9b21c1146c8e3..496abc71badcb 100644
--- a/paddle/cinn/adt/equation_value_match_trait.h
+++ b/paddle/cinn/adt/equation_value_match_trait.h
@@ -47,37 +47,38 @@ struct MatchTrait<Value, List<T>> final {
   }
 };
 
-#define DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(name, type0, type1)          \
-  template <typename T0, typename T1>                                         \
-  struct MatchTrait<Value, name<T0, T1>> final {                              \
-    using base_type = name<type0, type1>;                                     \
-                                                                              \
-    static constexpr int is_template = true;                                  \
-                                                                              \
-    template <template <typename, typename> class Matcher>                              \
-    static bool MatchChildren(const base_type& value) {                       \
+#define DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(name, type0, type1) \
+  template <typename T0, typename T1>                                \
+  struct MatchTrait<Value, name<T0, T1>> final {                     \
+    using base_type = name<type0, type1>;                            \
+                                                                     \
+    static constexpr int is_template = true;                         \
+                                                                     \
+    template <template <typename, typename> class Matcher>           \
+    static bool MatchChildren(const base_type& value) {              \
       return Matcher<T0, type0>::Call(std::get<0>(value.tuple())) && \
              Matcher<T1, type1>::Call(std::get<1>(value.tuple()));   \
-    }                                                                         \
+    }                                                                \
   };
 
 DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(ListGetItem, Value, DimExpr);
 DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(BroadcastedIterator, Value, DimExpr);
+DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(SubValue, Value, DimExpr);
 DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(IndexDotValue, Value, List<DimExpr>);
 DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2(IndexUnDotValue, Value, List<DimExpr>);
 #undef DEFINE_MATCH_TRAIT_VALUE_UNION_ARGSIZE_2
 
-#define DEFINE_ADT_MATCH_TRAIT_EQUATION(name)                              \
-  template <typename T>                                                    \
-  struct MatchTrait<Value, name<T>> final {                                \
-    using base_type = name<Value>;                                         \
-                                                                           \
-    static constexpr int is_template = true;                               \
-                                                                           \
-    template <template <typename, typename> class Matcher>                           \
-    static bool MatchChildren(const base_type& value) {                    \
+#define DEFINE_ADT_MATCH_TRAIT_EQUATION(name)                     \
+  template <typename T>                                           \
+  struct MatchTrait<Value, name<T>> final {                       \
+    using base_type = name<Value>;                                \
+                                                                  \
+    static constexpr int is_template = true;                      \
+                                                                  \
+    template <template <typename, typename> class Matcher>        \
+    static bool MatchChildren(const base_type& value) {           \
       return Matcher<T, Value>::Call(std::get<0>(value.tuple())); \
-    }                                                                      \
+    }                                                             \
   };
 
 DEFINE_ADT_MATCH_TRAIT_EQUATION(PtrGetItem);
diff --git a/paddle/cinn/adt/m_ir.cc b/paddle/cinn/adt/m_ir.cc
index 003b6880c813a..8c32cef5394c1 100644
--- a/paddle/cinn/adt/m_ir.cc
+++ b/paddle/cinn/adt/m_ir.cc
@@ -82,9 +82,15 @@ void CollectTensorIndexIteratorsImpl(
 }
 
 void CollectTensorIndexIteratorsImpl(
-    const BroadcastedIterator<Value, DimExpr>& broadcasted_iterator,
+    const BroadcastedIterator<Value, DimExpr>& tensor_index_expr,
     std::unordered_set<Iterator>* ret) {
-  CollectTensorIndexIterators(broadcasted_iterator.GetArg0(), ret);
+  CollectTensorIndexIterators(tensor_index_expr.GetArg0(), ret);
+}
+
+void CollectTensorIndexIteratorsImpl(
+    const SubValue<Value, DimExpr>& tensor_index_expr,
+    std::unordered_set<Iterator>* ret) {
+  CollectTensorIndexIterators(tensor_index_expr.GetArg0(), ret);
 }
 
 void CollectTensorIndexIteratorsImpl(const PtrGetItem<Value>& tensor_index_expr,
diff --git a/paddle/cinn/adt/naive_op_equation_context.h b/paddle/cinn/adt/naive_op_equation_context.h
index 7f2b2bc8fe95e..a25ba47dbf2a2 100644
--- a/paddle/cinn/adt/naive_op_equation_context.h
+++ b/paddle/cinn/adt/naive_op_equation_context.h
@@ -102,6 +102,22 @@ class NaiveOpEquationContext final : public OpEquationContext {
     return input_tensor_iterator;
   }
 
+  Iterator Sub(const Iterator& out_tensor_iterator,
+               const DimExpr& dim) override {
+    Iterator input_tensor_iterator{UniqueId::New()};
+    using Function = SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>;
+    equations_->emplace_back(
+        Function{dim, input_tensor_iterator, out_tensor_iterator});
+    return input_tensor_iterator;
+  }
+
+  Iterator Identity(const Iterator& iterator) override {
+    Iterator new_iterator{UniqueId::New()};
+    using Function = cinn::adt::Identity<tOut<Iterator>, tIn<Iterator>>;
+    equations_->emplace_back(Function{new_iterator, iterator});
+    return new_iterator;
+  }
+
   Iterator MakeConstantIterator(std::size_t constant,
                                 Equations* equations) const {
     using ConstF = ConstantFunction<tOut<Iterator>, tIn<Index>>;
@@ -227,8 +243,8 @@ class NaiveOpEquationContext final : public OpEquationContext {
   }
 
   std::optional<DimExpr> GetSymbolicDimSize(bool is_out,
-                                                    std::size_t arg_idx,
-                                                    std::size_t axis) const {
+                                            std::size_t arg_idx,
+                                            std::size_t axis) const {
     const auto* Get = (is_out ? &GetSymbolicOutDim_ : &GetSymbolicInDim_);
     const auto& opt_dim = (*Get)(arg_idx, axis);
     return opt_dim;
@@ -247,7 +263,8 @@ class NaiveOpEquationContext final : public OpEquationContext {
     }
   }
 
-  void InitInputDimExpr(std::vector<DimTuple>* vec, const std::vector<std::uint64_t>& tensors_ranks) {
+  void InitInputDimExpr(std::vector<DimTuple>* vec,
+                        const std::vector<std::uint64_t>& tensors_ranks) {
     for (std::size_t i = 0; i < tensors_ranks.size(); ++i) {
       vec->push_back(DimTuple{});
       for (std::size_t j = 0; j < tensors_ranks.at(i); ++j) {
@@ -258,7 +275,8 @@ class NaiveOpEquationContext final : public OpEquationContext {
     }
   }
 
-  void InitOutputDimExpr(std::vector<DimTuple>* vec, const std::vector<std::uint64_t>& tensors_ranks) {
+  void InitOutputDimExpr(std::vector<DimTuple>* vec,
+                         const std::vector<std::uint64_t>& tensors_ranks) {
     for (std::size_t i = 0; i < tensors_ranks.size(); ++i) {
       vec->push_back(DimTuple{});
       for (std::size_t j = 0; j < tensors_ranks.at(i); ++j) {
@@ -303,8 +321,8 @@ class NaiveOpEquationContext final : public OpEquationContext {
 
   template <typename T>
   void Equal(const T& lhs, const T& rhs) {
-    equations_->emplace_back(Identity<tOut<T>, tIn<T>>(lhs, rhs));
-    equations_->emplace_back(Identity<tOut<T>, tIn<T>>(rhs, lhs));
+    equations_->emplace_back(cinn::adt::Identity<tOut<T>, tIn<T>>(lhs, rhs));
+    equations_->emplace_back(cinn::adt::Identity<tOut<T>, tIn<T>>(rhs, lhs));
   }
 
   static std::optional<std::size_t> FindPos(const List<Index>& vector,
@@ -317,13 +335,17 @@ class NaiveOpEquationContext final : public OpEquationContext {
     return std::nullopt;
   }
 
-  const utils::Attribute& GetAttribute(const std::string& name) const {
+  const utils::Attribute& GetAttribute(const std::string& name) const override {
     const auto& iter = attr_map_type_->find(name);
     CHECK(iter != attr_map_type_->end())
         << "Can't find Attribute with this name";
     return iter->second;
   }
 
+  bool HasAttr(const std::string& name) const override {
+    return attr_map_type_->find(name) != attr_map_type_->end();
+  }
+
   std::vector<std::uint64_t> in_tensors_ranks_;
   std::vector<std::uint64_t> out_tensors_ranks_;
   GetArgStaticDimT GetInDim_;
diff --git a/paddle/cinn/adt/op_equation_context.h b/paddle/cinn/adt/op_equation_context.h
index f09b1335937f6..cd099a4268407 100644
--- a/paddle/cinn/adt/op_equation_context.h
+++ b/paddle/cinn/adt/op_equation_context.h
@@ -19,8 +19,8 @@
 
 #include "glog/logging.h"
 #include "paddle/cinn/adt/adt.h"
-#include "paddle/cinn/adt/equation.h"
 #include "paddle/cinn/adt/dim_expr.h"
+#include "paddle/cinn/adt/equation.h"
 #include "paddle/cinn/hlir/framework/node.h"
 
 namespace cinn::adt::config {
@@ -43,8 +43,13 @@ class OpEquationContext {
 
   virtual void Equal(const IteratorTuple& lhs, const IteratorTuple& rhs) = 0;
 
-  virtual Iterator GetBroadcastedInputIterator(const Iterator& out_iterator,
-                                               const DimExpr& dim) = 0;
+  virtual Iterator GetBroadcastedInputIterator(
+      const Iterator& out_tensor_iterator, const DimExpr& dim) = 0;
+
+  virtual Iterator Sub(const Iterator& out_tensor_iterator,
+                       const DimExpr& dim) = 0;
+
+  virtual Iterator Identity(const Iterator& iterator) = 0;
 
   virtual const IteratorTuple& GetInIteratorTuple(
       std::size_t input_idx) const = 0;
@@ -65,6 +70,8 @@ class OpEquationContext {
     return absl::get<T>(GetAttribute(name));
   }
 
+  virtual bool HasAttr(const std::string& name) const = 0;
+
  protected:
   OpEquationContext() = default;
 
diff --git a/paddle/cinn/adt/print_equations.cc b/paddle/cinn/adt/print_equations.cc
index 859957b6d755d..9f298a8ab2c92 100644
--- a/paddle/cinn/adt/print_equations.cc
+++ b/paddle/cinn/adt/print_equations.cc
@@ -185,8 +185,17 @@ struct ToTxtStringStruct {
   }
 
   std::string operator()(
-      const IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>&
-          undot) const {
+      const SubFunction<DimExpr, tOut<Iterator>, tIn<Iterator>>& sub) const {
+    std::string ret;
+    const auto& [dim, out_iterator, in_iterator] = sub.tuple();
+    ret += ToTxtString(out_iterator.value()) + " = Sub(" +
+           ToTxtString(in_iterator.value()) + ", " + ToTxtString(dim) + ")";
+    return ret;
+  }
+
+  std::string operator()(
+      const IndexUnDot<List<DimExpr>, tOut<List<Iterator>>, tIn<Index>>& undot)
+      const {
     std::string ret;
     const auto& [dim_list, out_iterator_list_tag, in_index_tag] = undot.tuple();
     const List<Iterator>& out_iterator_list = out_iterator_list_tag.value();
diff --git a/paddle/cinn/adt/print_value.cc b/paddle/cinn/adt/print_value.cc
index ff174f5f630bb..ce9fb63d0b68e 100644
--- a/paddle/cinn/adt/print_value.cc
+++ b/paddle/cinn/adt/print_value.cc
@@ -83,14 +83,20 @@ struct ToTxtStringStruct {
     return ret;
   }
 
-  std::string operator()(
-      const BroadcastedIterator<Value, DimExpr>& broadcast) {
+  std::string operator()(const BroadcastedIterator<Value, DimExpr>& broadcast) {
     std::string ret;
     const auto& [value, constant] = broadcast.tuple();
     ret += "BI(" + ToTxtString(value) + ", " + ToTxtString(constant) + ")";
     return ret;
   }
 
+  std::string operator()(const SubValue<Value, DimExpr>& sub_value) {
+    std::string ret;
+    const auto& [value, constant] = sub_value.tuple();
+    ret += "(" + ToTxtString(value) + " - " + ToTxtString(constant) + ")";
+    return ret;
+  }
+
   std::string operator()(const PtrGetItem<Value>& ptr_get_item) {
     std::string ret;
     const auto& [ptr_tag, value] = ptr_get_item.tuple();
diff --git a/paddle/cinn/hlir/op/transform.cc b/paddle/cinn/hlir/op/transform.cc
index ed22f50c644b0..c66d354f3b627 100644
--- a/paddle/cinn/hlir/op/transform.cc
+++ b/paddle/cinn/hlir/op/transform.cc
@@ -16,6 +16,8 @@
 
 #include <algorithm>
 
+#include "paddle/cinn/adt/op_equation_context.h"
+#include "paddle/cinn/adt/symbolic_dim_infer_ctx.h"
 #include "paddle/cinn/common/cas.h"
 #include "paddle/cinn/hlir/framework/node.h"
 #include "paddle/cinn/hlir/framework/op.h"
@@ -508,6 +510,44 @@ std::vector<Type> InferDtypeForConcat(const std::vector<Type> &inputs_type,
   return res;
 }
 
+void GenerateEquationsForConcat(cinn::adt::config::OpEquationContext *ctx) {
+  CHECK_EQ(ctx->GetInTensorsRanks().size(), 2UL)
+      << "Strong constraint branch only supports two inputs temporarily";
+  CHECK_EQ(ctx->GetInTensorsRanks().at(0), ctx->GetInTensorsRanks().at(1));
+  int axis = 0;
+  if (ctx->HasAttr("axis")) {
+    axis = ctx->Attr<int>("axis");
+  }
+
+  if (axis < 0) axis += ctx->GetInTensorsRanks().at(0);
+  CHECK(axis >= 0 && axis < ctx->GetInTensorsRanks().at(0))
+      << "In Concat op, the attribute `axis` should be >= 0 and < input "
+         "shape's size, please check!";
+
+  for (int i = 0; i < ctx->GetInTensorsRanks().at(0); ++i) {
+    if (i == axis) {
+      ctx->Equal(ctx->GetInIteratorTuple(0)->at(i),
+                 ctx->Identity(ctx->GetOutIteratorTuple(0)->at(i)));
+      ctx->Equal(ctx->GetInIteratorTuple(1)->at(i),
+                 ctx->Sub(ctx->GetOutIteratorTuple(0)->at(i),
+                          ctx->GetOutDimTuple(0)->at(i)));
+    } else {
+      // TODO(Hongyu Jia) : Add Check interface:
+      // ctx->CheckEqual(ctx->GetInDimTuple(0).at(i),
+      // ctx->GetInDimTuple(1).at(i))
+      ctx->Equal(ctx->GetInIteratorTuple(0)->at(i),
+                 ctx->GetOutIteratorTuple(0)->at(i));
+      ctx->Equal(ctx->GetInIteratorTuple(1)->at(i),
+                 ctx->GetOutIteratorTuple(0)->at(i));
+    }
+  }
+}
+
+void InferSymbolicDimForConcat(cinn::adt::config::SymbolicDimInferCtx *ctx) {
+  // TODO(Hongyu Jia): Support dynamic shape of concat, support Add in
+  // SymbolicDimInferCtx
+}
+
 std::vector<std::vector<std::string>> InferLayoutForConcat(
     const std::vector<framework::shape_t> &input_shapes,
     const std::vector<std::string> &input_layouts,
@@ -1844,6 +1884,10 @@ CINN_REGISTER_HELPER(transform_ops) {
                 MakeOpFunction(cinn::hlir::op::InferShapeForConcat))
       .set_attr("inferdtype",
                 MakeOpFunction(cinn::hlir::op::InferDtypeForConcat))
+      .set_attr("generate_equations",
+                MakeOpFunction(cinn::hlir::op::GenerateEquationsForConcat))
+      .set_attr("infer_symbolic_dim",
+                MakeOpFunction(cinn::hlir::op::InferSymbolicDimForConcat))
 #ifndef CINN_WITH_CUDA
       .set_attr("inferlayout",
                 MakeOpFunction(cinn::hlir::op::InferLayoutForConcat))
diff --git a/test/cinn/adt/test_naive_concat.py b/test/cinn/adt/test_naive_concat.py
new file mode 100644
index 0000000000000..306b0d71d6a29
--- /dev/null
+++ b/test/cinn/adt/test_naive_concat.py
@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+
+# Copyright (c) 2021 CINN Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from cinn.common import DefaultNVGPUTarget, Float
+from cinn.frontend import NetBuilder
+
+
+class TestMapExprNaiveAdd(unittest.TestCase):
+    def setUp(self):
+        self.inputs = {
+            "x": np.random.uniform(-1.0, 1.0, [64, 16]).astype("float32"),
+            "y": np.random.uniform(-1.0, 1.0, [64, 16]).astype("float32"),
+        }
+
+    def test_naive_add(self):
+        builder = NetBuilder("TestMapExprNaiveAdd")
+        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
+        y = builder.create_input(Float(32), self.inputs["y"].shape, "y")
+
+        x = builder.relu(x)
+        y = builder.relu(y)
+        out = builder.concat([x, y], axis=1)
+        prog = builder.build()
+        target = DefaultNVGPUTarget()
+
+        result = prog.build_and_get_output(
+            target,
+            [x, y],
+            [self.inputs["x"], self.inputs["y"]],
+            [out],
+            passes=[],
+            scope=None,
+        )
+
+        # np.testing.assert_allclose(
+        #     result[0].numpy(target),
+        #     np.concatenate((self.inputs["x"], self.inputs["y"]), axis=1),
+        #     err_msg="TestMapExprNaiveAdd failed!",
+        # )
+
+
+if __name__ == "__main__":
+    unittest.main()