From 0339776519b95ac5679bd889ffb9ac0179d0c341 Mon Sep 17 00:00:00 2001
From: HongyuJia <jiahongyu@baidu.com>
Date: Tue, 21 Nov 2023 11:13:09 +0800
Subject: [PATCH] [CINN] Strong constraint branch adapt to pir (#58993)

* Strong Constraint Branch

* NoInlineTranslator (#84)

* Adapt adt to pir

* Move FLAGS_cinn_enable_map_expr_schedule location

* Apply new group schedule

* Remove useless log

* Remove adt unittest

* Solve merge conflicts

* Fix typo

* Fix merge conflicts

* Add unit test

* Fix cmake

* Add test_cinn_sub_graph_map_expr

* Save current workspace

* Remove unittest to cinn directory

* Refactor unittest cmake

* Refactor unittest cmake

* Add unit test at Cmake

* Restore test_cinn_sub_graph.py

* Fix unittest

* Refine codes according to comment

* Refine codes according to comment
---
 paddle/cinn/adt/CMakeLists.txt                |  62 ++++-----
 paddle/cinn/adt/adapter_tensor.cc             |  44 +++++++
 paddle/cinn/adt/adapter_tensor.h              |  45 ++-----
 paddle/cinn/adt/equation_value.h              |   5 -
 paddle/cinn/adt/generate_map_expr.cc          | 118 +++++++++---------
 paddle/cinn/adt/generate_map_expr.h           |  22 ++--
 paddle/cinn/adt/inline_translator.h           |  37 +-----
 paddle/cinn/adt/inline_translator_trait.h     |  58 +++++++++
 paddle/cinn/adt/kgroup.cc                     |   4 +-
 paddle/cinn/adt/kgroup.h                      |  12 +-
 paddle/cinn/adt/m_expr.h                      |  23 ++--
 paddle/cinn/adt/map_expr_ctx.h                |   6 +-
 paddle/cinn/adt/naive_op_equation_context.cc  |  45 ++++---
 paddle/cinn/adt/naive_op_equation_context.h   |   8 +-
 paddle/cinn/adt/no_inline_translator.h        |  83 ++++++++++++
 paddle/cinn/adt/print_utils/CMakeLists.txt    |  27 ++--
 .../cinn/adt/print_utils/print_equations.cc   |  14 ++-
 paddle/cinn/adt/print_utils/print_map_expr.cc |  18 +--
 paddle/cinn/adt/schedule_descriptor.cc        |   4 +-
 .../transforms/cinn_group_lowering_pass.cc    |   7 ++
 paddle/cinn/hlir/framework/graph.h            |  16 ---
 .../cinn/hlir/framework/op_lowering_impl.cc   |  94 --------------
 paddle/cinn/hlir/framework/op_lowering_impl.h |  38 ------
 paddle/cinn/hlir/framework/pir/group.h        |  31 ++++-
 .../hlir/framework/pir/op_lowering_impl.cc    | 105 +++++++++++++++-
 .../hlir/framework/pir/op_lowering_impl.h     |  33 +++++
 paddle/cinn/hlir/pe/CMakeLists.txt            |   7 +-
 paddle/cinn/hlir/pe/map_expr_to_ir.cc         | 108 +++++++++++-----
 paddle/cinn/hlir/pe/map_expr_to_ir.h          |   1 +
 .../st_shape_group_scheduler.cc               |   7 ++
 .../group_schedule/st_shape_group_scheduler.h |   2 +
 paddle/cinn/pybind/frontend.cc                |   2 -
 paddle/cinn/runtime/flags.cc                  |  11 +-
 .../framework/paddle2cinn/cinn_compiler.cc    |   2 -
 test/CMakeLists.txt                           |   2 +-
 test/cinn/CMakeLists.txt                      |  25 +---
 test/cinn/adt/CMakeLists.txt                  |  21 ++++
 test/cinn/adt/test_add_inline.py              |  60 ---------
 test/cinn/adt/test_broadcast_expr.py          |  59 ---------
 test/cinn/adt/test_cinn_sub_graph_map_expr.py |  76 +++++++++++
 test/cinn/adt/test_naive_add.py               |  57 ---------
 test/cinn/adt/test_naive_reduce.py            |  50 --------
 test/cinn/adt/test_reduce_fusion.py           |  62 ---------
 test/cinn/adt/test_reduce_schedule_mesh.py    |  59 ---------
 44 files changed, 752 insertions(+), 818 deletions(-)
 create mode 100644 paddle/cinn/adt/adapter_tensor.cc
 create mode 100644 paddle/cinn/adt/inline_translator_trait.h
 create mode 100644 paddle/cinn/adt/no_inline_translator.h
 create mode 100644 test/cinn/adt/CMakeLists.txt
 delete mode 100755 test/cinn/adt/test_add_inline.py
 delete mode 100755 test/cinn/adt/test_broadcast_expr.py
 create mode 100644 test/cinn/adt/test_cinn_sub_graph_map_expr.py
 delete mode 100755 test/cinn/adt/test_naive_add.py
 delete mode 100644 test/cinn/adt/test_naive_reduce.py
 delete mode 100644 test/cinn/adt/test_reduce_fusion.py
 delete mode 100644 test/cinn/adt/test_reduce_schedule_mesh.py

diff --git a/paddle/cinn/adt/CMakeLists.txt b/paddle/cinn/adt/CMakeLists.txt
index 8dbb3bc8769839..e74c21bb0e7949 100644
--- a/paddle/cinn/adt/CMakeLists.txt
+++ b/paddle/cinn/adt/CMakeLists.txt
@@ -1,35 +1,39 @@
-add_subdirectory(print_utils)
+if(NOT CINN_ONLY)
+  add_subdirectory(print_utils)
 
-core_gather_headers()
+  core_gather_headers()
 
-gather_srcs(
-  cinnapi_src
-  SRCS
-  anchor_sd_equation_context.cc
-  equation_function.cc
-  equation_solver.cc
-  equation_value.cc
-  generate_map_expr.cc
-  get_sub_reshape_dim_ranges.cc
-  igroup.cc
-  index_expr_infer_context.cc
-  kgroup.cc
-  m_ir.cc
-  naive_bidirection_equation_generator.cc
-  naive_op_equation_context.cc
-  partition_op_stmts.cc
-  schedule_descriptor.cc
-  schedule_dim.cc
-  schedule_mesh.cc
-  simplify_value.cc
-  write_broadcast_disabled_bidirection_equation_generator.cc)
+  gather_srcs(
+    cinnapi_src
+    SRCS
+    adapter_tensor.cc
+    anchor_sd_equation_context.cc
+    equation_function.cc
+    equation_solver.cc
+    equation_value.cc
+    generate_map_expr.cc
+    get_sub_reshape_dim_ranges.cc
+    igroup.cc
+    index_expr_infer_context.cc
+    kgroup.cc
+    m_ir.cc
+    naive_bidirection_equation_generator.cc
+    naive_op_equation_context.cc
+    partition_op_stmts.cc
+    schedule_descriptor.cc
+    schedule_dim.cc
+    schedule_mesh.cc
+    simplify_value.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)
+  cinn_cc_test(equation_value_match_trait_test SRCS
+               equation_value_match_trait_test.cc DEPS gtest glog)
 
-cinn_cc_test(tree_test SRCS tree_test.cc DEPS gtest glog)
+  cinn_cc_test(tree_test SRCS tree_test.cc DEPS gtest glog)
 
-cinn_cc_test(inline_translator_test SRCS inline_translator_test.cc DEPS
-             cinncore)
+  cinn_cc_test(inline_translator_test SRCS inline_translator_test.cc DEPS
+               cinncore)
 
-message(STATUS "ADT srcs: ${cinnapi_src}")
+  message(STATUS "ADT srcs: ${cinnapi_src}")
+
+endif()
diff --git a/paddle/cinn/adt/adapter_tensor.cc b/paddle/cinn/adt/adapter_tensor.cc
new file mode 100644
index 00000000000000..464c45780dbecd
--- /dev/null
+++ b/paddle/cinn/adt/adapter_tensor.cc
@@ -0,0 +1,44 @@
+// Copyright (c) 2023 PaddlePaddle 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.
+
+#include "paddle/cinn/adt/adapter_tensor.h"
+#include "glog/logging.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
+
+namespace cinn::adt::adapter {
+
+std::size_t Tensor::GetRank() const {
+  return cinn::hlir::framework::pir::CompatibleInfo::ValueShape(node_data)
+      .size();
+}
+
+std::vector<int32_t> Tensor::GetShape() const {
+  std::vector<int32_t> ret{};
+  for (int dim_size :
+       cinn::hlir::framework::pir::CompatibleInfo::ValueShape(node_data)) {
+    ret.emplace_back(dim_size);
+  }
+  return ret;
+}
+
+std::size_t Tensor::GetNumel() const {
+  std::size_t ret = 1;
+  for (int dim_size :
+       cinn::hlir::framework::pir::CompatibleInfo::ValueShape(node_data)) {
+    ret = ret * dim_size;
+  }
+  return ret;
+}
+
+}  // namespace cinn::adt::adapter
diff --git a/paddle/cinn/adt/adapter_tensor.h b/paddle/cinn/adt/adapter_tensor.h
index 2a6cc941afb89e..dbd2c2dcecfdbb 100644
--- a/paddle/cinn/adt/adapter_tensor.h
+++ b/paddle/cinn/adt/adapter_tensor.h
@@ -13,59 +13,28 @@
 // limitations under the License.
 
 #pragma once
-#include "glog/logging.h"
 
 #include "paddle/cinn/adt/adt.h"
-#include "paddle/cinn/hlir/framework/graph.h"
-#include "paddle/cinn/hlir/framework/node.h"
+#include "paddle/pir/core/value.h"
 
 namespace cinn::adt::adapter {
 
 struct Tensor final {
-  const hlir::framework::NodeData* node_data;
-  const hlir::framework::Graph* graph;
+  ::pir::Value node_data;
 
   bool operator==(const Tensor& other) const {
-    return this->node_data == other.node_data && this->graph == other.graph;
+    return this->node_data == other.node_data;
   }
 
-  std::size_t GetRank() const {
-    const auto& shape_dict =
-        graph->GetAttrs<absl::flat_hash_map<std::string, utils::ShapeType>>(
-            "infershape");
-    CHECK(shape_dict.count(node_data->id()))
-        << "Can't find " << node_data->id() << " 's shape!";
-    return shape_dict.at(node_data->id()).size();
-  }
+  std::size_t GetRank() const;
 
-  const std::vector<int32_t>& GetShape() const {
-    const auto& shape_dict =
-        graph->GetAttrs<absl::flat_hash_map<std::string, utils::ShapeType>>(
-            "infershape");
-    CHECK(shape_dict.count(node_data->id()))
-        << "Can't find " << node_data->id() << " 's shape!";
-    return shape_dict.at(node_data->id());
-  }
+  std::vector<int32_t> GetShape() const;
 
-  std::size_t GetNumel() const {
-    const auto& shape_dict =
-        graph->GetAttrs<absl::flat_hash_map<std::string, utils::ShapeType>>(
-            "infershape");
-    CHECK(shape_dict.count(node_data->id()))
-        << "Can't find " << node_data->id() << " 's shape!";
-    std::vector<int32_t> shape = shape_dict.at(node_data->id());
-    std::size_t ret = 1;
-    for (int32_t dim_size : shape) {
-      ret = ret * dim_size;
-    }
-    return ret;
-  }
+  std::size_t GetNumel() const;
 };
 
 inline std::size_t GetHashValueImpl(const Tensor& tensor) {
-  return hash_combine(
-      std::hash<const hlir::framework::NodeData*>()(tensor.node_data),
-      std::hash<const hlir::framework::Graph*>()(tensor.graph));
+  return std::hash<::pir::Value>()(tensor.node_data);
 }
 
 }  // namespace cinn::adt::adapter
diff --git a/paddle/cinn/adt/equation_value.h b/paddle/cinn/adt/equation_value.h
index 7aa6c2b7c3155b..6c1fef21a93dde 100644
--- a/paddle/cinn/adt/equation_value.h
+++ b/paddle/cinn/adt/equation_value.h
@@ -20,11 +20,6 @@
 #include "paddle/cinn/adt/equation.h"
 #include "paddle/cinn/adt/match.h"
 
-namespace cinn::hlir::framework {
-class Node;
-class NodeData;
-}  // namespace cinn::hlir::framework
-
 namespace cinn::adt {
 
 DEFINE_ADT_TAG(tPointer);
diff --git a/paddle/cinn/adt/generate_map_expr.cc b/paddle/cinn/adt/generate_map_expr.cc
index b435acbcbcfb95..4180c9174a45fb 100644
--- a/paddle/cinn/adt/generate_map_expr.cc
+++ b/paddle/cinn/adt/generate_map_expr.cc
@@ -26,7 +26,11 @@
 #include "paddle/cinn/adt/print.h"
 #include "paddle/cinn/adt/schedule_descriptor.h"
 #include "paddle/cinn/adt/tree.h"
+#include "paddle/cinn/hlir/framework/pir/group.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/cinn/runtime/flags.h"
+#include "paddle/pir/core/operation.h"
+#include "paddle/pir/core/value.h"
 
 #include "glog/logging.h"
 
@@ -84,79 +88,65 @@ using LoopDescriptor4IterVarT = std::function<LoopDescriptor(const Iterator&)>;
 using AnchorTensor = Variable;
 using FakeOpPlaceHolders = List<FakeOpPlaceHolder>;
 
-Op MakeOp(const hlir::framework::Node* op) { return {op}; }
+Op MakeOp(const ::pir::Operation* op) { return {op}; }
 
 template <typename DoEachT>
-void VisitEachInputTensor(const hlir::framework::Node* op,
-                          const DoEachT& DoEach) {
-  for (const auto& graph_edge : op->inlinks_in_order()) {
-    DoEach(graph_edge->source()->safe_as<hlir::framework::NodeData>());
+void VisitEachInputTensor(const ::pir::Operation* op, const DoEachT& DoEach) {
+  for (std::size_t i = 0; i < op->num_operands(); ++i) {
+    DoEach(op->operand_source(i));
   }
 }
 
-List<Arg> MakeOpStmtInputList(const hlir::framework::Node* op,
-                              const hlir::framework::Graph* graph) {
+List<Arg> MakeOpStmtInputList(const ::pir::Operation* op) {
   List<Arg> ret{};
 
-  VisitEachInputTensor(op, [&](const auto* tensor) {
-    ret->emplace_back(adapter::Tensor{tensor, graph});
+  VisitEachInputTensor(op, [&](const ::pir::Value& tensor) {
+    ret->emplace_back(adapter::Tensor{tensor});
   });
 
   return ret;
 }
 
 template <typename DoEachT>
-void VisitEachOutputTensor(const hlir::framework::Node* op,
-                           const DoEachT& DoEach) {
-  for (const auto& graph_edge : op->outlinks_in_order()) {
-    DoEach(graph_edge->sink()->safe_as<hlir::framework::NodeData>());
+void VisitEachOutputTensor(const ::pir::Operation* op, const DoEachT& DoEach) {
+  for (std::size_t i = 0; i < op->num_results(); ++i) {
+    DoEach(const_cast<::pir::Operation*>(op)->result(i));
   }
 }
 
-List<Arg> MakeOpStmtOutputList(const hlir::framework::Node* op,
-                               const hlir::framework::Graph* graph) {
+List<Arg> MakeOpStmtOutputList(const ::pir::Operation* op) {
   List<Arg> ret{};
 
-  VisitEachOutputTensor(op, [&](const auto* tensor) {
-    ret->emplace_back(adapter::Tensor{tensor, graph});
+  VisitEachOutputTensor(op, [&](const ::pir::Value& tensor) {
+    ret->emplace_back(adapter::Tensor{tensor});
   });
 
   return ret;
 }
 
 template <typename DoEachT>
-void VisitEachOpStmt(
-    const std::shared_ptr<hlir::framework::Graph::Group>& group,
-    const DoEachT& DoEach) {
-  // Note
-  for (const auto* op : group->CollectNodes()) {
-    DoEach(OpStmt{MakeOp(op),
-                  MakeOpStmtInputList(op, group->graph_),
-                  MakeOpStmtOutputList(op, group->graph_)});
+void VisitEachOpStmt(const std::shared_ptr<hlir::framework::pir::Group>& group,
+                     const DoEachT& DoEach) {
+  for (const auto* op : group->CollectOps()) {
+    DoEach(
+        OpStmt{MakeOp(op), MakeOpStmtInputList(op), MakeOpStmtOutputList(op)});
   }
 }
 
-hlir::framework::OpPatternKind GetOpPatternKind(
-    const hlir::framework::Node* node) {
-  static const hlir::framework::OpValueType<hlir::framework::OpPatternKind>&
-      op_pattern_dict =
-          hlir::framework::Operator::GetAttrs<hlir::framework::OpPatternKind>(
-              "OpPattern");
-  auto kind = op_pattern_dict[node->op()];
-  return kind;
+hlir::framework::OpPatternKind GetOpPatternKind(const ::pir::Operation* node) {
+  return hlir::framework::pir::CompatibleInfo::OpKind(*node);
 }
 
 bool CollectRewritedReductionOpStmts(const OpStmt& op_stmt, List<OpStmt>* ret) {
   const auto& [op, inputs, outputs] = op_stmt.tuple();
-  CHECK(op.Has<const hlir::framework::Node*>());
-  if (GetOpPatternKind(op.Get<const hlir::framework::Node*>()) ==
+  CHECK(op.Has<const ::pir::Operation*>());
+  if (GetOpPatternKind(op.Get<const ::pir::Operation*>()) ==
       hlir::framework::OpPatternKind::kReduction) {
-    tReduceInit<const hlir::framework::Node*> init_op{
-        op.Get<const hlir::framework::Node*>()};
+    tReduceInit<const ::pir::Operation*> init_op{
+        op.Get<const ::pir::Operation*>()};
     (*ret)->emplace_back(OpStmt{init_op, List<Arg>{}, outputs});
 
-    tReduceAcc<const hlir::framework::Node*> acc_op{
-        op.Get<const hlir::framework::Node*>()};
+    tReduceAcc<const pir::Operation*> acc_op{op.Get<const ::pir::Operation*>()};
     (*ret)->emplace_back(OpStmt{acc_op, inputs, outputs});
     return true;
   } else {
@@ -172,7 +162,7 @@ void CollectRewritedOpStmts(const OpStmt& op_stmt, List<OpStmt>* ret) {
 }
 
 List<OpStmt> MakeOpStmts(
-    const std::shared_ptr<hlir::framework::Graph::Group>& group) {
+    const std::shared_ptr<hlir::framework::pir::Group>& group) {
   List<OpStmt> ret{};
 
   VisitEachOpStmt(group, [&](const auto& op_stmt) {
@@ -213,7 +203,7 @@ std::shared_ptr<IGroup> MakeIGroup(const AnchorGroup& igroup_spec) {
 }
 
 std::vector<std::shared_ptr<IGroup>> GenerateIGroups(
-    const std::shared_ptr<hlir::framework::Graph::Group>& group) {
+    const std::shared_ptr<hlir::framework::pir::Group>& group) {
   std::vector<std::shared_ptr<IGroup>> ret{};
 
   List<OpStmt> op_stmts = MakeOpStmts(group);
@@ -227,7 +217,7 @@ std::vector<std::shared_ptr<IGroup>> GenerateIGroups(
 }
 
 std::shared_ptr<KGroup> GenerateKGroups(
-    const std::shared_ptr<hlir::framework::Graph::Group>& group,
+    const std::shared_ptr<hlir::framework::pir::Group>& group,
     const std::vector<std::shared_ptr<IGroup>>& igroups) {
   CHECK_EQ(igroups.size(), 1);
   return std::make_shared<KGroup>(group, igroups);
@@ -343,36 +333,34 @@ Tensor GetAnchorTensor(const std::shared_ptr<IGroup>& igroup) {
 }
 
 template <typename DoEachT>
-void VisitInputTensor(const hlir::framework::Graph::Group& group,
+void VisitInputTensor(const hlir::framework::pir::Group& group,
                       const DoEachT& DoEach) {
-  for (const auto* node_data : group.GetInputNodeDatas()) {
-    DoEach(node_data, group.graph_);
+  for (const ::pir::Value& node_data : group.GetInputOpValues()) {
+    DoEach(node_data);
   }
 }
 
 template <typename DoEachT>
-void VisitOutputTensor(const hlir::framework::Graph::Group& group,
+void VisitOutputTensor(const hlir::framework::pir::Group& group,
                        const DoEachT& DoEach) {
-  for (const auto& node_data : group.GetOutputNodeDatas()) {
-    DoEach(node_data, group.graph_);
+  for (const ::pir::Value& node_data : group.GetOutputOpValues()) {
+    DoEach(node_data);
   }
 }
 
 List<Tensor> MakeInputTensors(const std::shared_ptr<KGroup>& kgroup) {
   List<Tensor> ret{};
-  VisitInputTensor(*kgroup->cinn_group(),
-                   [&](const auto* node_data, const auto* graph) {
-                     ret->emplace_back(adapter::Tensor{node_data, graph});
-                   });
+  VisitInputTensor(*kgroup->cinn_group(), [&](const ::pir::Value& node_data) {
+    ret->emplace_back(adapter::Tensor{node_data});
+  });
   return ret;
 }
 
 List<Tensor> MakeOutputTensors(const std::shared_ptr<KGroup>& kgroup) {
   List<Tensor> ret{};
-  VisitOutputTensor(*kgroup->cinn_group(),
-                    [&](const auto* node_data, const auto* graph) {
-                      ret->emplace_back(adapter::Tensor{node_data, graph});
-                    });
+  VisitOutputTensor(*kgroup->cinn_group(), [&](const ::pir::Value& node_data) {
+    ret->emplace_back(adapter::Tensor{node_data});
+  });
   return ret;
 }
 
@@ -437,7 +425,7 @@ MapExpr GenerateMapExpr(const std::shared_ptr<KGroup>& kgroup) {
 }  // namespace
 
 MapExpr GenerateMapExpr(
-    const std::shared_ptr<hlir::framework::Graph::Group>& group) {
+    const std::shared_ptr<hlir::framework::pir::Group>& group) {
   const auto& igroups = GenerateIGroups(group);
 
   const auto& kgroup = GenerateKGroups(group, igroups);
@@ -445,18 +433,26 @@ MapExpr GenerateMapExpr(
   return GenerateMapExpr(kgroup);
 }
 
-namespace {}  // namespace
-
 void TryGenerateMapExprFromGraph(
-    const std::shared_ptr<cinn::hlir::framework::Graph>& graph) {
+    const hlir::framework::pir::GroupList& groups) {
   if (!FLAGS_cinn_enable_map_expr) {
     return;
   }
-  for (const auto& fusion_group : graph->fusion_groups) {
+  for (const auto& fusion_group : groups) {
     const auto& map_expr = GenerateMapExpr(fusion_group);
     VLOG(1) << ToTxtString(map_expr, fusion_group->group_id);
     fusion_group->set_map_expr_ctx(std::make_shared<MapExprCtx>(map_expr));
   }
 }
 
+void TryGenerateMapExprFromGroup(
+    const std::shared_ptr<hlir::framework::pir::Group>& fusion_group) {
+  if (!FLAGS_cinn_enable_map_expr) {
+    return;
+  }
+  const auto& map_expr = GenerateMapExpr(fusion_group);
+  VLOG(1) << ToTxtString(map_expr, fusion_group->group_id);
+  fusion_group->set_map_expr_ctx(std::make_shared<MapExprCtx>(map_expr));
+}
+
 }  // namespace cinn::adt
diff --git a/paddle/cinn/adt/generate_map_expr.h b/paddle/cinn/adt/generate_map_expr.h
index c604dd9f070c06..61b5906c8138a3 100644
--- a/paddle/cinn/adt/generate_map_expr.h
+++ b/paddle/cinn/adt/generate_map_expr.h
@@ -14,19 +14,25 @@
 
 #pragma once
 
+#include <vector>
+
 #include "paddle/cinn/adt/m_expr.h"
-#include "paddle/cinn/adt/m_ir.h"
-#include "paddle/cinn/hlir/framework/graph.h"
 
-namespace cinn::adt {
+namespace cinn::hlir::framework::pir {
+
+struct Group;
+using GroupList = std::vector<std::shared_ptr<Group>>;
 
-class IGroup;
-class KGroup;
+}  // namespace cinn::hlir::framework::pir
+
+namespace cinn::adt {
 
 MapExpr GenerateMapExpr(
-    const std::shared_ptr<cinn::hlir::framework::Graph::Group>& group);
+    const std::shared_ptr<hlir::framework::pir::Group>& group);
+
+void TryGenerateMapExprFromGraph(const hlir::framework::pir::GroupList& groups);
 
-void TryGenerateMapExprFromGraph(
-    const std::shared_ptr<cinn::hlir::framework::Graph>& graph);
+void TryGenerateMapExprFromGroup(
+    const std::shared_ptr<hlir::framework::pir::Group>& fusion_group);
 
 }  // namespace cinn::adt
diff --git a/paddle/cinn/adt/inline_translator.h b/paddle/cinn/adt/inline_translator.h
index 5298d17ffadcda..2cd3a44bc7dd05 100644
--- a/paddle/cinn/adt/inline_translator.h
+++ b/paddle/cinn/adt/inline_translator.h
@@ -15,47 +15,12 @@
 #pragma once
 
 #include "paddle/cinn/adt/adt.h"
+#include "paddle/cinn/adt/inline_translator_trait.h"
 #include "paddle/cinn/adt/m_expr.h"
 #include "paddle/cinn/adt/tree.h"
 
 namespace cinn::adt {
 
-template <template <typename> class MapT>
-struct InlineTranslatorTrait;
-
-template <>
-struct InlineTranslatorTrait<MapStmt> final {
-  template <typename T>
-  static List<T> GetTreeInnerNodeChildren(const MapStmt<T>& map_stmt) {
-    const auto& [iterators, stmts] = map_stmt.tuple();
-    return stmts;
-  }
-
-  template <typename SrcTreeT, typename DstTreeT>
-  static MapStmt<DstTreeT> ConvertMap(const MapStmt<SrcTreeT>& src_map,
-                                      const List<DstTreeT>& dst_children) {
-    const auto& [iterators, src_children] = src_map.tuple();
-    return MapStmt<DstTreeT>{iterators, dst_children};
-  }
-};
-
-// OpCall T = (Op, [T])
-template <>
-struct InlineTranslatorTrait<OpCall> final {
-  template <typename T>
-  static List<T> GetTreeInnerNodeChildren(const OpCall<T>& op_call) {
-    const auto& [op, tensors] = op_call.tuple();
-    return tensors;
-  }
-
-  template <typename SrcTreeT, typename DstTreeT>
-  static OpCall<DstTreeT> ConvertMap(const OpCall<SrcTreeT>& src_map,
-                                     const List<DstTreeT>& dst_children) {
-    const auto& [op, _] = src_map.tuple();
-    return OpCall<DstTreeT>{op, dst_children};
-  }
-};
-
 template <template <typename> class MapT,
           template <typename>
           class OpCallT,
diff --git a/paddle/cinn/adt/inline_translator_trait.h b/paddle/cinn/adt/inline_translator_trait.h
new file mode 100644
index 00000000000000..665a27c8a18a53
--- /dev/null
+++ b/paddle/cinn/adt/inline_translator_trait.h
@@ -0,0 +1,58 @@
+// Copyright (c) 2023 PaddlePaddle 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.
+
+#pragma once
+
+#include "paddle/cinn/adt/adt.h"
+#include "paddle/cinn/adt/m_expr.h"
+
+namespace cinn::adt {
+
+template <template <typename> class MapT>
+struct InlineTranslatorTrait;
+
+template <>
+struct InlineTranslatorTrait<MapStmt> final {
+  template <typename T>
+  static List<T> GetTreeInnerNodeChildren(const MapStmt<T>& map_stmt) {
+    const auto& [iterators, stmts] = map_stmt.tuple();
+    return stmts;
+  }
+
+  template <typename SrcTreeT, typename DstTreeT>
+  static MapStmt<DstTreeT> ConvertMap(const MapStmt<SrcTreeT>& src_map,
+                                      const List<DstTreeT>& dst_children) {
+    const auto& [iterators, src_children] = src_map.tuple();
+    return MapStmt<DstTreeT>{iterators, dst_children};
+  }
+};
+
+// OpCall T = (Op, [T])
+template <>
+struct InlineTranslatorTrait<OpCall> final {
+  template <typename T>
+  static List<T> GetTreeInnerNodeChildren(const OpCall<T>& op_call) {
+    const auto& [op, tensors] = op_call.tuple();
+    return tensors;
+  }
+
+  template <typename SrcTreeT, typename DstTreeT>
+  static OpCall<DstTreeT> ConvertMap(const OpCall<SrcTreeT>& src_map,
+                                     const List<DstTreeT>& dst_children) {
+    const auto& [op, _] = src_map.tuple();
+    return OpCall<DstTreeT>{op, dst_children};
+  }
+};
+
+}  // namespace cinn::adt
diff --git a/paddle/cinn/adt/kgroup.cc b/paddle/cinn/adt/kgroup.cc
index a2456092dfb1e0..39edc167429a3b 100644
--- a/paddle/cinn/adt/kgroup.cc
+++ b/paddle/cinn/adt/kgroup.cc
@@ -32,7 +32,7 @@ std::size_t GetTensorNumel(const Tensor& tensor) {
   return tensor.Get<adapter::Tensor>().GetNumel();
 }
 
-const std::vector<int32_t>& GetTensorShape(const Tensor& tensor) {
+std::vector<int32_t> GetTensorShape(const Tensor& tensor) {
   CHECK(tensor.Has<adapter::Tensor>());
   return tensor.Get<adapter::Tensor>().GetShape();
 }
@@ -44,7 +44,7 @@ List<LoopSize> KGroup::GetDefaultScheduleSizes(
   List<LoopSize> ret{};
 
   const Tensor& tensor = igroup->anchor_tensor();
-  const auto tensor_shape = GetTensorShape(tensor);
+  const std::vector<int32_t> tensor_shape = GetTensorShape(tensor);
   for (int32_t dim : tensor_shape) {
     ret->emplace_back(LoopSize{dim});
   }
diff --git a/paddle/cinn/adt/kgroup.h b/paddle/cinn/adt/kgroup.h
index 83d1885873151e..60c1df3cad9fb4 100644
--- a/paddle/cinn/adt/kgroup.h
+++ b/paddle/cinn/adt/kgroup.h
@@ -19,6 +19,12 @@
 
 #include "paddle/cinn/adt/m_expr.h"
 
+namespace cinn::hlir::framework::pir {
+
+struct Group;
+
+}  // namespace cinn::hlir::framework::pir
+
 namespace cinn::adt {
 
 class IGroup;
@@ -32,11 +38,11 @@ using cinn::adt::LoopDescriptors;
 class KGroup final {
  public:
   explicit KGroup(
-      const std::shared_ptr<hlir::framework::Graph::Group>& cinn_group,
+      const std::shared_ptr<hlir::framework::pir::Group>& cinn_group,
       const std::vector<std::shared_ptr<IGroup>>& igroups)
       : cinn_group_(cinn_group), igroups_(igroups) {}
 
-  std::shared_ptr<hlir::framework::Graph::Group> cinn_group() const {
+  std::shared_ptr<hlir::framework::pir::Group> cinn_group() const {
     return CHECK_NOTNULL(cinn_group_.lock());
   }
 
@@ -51,7 +57,7 @@ class KGroup final {
       const std::shared_ptr<IGroup>& igroup) const;
 
  private:
-  std::weak_ptr<hlir::framework::Graph::Group> cinn_group_;
+  std::weak_ptr<hlir::framework::pir::Group> cinn_group_;
   // NOTE: Use single igroup temporarily. Actually KGroup contains
   // multiple IGroups
   std::vector<std::shared_ptr<IGroup>> igroups_;
diff --git a/paddle/cinn/adt/m_expr.h b/paddle/cinn/adt/m_expr.h
index fb4d4a57101a86..588b59d6faa512 100644
--- a/paddle/cinn/adt/m_expr.h
+++ b/paddle/cinn/adt/m_expr.h
@@ -24,14 +24,9 @@
 #include "paddle/cinn/adt/tags.h"
 #include "paddle/cinn/adt/tree.h"
 
-namespace cinn {
-namespace hlir {
-namespace framework {
-class Node;
-class NodeData;
-}  // namespace framework
-}  // namespace hlir
-}  // namespace cinn
+namespace pir {
+class Operation;
+}
 
 namespace cinn {
 namespace adt {
@@ -82,7 +77,7 @@ inline std::size_t GetHashValueImpl(const TempStorage& temp_storage) {
   return hash_value;
 }
 
-// SSAShadowTensor = (tSSAShadow Name, const Graph::NodeData*)
+// SSAShadowTensor = (tSSAShadow Name, adapter::Tensor)
 class SSAShadowTensor final : public Tuple<tSSAShadow<Name>, adapter::Tensor> {
  public:
   using Tuple<tSSAShadow<Name>, adapter::Tensor>::Tuple;
@@ -102,11 +97,13 @@ DEFINE_ADT_UNION(Tensor, adapter::Tensor, SSAShadowTensor, TempStorage);
 OVERRIDE_UNION_GET_HASH_VALUE(Tensor);
 OVERLOAD_OPERATOR_EQ_NE(Tensor, UnionEqual);
 
-// Op = const Node*
+// Op = const pir::Operation*
+//    | tReduceInit<const pir::Operation*>
+//    | tReduceAcc<const pir::Operation*>
 DEFINE_ADT_UNION(Op,
-                 const hlir::framework::Node*,
-                 tReduceInit<const hlir::framework::Node*>,
-                 tReduceAcc<const hlir::framework::Node*>);
+                 const ::pir::Operation*,
+                 tReduceInit<const ::pir::Operation*>,
+                 tReduceAcc<const ::pir::Operation*>);
 
 using Arg = Tensor;
 
diff --git a/paddle/cinn/adt/map_expr_ctx.h b/paddle/cinn/adt/map_expr_ctx.h
index 1438228cd346ad..bf0c8d001f2ffb 100644
--- a/paddle/cinn/adt/map_expr_ctx.h
+++ b/paddle/cinn/adt/map_expr_ctx.h
@@ -18,16 +18,16 @@
 #include <unordered_map>
 
 #include "paddle/cinn/adt/m_expr.h"
-#include "paddle/cinn/hlir/framework/node.h"
 #include "paddle/cinn/ir/lowered_func.h"
 #include "paddle/cinn/ir/utils/ir_copy.h"
+#include "paddle/pir/core/operation.h"
 
 namespace cinn::adt {
 
 class MapExprCtx final {
  public:
   using Node2LoweredFuncs =
-      std::unordered_map<hlir::framework::Node*, std::vector<ir::LoweredFunc>>;
+      std::unordered_map<::pir::Operation*, std::vector<ir::LoweredFunc>>;
 
   MapExprCtx(const MapExprCtx&) = delete;
   MapExprCtx(MapExprCtx&&) = delete;
@@ -37,7 +37,7 @@ class MapExprCtx final {
   const MapExpr& map_expr() const { return map_expr_; }
 
   void UpdateOpLoweredFuncKey(
-      hlir::framework::Node* node,
+      ::pir::Operation* node,
       const std::vector<ir::LoweredFunc>& lowered_funcs) {
     Node2LoweredFuncs* map = &node2lowered_funcs_;
     CHECK(map->emplace(node, ir::ir_utils::IRCopy(lowered_funcs)).second);
diff --git a/paddle/cinn/adt/naive_op_equation_context.cc b/paddle/cinn/adt/naive_op_equation_context.cc
index a215cb0c845bb0..c3c4327fe378e5 100644
--- a/paddle/cinn/adt/naive_op_equation_context.cc
+++ b/paddle/cinn/adt/naive_op_equation_context.cc
@@ -19,6 +19,8 @@
 #include "paddle/cinn/adt/naive_op_equation_context.h"
 #include "paddle/cinn/adt/op_arg_pos.h"
 #include "paddle/cinn/adt/print.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
+#include "paddle/cinn/utils/type_defs.h"
 
 #include "glog/logging.h"
 
@@ -37,7 +39,7 @@ std::vector<std::uint64_t> MakeTensorRanks(const List<Arg>& arg_lists) {
   return ret;
 }
 
-void GenerateOpEquationsImpl(const hlir::framework::Node* op_node,
+void GenerateOpEquationsImpl(const ::pir::Operation* op_node,
                              const OpStmt& op_stmt,
                              config::NaiveOpEquationContext* ctx) {
   const auto& [_, inputs, outputs] = op_stmt.tuple();
@@ -48,8 +50,10 @@ void GenerateOpEquationsImpl(const hlir::framework::Node* op_node,
   const auto& generate_equations =
       hlir::framework::Operator::GetAttrs<GenerateEquationFunc>(
           "generate_equations");
-  CHECK(generate_equations.Find(op_node->op()));
-  generate_equations[op_node->op()](ctx);
+  const hlir::framework::Operator* cinn_op = hlir::framework::Operator::Get(
+      hlir::framework::pir::CompatibleInfo::OpName(*op_node));
+  CHECK(generate_equations.Find(cinn_op));
+  generate_equations[cinn_op](ctx);
 }
 
 using GetArgStaticDimT = std::function<std::optional<std::int64_t>(
@@ -62,7 +66,7 @@ GetArgStaticDimT MakeGetArgStaticDimT(const List<Tensor>& tensors) {
       return std::nullopt;
     }
     CHECK(tensors->at(tensor_idx).Has<adapter::Tensor>());
-    const auto& tensor_shape =
+    const std::vector<int32_t> tensor_shape =
         tensors->at(tensor_idx).Get<adapter::Tensor>().GetShape();
     if (dim_idx >= tensor_shape.size()) {
       return std::nullopt;
@@ -71,15 +75,14 @@ GetArgStaticDimT MakeGetArgStaticDimT(const List<Tensor>& tensors) {
   };
 }
 
-void GenerateOpEquationsImpl(
-    const tReduceAcc<const hlir::framework::Node*>& op_node,
-    const OpStmt& op_stmt,
-    config::NaiveOpEquationContext* ctx) {
+void GenerateOpEquationsImpl(const tReduceAcc<const ::pir::Operation*>& op_node,
+                             const OpStmt& op_stmt,
+                             config::NaiveOpEquationContext* ctx) {
   GenerateOpEquationsImpl(op_node.value(), op_stmt, ctx);
 }
 
 void GenerateOpEquationsImpl(
-    const tReduceInit<const hlir::framework::Node*>& op_node,
+    const tReduceInit<const ::pir::Operation*>& op_node,
     const OpStmt& op_stmt,
     config::NaiveOpEquationContext* ctx) {
   // Do nothing
@@ -96,29 +99,25 @@ void GenerateOpEquations(const OpStmt& op_stmt,
       op.variant());
 }
 
-const hlir::framework::AttrMapType* GetOpAttrImpl(
-    const hlir::framework::Node* op_node) {
-  return &op_node->attrs.attr_store;
+cinn::utils::AttributeMap GetOpAttrImpl(const ::pir::Operation* op_node) {
+  return hlir::framework::pir::CompatibleInfo::ConvertAttributes(*op_node);
 }
 
-const hlir::framework::AttrMapType* GetOpAttrImpl(
-    const tReduceInit<const hlir::framework::Node*>&) {
-  static hlir::framework::AttrMapType empty{};
-  return &empty;
+cinn::utils::AttributeMap GetOpAttrImpl(
+    const tReduceInit<const ::pir::Operation*>&) {
+  return cinn::utils::AttributeMap{};
 }
 
-const hlir::framework::AttrMapType* GetOpAttrImpl(
-    const tReduceAcc<const hlir::framework::Node*>& op_node) {
+cinn::utils::AttributeMap GetOpAttrImpl(
+    const tReduceAcc<const ::pir::Operation*>& op_node) {
   return GetOpAttrImpl(op_node.value());
 }
 
-const hlir::framework::AttrMapType* GetOpAttr(const OpStmt& op_stmt) {
+cinn::utils::AttributeMap GetOpAttr(const OpStmt& op_stmt) {
   const auto& [op_node, inputs, outputs] = op_stmt.tuple();
 
-  const auto* attr = std::visit(
-      [&](const auto& impl) { return GetOpAttrImpl(impl); }, op_node.variant());
-
-  return attr;
+  return std::visit([&](const auto& impl) { return GetOpAttrImpl(impl); },
+                    op_node.variant());
 }
 
 std::shared_ptr<config::NaiveOpEquationContext> MakeContextAndGenerateEquations(
diff --git a/paddle/cinn/adt/naive_op_equation_context.h b/paddle/cinn/adt/naive_op_equation_context.h
index 29e0ee13a54900..6a8f849e47cdba 100644
--- a/paddle/cinn/adt/naive_op_equation_context.h
+++ b/paddle/cinn/adt/naive_op_equation_context.h
@@ -45,7 +45,7 @@ class NaiveOpEquationContext final : public OpEquationContext {
       const std::vector<std::uint64_t>& out_tensors_ranks,
       GetArgStaticDimT GetInDim,
       GetArgStaticDimT GetOutDim,
-      const hlir::framework::AttrMapType* attr_map_type)
+      cinn::utils::AttributeMap attr_map_type)
       : in_tensors_ranks_(in_tensors_ranks),
         out_tensors_ranks_(out_tensors_ranks),
         GetInDim_(GetInDim),
@@ -309,8 +309,8 @@ class NaiveOpEquationContext final : public OpEquationContext {
   }
 
   const utils::Attribute& GetAttribute(const std::string& name) const {
-    const auto& iter = attr_map_type_->find(name);
-    CHECK(iter != attr_map_type_->end())
+    const auto& iter = attr_map_type_.find(name);
+    CHECK(iter != attr_map_type_.end())
         << "Can't find Attribute with this name";
     return iter->second;
   }
@@ -320,7 +320,7 @@ class NaiveOpEquationContext final : public OpEquationContext {
   GetArgStaticDimT GetInDim_;
   GetArgStaticDimT GetOutDim_;
   Equations equations_;
-  const hlir::framework::AttrMapType* attr_map_type_;
+  const cinn::utils::AttributeMap attr_map_type_;
   FakeOpPlaceHolder fake_op_placeholder_;
 
   std::vector<IteratorTuple> in_iterator_tuples_;
diff --git a/paddle/cinn/adt/no_inline_translator.h b/paddle/cinn/adt/no_inline_translator.h
new file mode 100644
index 00000000000000..8daa07bd7bcb6c
--- /dev/null
+++ b/paddle/cinn/adt/no_inline_translator.h
@@ -0,0 +1,83 @@
+// Copyright (c) 2023 PaddlePaddle 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.
+
+#pragma once
+
+#include "paddle/cinn/adt/adt.h"
+#include "paddle/cinn/adt/inline_translator_trait.h"
+#include "paddle/cinn/adt/m_expr.h"
+#include "paddle/cinn/adt/tree.h"
+
+namespace cinn::adt {
+
+template <template <typename> class MapT,
+          template <typename>
+          class OpCallT,
+          typename TensorT>
+struct NoInlineTranslator final {
+  using SrcLeaf = Store<TensorT, OpCallT<Load<TensorT>>>;
+  using OpExpr = Tree<OpCallT, Load<TensorT>>;
+  using DstLeaf = Store<TensorT, OpExpr>;
+  using SrcTree = Tree<MapT, SrcLeaf>;
+  using DstTree = Tree<MapT, DstLeaf>;
+
+  static DstTree Call(const SrcTree& src_tree) { return Translate(src_tree); }
+
+ private:
+  static DstTree Translate(const SrcTree& src_tree) {
+    return std::visit([&](const auto& impl) { return TranslateImpl(impl); },
+                      src_tree.variant());
+  }
+
+  static DstTree TranslateImpl(const MapT<SrcTree>& src_map) {
+    return DstTree{TranslateMap(src_map)};
+  }
+
+  static MapT<DstTree> TranslateMap(const MapT<SrcTree>& src_map) {
+    const List<SrcTree> src_children =
+        InlineTranslatorTrait<MapT>::GetTreeInnerNodeChildren(src_map);
+    const List<DstTree> dst_children = TranslateList(src_children);
+    return InlineTranslatorTrait<MapT>::ConvertMap(src_map, dst_children);
+  }
+
+  static List<DstTree> TranslateList(const List<SrcTree>& src_children) {
+    List<DstTree> ret{};
+    for (const auto& src_child : *src_children) {
+      ret->emplace_back(Translate(src_child));
+    }
+    return ret;
+  }
+
+  static DstTree TranslateImpl(const SrcLeaf& src_leaf) {
+    return DstTree{TranslateLeaf(src_leaf)};
+  }
+
+  // using SrcLeaf = Store<TensorT, OpCallT<Load<TensorT>>>;
+  // using DstLeaf = Store<TensorT, OpExpr>;
+  static DstLeaf TranslateLeaf(const SrcLeaf& src_leaf) {
+    const auto& [tensor, op_call] = src_leaf.tuple();
+    const List<Load<TensorT>>& src_loads =
+        InlineTranslatorTrait<OpCallT>::GetTreeInnerNodeChildren(op_call);
+    List<OpExpr> dst_loads{};
+    for (const auto& src_load : *src_loads) {
+      dst_loads->emplace_back(src_load);
+    }
+    OpCallT<OpExpr> dst_op_call =
+        InlineTranslatorTrait<OpCallT>::ConvertMap(op_call, dst_loads);
+    OpExpr dst_op_call_tree = dst_op_call;
+    return DstLeaf{tensor, dst_op_call_tree};
+  }
+};
+
+}  // namespace cinn::adt
diff --git a/paddle/cinn/adt/print_utils/CMakeLists.txt b/paddle/cinn/adt/print_utils/CMakeLists.txt
index 006d043472ecbb..bf53a5cae4e35c 100644
--- a/paddle/cinn/adt/print_utils/CMakeLists.txt
+++ b/paddle/cinn/adt/print_utils/CMakeLists.txt
@@ -1,13 +1,16 @@
-core_gather_headers()
+if(NOT CINN_ONLY)
+  core_gather_headers()
 
-gather_srcs(
-  cinnapi_src
-  SRCS
-  print_constant.cc
-  print_equations.cc
-  print_loop_size.cc
-  print_map_expr.cc
-  print_schedule_descriptor.cc
-  print_schedule_dim.cc
-  print_schedule_mesh.cc
-  print_value.cc)
+  gather_srcs(
+    cinnapi_src
+    SRCS
+    print_constant.cc
+    print_equations.cc
+    print_loop_size.cc
+    print_map_expr.cc
+    print_schedule_descriptor.cc
+    print_schedule_dim.cc
+    print_schedule_mesh.cc
+    print_value.cc)
+
+endif()
diff --git a/paddle/cinn/adt/print_utils/print_equations.cc b/paddle/cinn/adt/print_utils/print_equations.cc
index 587607ec0e56b5..524d80c40b7681 100644
--- a/paddle/cinn/adt/print_utils/print_equations.cc
+++ b/paddle/cinn/adt/print_utils/print_equations.cc
@@ -18,6 +18,8 @@
 #include <string>
 
 #include "paddle/cinn/adt/equation_function.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
+#include "paddle/pir/core/operation.h"
 
 namespace cinn::adt {
 
@@ -28,14 +30,16 @@ std::string ToTxtString(const tDim<UniqueId>& constant) {
   return "dim_" + std::to_string(constant_unique_id);
 }
 
-std::string OpImpl(const hlir::framework::Node* op) { return op->op()->name; }
+std::string OpImpl(const ::pir::Operation* op) {
+  return hlir::framework::pir::CompatibleInfo::OpName(*op);
+}
 
-std::string OpImpl(const tReduceInit<const hlir::framework::Node*>& op) {
-  return op.value()->op()->name + "_init";
+std::string OpImpl(const tReduceInit<const ::pir::Operation*>& op) {
+  return OpImpl(op.value()) + "_init";
 }
 
-std::string OpImpl(const tReduceAcc<const hlir::framework::Node*>& op) {
-  return op.value()->op()->name + "_acc";
+std::string OpImpl(const tReduceAcc<const ::pir::Operation*>& op) {
+  return OpImpl(op.value()) + "_acc";
 }
 
 }  // namespace
diff --git a/paddle/cinn/adt/print_utils/print_map_expr.cc b/paddle/cinn/adt/print_utils/print_map_expr.cc
index e537b71fab343c..68e15e4b8a4c15 100644
--- a/paddle/cinn/adt/print_utils/print_map_expr.cc
+++ b/paddle/cinn/adt/print_utils/print_map_expr.cc
@@ -21,6 +21,7 @@
 #include "paddle/cinn/adt/print_utils/print_schedule_mesh.h"
 #include "paddle/cinn/adt/print_utils/print_value.h"
 #include "paddle/cinn/adt/schedule_descriptor.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
 
 namespace cinn::adt {
 
@@ -54,7 +55,8 @@ std::string ToTxtString(const Tensor& tensor) {
   CHECK(tensor.Has<adapter::Tensor>());
   std::string ret;
   ret += "t_";
-  ret += tensor.Get<adapter::Tensor>().node_data->id();
+  ret += hlir::framework::pir::CompatibleInfo::ValueName(
+      tensor.Get<adapter::Tensor>().node_data);
   return ret;
 }
 
@@ -86,18 +88,16 @@ std::string ToTxtString(const List<Arg>& out_args,
   return ret;
 }
 
-std::string ToTxtStringOpImpl(const hlir::framework::Node* op) {
-  return op->op()->name;
+std::string ToTxtStringOpImpl(const ::pir::Operation* op) {
+  return hlir::framework::pir::CompatibleInfo::OpName(*op);
 }
 
-std::string ToTxtStringOpImpl(
-    const tReduceInit<const hlir::framework::Node*>& op) {
-  return op.value()->op()->name + "_init";
+std::string ToTxtStringOpImpl(const tReduceInit<const ::pir::Operation*>& op) {
+  return ToTxtStringOpImpl(op.value()) + "_init";
 }
 
-std::string ToTxtStringOpImpl(
-    const tReduceAcc<const hlir::framework::Node*>& op) {
-  return op.value()->op()->name + "_acc";
+std::string ToTxtStringOpImpl(const tReduceAcc<const ::pir::Operation*>& op) {
+  return ToTxtStringOpImpl(op.value()) + "_acc";
 }
 
 std::string ToTxtString(const Op& op) {
diff --git a/paddle/cinn/adt/schedule_descriptor.cc b/paddle/cinn/adt/schedule_descriptor.cc
index c8ef34f43916e8..6c5b9d9185521f 100644
--- a/paddle/cinn/adt/schedule_descriptor.cc
+++ b/paddle/cinn/adt/schedule_descriptor.cc
@@ -23,7 +23,7 @@ namespace cinn::adt {
 
 namespace {
 
-const std::vector<int32_t>& GetTensorShape(const Tensor& tensor) {
+std::vector<int32_t> GetTensorShape(const Tensor& tensor) {
   CHECK(tensor.Has<adapter::Tensor>());
   return tensor.Get<adapter::Tensor>().GetShape();
 }
@@ -50,7 +50,7 @@ LoopDescriptors MakeNaiveScheduleDescriptor(
   const Tensor& tensor = igroup->anchor_tensor();
 
   List<LoopDescriptor> ret{};
-  const auto tensor_shape = GetTensorShape(tensor);
+  const std::vector<int32_t> tensor_shape = GetTensorShape(tensor);
   for (int32_t dim : tensor_shape) {
     ret->emplace_back(LoopDescriptor{Temporal{}, dim});
   }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc
index 7b08187cfe9b6d..e5b163a9ef8e7b 100644
--- a/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc
+++ b/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc
@@ -18,6 +18,7 @@
 
 #include <unordered_map>
 
+#include "paddle/cinn/adt/generate_map_expr.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/op_attribute.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
@@ -25,9 +26,12 @@
 #include "paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h"
 #include "paddle/cinn/hlir/dialect/runtime/ir/runtime_dialect.h"
 #include "paddle/cinn/hlir/framework/pir_compiler.h"
+#include "paddle/cinn/runtime/flags.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
 
+PD_DECLARE_bool(cinn_enable_map_expr);
+
 namespace cinn {
 namespace dialect {
 namespace ir {
@@ -165,6 +169,9 @@ std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
         auto ir_compiler = std::make_shared<cinn::hlir::framework::PirCompiler>(
             *program, target, scope);
         hlir::framework::PirCompilerManager::Instance().insert(ir_compiler);
+        if (FLAGS_cinn_enable_map_expr) {
+          adt::TryGenerateMapExprFromGroup(group);
+        }
         auto fn_ptr_res = ir_compiler->BuildCUDAJITInfo({group});
         std::unordered_map<std::string, ::pir::Attribute> op_attrs{
             {cinn::dialect::JitKernelOp::kAttrName,
diff --git a/paddle/cinn/hlir/framework/graph.h b/paddle/cinn/hlir/framework/graph.h
index c090397f1608c8..4c014043a1e248 100644
--- a/paddle/cinn/hlir/framework/graph.h
+++ b/paddle/cinn/hlir/framework/graph.h
@@ -27,10 +27,6 @@
 
 namespace cinn {
 
-namespace adt {
-class MapExprCtx;
-}  // namespace adt
-
 namespace hlir {
 namespace framework {
 
@@ -187,17 +183,6 @@ class Graph : public cinn::common::Graph {
 
     hlir::framework::OpPatternKind kind() const { return op_pattern_kind; }
 
-    adt::MapExprCtx* mut_map_expr_ctx() { return map_expr_ctx_.get(); }
-
-    const adt::MapExprCtx& map_expr_ctx() const {
-      return *CHECK_NOTNULL(map_expr_ctx_);
-    }
-
-    void set_map_expr_ctx(
-        const std::shared_ptr<adt::MapExprCtx>& map_expr_ctx) {
-      map_expr_ctx_ = map_expr_ctx;
-    }
-
    private:
     // input groups
     std::unordered_set<std::shared_ptr<Group>,
@@ -209,7 +194,6 @@ class Graph : public cinn::common::Graph {
                        SharedGroupHasher,
                        SharedGroupComparator>
         consumer_groups_;
-    std::shared_ptr<adt::MapExprCtx> map_expr_ctx_;
   };
   std::vector<std::shared_ptr<Group>> fusion_groups;
 
diff --git a/paddle/cinn/hlir/framework/op_lowering_impl.cc b/paddle/cinn/hlir/framework/op_lowering_impl.cc
index fb0cf399654517..1b47dbda611d79 100644
--- a/paddle/cinn/hlir/framework/op_lowering_impl.cc
+++ b/paddle/cinn/hlir/framework/op_lowering_impl.cc
@@ -14,13 +14,11 @@
 
 #include "paddle/cinn/hlir/framework/op_lowering_impl.h"
 
-#include "paddle/cinn/adt/map_expr_ctx.h"
 #include "paddle/cinn/ast_gen_ius/tensor_group.h"
 #include "paddle/cinn/hlir/framework/compile_error.h"
 #include "paddle/cinn/hlir/framework/graph_compiler_util.h"
 #include "paddle/cinn/hlir/framework/op_lowering_util.h"
 #include "paddle/cinn/hlir/op/external_api_registry.h"
-#include "paddle/cinn/hlir/pe/map_expr_to_ir.h"
 #include "paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h"
 #include "paddle/cinn/ir/schedule/ir_schedule.h"
 #include "paddle/cinn/optim/transform_gpu_forloop.h"
@@ -28,8 +26,6 @@
 
 PD_DECLARE_bool(cinn_use_cuda_vectorize);
 PD_DECLARE_bool(cinn_new_group_scheduler);
-PD_DECLARE_bool(cinn_enable_map_expr);
-PD_DECLARE_bool(cinn_map_expr_enable_schedule);
 
 namespace cinn {
 namespace hlir {
@@ -103,87 +99,6 @@ bool OpLowererImpl::NonFusibleScheduleDetermineFunction(Node* node) {
   return true;
 }
 
-void OpLowererImpl::LowerOpsForMapExpr(
-    const GroupPtr& group,
-    const std::vector<Node*>& nodes,
-    std::vector<ir::Tensor>* group_func_arg_tensors,
-    std::unordered_map<std::string, ir::Tensor>* tensor_map) {
-  auto& strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");
-  for (Node* node : nodes) {
-    // 1.Select Op impl
-    std::vector<Type> out_types;
-    std::vector<std::vector<int>> out_shapes;
-    std::vector<NodeData*> node_datas = GetAllNodeData(node);
-    for (const auto& node_data : node_datas) {
-      out_types.push_back(this->type_dict_.at(node_data->id()));
-      out_shapes.push_back(this->shape_dict_.at(node_data->id()));
-    }
-    std::vector<ir::Tensor> op_func_arg_tensors =
-        std::move(CollectInputTensor(node,
-                                     this->type_dict_,
-                                     this->shape_dict_,
-                                     group_func_arg_tensors,
-                                     tensor_map));
-    auto op_impl =
-        OpStrategy::SelectImpl(strategy[node->op()](node->attrs,
-                                                    op_func_arg_tensors,
-                                                    out_types,
-                                                    out_shapes,
-                                                    this->target_));
-
-    // 2.Perform the lower process of Op
-    std::vector<ir::LoweredFunc> funcs =
-        DoOpLower(op_impl, node, tensor_map, &op_func_arg_tensors);
-
-    group->mut_map_expr_ctx()->UpdateOpLoweredFuncKey(node, funcs);
-  }
-}
-
-/* Most of below codes copies from `PostProcess` function */
-std::vector<ir::LoweredFunc> OpLowererImpl::LowerMapExpr(
-    const GroupPtr& group,
-    const std::vector<Node*>& nodes,
-    bool do_op_schedule,
-    bool apply_group_schedule,
-    bool apply_pass,
-    std::vector<ir::Tensor>* group_func_arg_tensors,
-    std::unordered_map<std::string, ir::Tensor>* tensor_map) {
-  if (!FLAGS_cinn_map_expr_enable_schedule) {
-    do_op_schedule = false;
-    apply_group_schedule = false;
-    apply_pass = true;
-  }
-  VLOG(1) << "FLAGS_cinn_map_expr_enable_schedule = "
-          << FLAGS_cinn_map_expr_enable_schedule;
-  VLOG(1) << "do_op_schedule = " << do_op_schedule;
-  VLOG(1) << "apply_group_schedule = " << apply_group_schedule;
-  VLOG(1) << "apply_pass = " << apply_pass;
-  LowerOpsForMapExpr(group, nodes, group_func_arg_tensors, tensor_map);
-
-  ir::Expr func_body = adt::MapExprToIr(group->map_expr_ctx(), target_);
-
-  // 2.Do group schedule.
-  ir::ModuleExpr mod_expr({func_body});
-  ir::IRSchedule ir_sch(mod_expr);
-  ir_sch.MergeExprs();
-  VLOG(3) << "After lower, ir is: \n" << ir_sch.GetModule().GetExprs().at(0);
-  if (apply_group_schedule) {
-    DoGroupSchedule(ir_sch, group, *tensor_map);
-    VLOG(3) << "After group schedule, ir is: \n"
-            << ir_sch.GetModule().GetExprs().at(0);
-  }
-
-  // 3.Do post-processing,
-  // including preparing function args and temporary variables,
-  // applying low-level optimization passes, etc.
-  return PostProcess(group,
-                     *tensor_map,
-                     do_op_schedule,
-                     apply_pass,
-                     &ir_sch,
-                     group_func_arg_tensors);
-}
-
 std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
     const GroupPtr& group,
     bool apply_op_schedule,
@@ -200,15 +115,6 @@ std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
   std::vector<ir::Tensor> group_func_arg_tensors;
   std::unordered_map<std::string, ir::Tensor> tensor_map;
   bool do_op_schedule = apply_group_schedule || apply_op_schedule;
-  if (FLAGS_cinn_enable_map_expr) {
-    return LowerMapExpr(group,
-                        nodes,
-                        /*do_op_schedule=*/do_op_schedule,
-                        /*apply_group_schedule=*/apply_group_schedule,
-                        /*apply_pass=*/apply_pass,
-                        &group_func_arg_tensors,
-                        &tensor_map);
-  }
   std::vector<ir::Expr> func_bodies = LowerOps(nodes,
                                                do_op_schedule,
                                                schedule_determine_func,
diff --git a/paddle/cinn/hlir/framework/op_lowering_impl.h b/paddle/cinn/hlir/framework/op_lowering_impl.h
index 7c03c1d9903d0a..99be348d5be327 100644
--- a/paddle/cinn/hlir/framework/op_lowering_impl.h
+++ b/paddle/cinn/hlir/framework/op_lowering_impl.h
@@ -102,44 +102,6 @@ class OpLowererImpl : public OpLowererImplBase<GroupPtr> {
       ir::IRSchedule* ir_sch,
       std::vector<ir::Tensor>* group_func_arg_tensors);
 
-  /**
-   * @brief Lower an Op set to CINN IR.
-   * Compute and Lower will be performed one by one for each Op.
-   * @param group The group to be lowered.
-   * @param nodes The Op nodes to be lowered.
-   * @param apply_op_schedule Whether to schedule at Op level.
-   * @param schedule_determine_func Function used to determine which Ops to
-   * schedule.
-   * @param group_func_arg_tensors Tensors used as the group function arguments.
-   * @param tensor_map All tensors used for calculating the group.
-   * @return The lowered func bodies of Op set.
-   */
-  void LowerOpsForMapExpr(
-      const GroupPtr& group,
-      const std::vector<Node*>& nodes,
-      std::vector<ir::Tensor>* group_func_arg_tensors,
-      std::unordered_map<std::string, ir::Tensor>* tensor_map);
-
-  /**
-   * @brief Generate MapExpr and Lower it to std::vector<ir::LoweredFunc>
-   * @param group The group to be lowered.
-   * @param nodes The Op nodes to be lowered.
-   * @param do_op_schedule Mark whether the Op level schedule has been
-   * applied.
-   * @param apply_group_schedule Whether to schedule at group level.
-   * @param group_func_arg_tensors Tensors used as the group function arguments.
-   * @param tensor_map All tensors used for calculating the group.
-   * @return The lowered funcs after the post processing.
-   */
-  std::vector<ir::LoweredFunc> LowerMapExpr(
-      const GroupPtr& group,
-      const std::vector<Node*>& nodes,
-      bool do_op_schedule,
-      bool apply_group_schedule,
-      bool apply_pass,
-      std::vector<ir::Tensor>* group_func_arg_tensors,
-      std::unordered_map<std::string, ir::Tensor>* tensor_map);
-
   /**
    * @brief Lower an Op set to CINN IR.
    * Compute, Lower and optional Schedule will be performed one by one
diff --git a/paddle/cinn/hlir/framework/pir/group.h b/paddle/cinn/hlir/framework/pir/group.h
index 1d6813dad9d37a..1b34c22f641153 100644
--- a/paddle/cinn/hlir/framework/pir/group.h
+++ b/paddle/cinn/hlir/framework/pir/group.h
@@ -22,6 +22,11 @@
 #include "paddle/pir/core/value.h"
 
 namespace cinn {
+
+namespace adt {
+class MapExprCtx;
+}  // namespace adt
+
 namespace hlir {
 namespace framework {
 namespace pir {
@@ -31,6 +36,9 @@ using framework::OpPatternKind;
 struct Group {
  public:
   Group() = default;
+  Group(const Group&) = delete;
+  Group(Group&&) = delete;
+
   explicit Group(const std::vector<::pir::Operation*>& group_ops)
       : ops(group_ops) {}
 
@@ -83,7 +91,7 @@ struct Group {
     }
   };
 
-  std::vector<::pir::Operation*> CollectOps() {
+  std::vector<::pir::Operation*> CollectOps() const {
     if (fused_sub_groups.size()) {
       std::vector<::pir::Operation*> tmp_ops;
       for (auto& group : fused_sub_groups) {
@@ -109,7 +117,7 @@ struct Group {
     }
   }
 
-  std::unordered_set<::pir::Operation*> OpSet() {
+  std::unordered_set<::pir::Operation*> OpSet() const {
     std::unordered_set<::pir::Operation*> op_set;
     for (auto op : CollectOps()) {
       op_set.insert(op);
@@ -117,7 +125,7 @@ struct Group {
     return op_set;
   }
 
-  std::unordered_set<::pir::Value> GetInputOpValues() {
+  std::unordered_set<::pir::Value> GetInputOpValues() const {
     std::unordered_set<::pir::Value> group_inputs;
     auto ops_set = this->OpSet();
     // count all op's input Value
@@ -146,7 +154,8 @@ struct Group {
 
     return group_inputs;
   }
-  std::unordered_set<::pir::Value> GetOutputOpValues() {
+
+  std::unordered_set<::pir::Value> GetOutputOpValues() const {
     std::unordered_set<::pir::Value> group_outputs;
 
     for (auto op : this->output_ops) {
@@ -163,6 +172,19 @@ struct Group {
 
   std::string GetFuncName() { return "fn_" + group_id + unique_id; }
 
+  std::shared_ptr<adt::MapExprCtx> mut_map_expr_ctx() {
+    CHECK_NOTNULL(map_expr_ctx_);
+    return map_expr_ctx_;
+  }
+
+  const adt::MapExprCtx& map_expr_ctx() const {
+    return *CHECK_NOTNULL(map_expr_ctx_);
+  }
+
+  void set_map_expr_ctx(const std::shared_ptr<adt::MapExprCtx>& map_expr_ctx) {
+    map_expr_ctx_ = map_expr_ctx;
+  }
+
  public:
   const std::unordered_set<std::shared_ptr<Group>,
                            SharedGroupHasher,
@@ -213,6 +235,7 @@ struct Group {
                      SharedGroupHasher,
                      SharedGroupComparator>
       consumer_groups_;
+  std::shared_ptr<adt::MapExprCtx> map_expr_ctx_;
 };
 
 }  // namespace pir
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
index 349c8e70c036ef..585d7ed5d98ee0 100644
--- a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
+++ b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -16,19 +16,23 @@
 
 #include <string>
 
+#include "paddle/cinn/adt/map_expr_ctx.h"
 #include "paddle/cinn/ast_gen_ius/tensor_group.h"
+#include "paddle/cinn/hlir/framework/compile_error.h"
 #include "paddle/cinn/hlir/framework/pir/op_lowering_util.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/cinn/hlir/op/external_api_registry.h"
+#include "paddle/cinn/hlir/pe/map_expr_to_ir.h"
+#include "paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h"
 #include "paddle/cinn/ir/schedule/ir_schedule.h"
-#include "paddle/cinn/optim/transform_gpu_forloop.h"
-
-#include "paddle/cinn/hlir/framework/compile_error.h"
-#include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/cinn/lang/placeholder.h"
+#include "paddle/cinn/optim/transform_gpu_forloop.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/phi/core/ddim.h"
 
 PD_DECLARE_bool(cinn_use_cuda_vectorize);
+PD_DECLARE_bool(cinn_enable_map_expr);
+PD_DECLARE_bool(cinn_enable_map_expr_schedule);
 
 namespace cinn {
 namespace hlir {
@@ -182,6 +186,91 @@ bool OpLowererImpl::NonFusibleScheduleDetermineFunction(::pir::Operation* op) {
   return true;
 }
 
+void OpLowererImpl::LowerOpsForMapExpr(
+    const GroupPtr& group,
+    const std::vector<::pir::Operation*>& ops,
+    std::vector<ir::Tensor>* group_func_arg_tensors,
+    std::unordered_map<::pir::Value, ir::Tensor>* tensor_map) {
+  auto& strategy = Operator::GetAttrs<StrategyFunction>("CINNStrategy");
+  // for some op, it will output more tmp value and regard as
+  // XX_0, XX_1, so we log them in tmp_tensor_info;
+  std::unordered_map<std::string, ir::Tensor> tmp_tensor_info;
+  for (auto* op : ops) {
+    // 1.Select Op impl
+    std::vector<Type> out_types;
+    std::vector<std::vector<int>> out_shapes;
+
+    details::CollectOutputInfo(op, &out_types, &out_shapes);
+    VLOG(4) << "out_types.size(): " << out_types.size();
+    NodeAttr node_attrs = details::CollectAttrs(*op);
+
+    std::vector<ir::Tensor> op_func_arg_tensors =
+        details::CollectInputTensor(op, group_func_arg_tensors, tensor_map);
+    VLOG(4) << "input size:" << op_func_arg_tensors.size();
+
+    std::string cinn_op_name = CompatibleInfo::OpName(*op);
+    const hlir::framework::Operator* cinn_op = Operator::Get(cinn_op_name);
+    auto op_impl = OpStrategy::SelectImpl(strategy[cinn_op](
+        node_attrs, op_func_arg_tensors, out_types, out_shapes, this->target_));
+    // 2.Perform the lower process of Op
+    std::vector<ir::LoweredFunc> funcs = DoOpLower(
+        op_impl, op, tensor_map, &tmp_tensor_info, &op_func_arg_tensors);
+
+    group->mut_map_expr_ctx()->UpdateOpLoweredFuncKey(op, funcs);
+  }
+}
+
+/* Most of below codes copies from `PostProcess` function */
+std::vector<ir::LoweredFunc> OpLowererImpl::LowerMapExpr(
+    const GroupPtr& group,
+    const std::vector<::pir::Operation*>& ops,
+    bool apply_op_schedule,
+    bool apply_group_schedule,
+    std::vector<ir::Tensor>* group_func_arg_tensors,
+    std::unordered_map<::pir::Value, ir::Tensor>* tensor_map) {
+  if (FLAGS_cinn_enable_map_expr && FLAGS_cinn_enable_map_expr_schedule) {
+    apply_op_schedule = false;
+    apply_group_schedule = false;
+  }
+  VLOG(4) << "FLAGS_cinn_enable_map_expr_schedule = "
+          << FLAGS_cinn_enable_map_expr_schedule;
+  VLOG(4) << "apply_op_schedule = " << apply_op_schedule;
+  VLOG(4) << "apply_group_schedule = " << apply_group_schedule;
+
+  LowerOpsForMapExpr(group, ops, group_func_arg_tensors, tensor_map);
+
+  VLOG(4) << "Begin MapExprToIr";
+  ir::Expr func_body = adt::MapExprToIr(group->map_expr_ctx(), target_);
+
+  // 2.Do group schedule.
+  ir::ModuleExpr mod_expr({func_body});
+  ir::IRSchedule ir_sch(mod_expr);
+  ir_sch.MergeExprs();
+  VLOG(3) << "After lower, ir is: \n" << ir_sch.GetModule().GetExprs().at(0);
+  if (apply_group_schedule) {
+    std::unordered_set<std::string> output_tensor_names;
+    std::transform(
+        group->output_ops.begin(),
+        group->output_ops.end(),
+        std::inserter(output_tensor_names, output_tensor_names.begin()),
+        [](::pir::Operation* node) {
+          ::pir::Value node_data = node->result(0);
+          return hlir::framework::pir::CompatibleInfo::ValueName(node_data);
+        });
+    ir::StaticShapeGroupScheduler group_scheduler(
+        &ir_sch, output_tensor_names, target_);
+    group_scheduler.MapExprSchedule();
+    VLOG(3) << "After group schedule, ir is: \n"
+            << ir_sch.GetModule().GetExprs().at(0);
+  }
+
+  // 3.Do post-processing,
+  // including preparing function args and temporary variables,
+  // applying low-level optimization passes, etc.
+  return PostProcess(
+      group, *tensor_map, apply_op_schedule, &ir_sch, group_func_arg_tensors);
+}
+
 std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
     const GroupPtr& group,
     bool apply_op_schedule,
@@ -198,6 +287,14 @@ std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
   // XX_0, XX_1, so we log them in tmp_tensor_info;
   std::unordered_map<std::string, ir::Tensor> tmp_tensor_info;
   bool do_op_schedule = apply_group_schedule || apply_op_schedule;
+  if (FLAGS_cinn_enable_map_expr) {
+    return LowerMapExpr(group,
+                        ops,
+                        /*do_op_schedule=*/do_op_schedule,
+                        /*apply_group_schedule=*/apply_group_schedule,
+                        &group_func_arg_tensors,
+                        &tensor_map);
+  }
   std::vector<ir::Expr> func_bodies = LowerOps(ops,
                                                do_op_schedule,
                                                schedule_determine_func,
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.h b/paddle/cinn/hlir/framework/pir/op_lowering_impl.h
index 156e7a399ced51..aa86e580caa023 100644
--- a/paddle/cinn/hlir/framework/pir/op_lowering_impl.h
+++ b/paddle/cinn/hlir/framework/pir/op_lowering_impl.h
@@ -102,6 +102,39 @@ class OpLowererImpl : public OpLowererImplBase<GroupPtr> {
       ir::IRSchedule* ir_sch,
       std::vector<ir::Tensor>* group_func_arg_tensors);
 
+  /**
+   * @brief Lower an Op set to CINN IR.
+   * Compute and Lower will be performed one by one for each Op.
+   * @param group The group to be lowered.
+   * @param ops The Op to be lowered.
+   * @param group_func_arg_tensors Tensors used as the group function arguments.
+   * @param tensor_map All tensors used for calculating the group.
+   * @return The lowered func bodies of Op set.
+   */
+  void LowerOpsForMapExpr(
+      const GroupPtr& group,
+      const std::vector<::pir::Operation*>& ops,
+      std::vector<ir::Tensor>* group_func_arg_tensors,
+      std::unordered_map<::pir::Value, ir::Tensor>* tensor_map);
+
+  /**
+   * @brief Generate MapExpr and Lower it to std::vector<ir::LoweredFunc>
+   * @param group The group to be lowered.
+   * @param ops The Op to be lowered.
+   * @param apply_op_schedule Whether to schedule at Op level.
+   * @param apply_group_schedule Whether to schedule at group level.
+   * @param group_func_arg_tensors Tensors used as the group function arguments.
+   * @param tensor_map All tensors used for calculating the group.
+   * @return The lowered funcs after the post processing.
+   */
+  std::vector<ir::LoweredFunc> LowerMapExpr(
+      const GroupPtr& group,
+      const std::vector<::pir::Operation*>& ops,
+      bool apply_op_schedule,
+      bool apply_group_schedule,
+      std::vector<ir::Tensor>* group_func_arg_tensors,
+      std::unordered_map<::pir::Value, ir::Tensor>* tensor_map);
+
   /**
    * @brief Lower an Op set to CINN IR.
    * Compute, Lower and optional Schedule will be performed one by one
diff --git a/paddle/cinn/hlir/pe/CMakeLists.txt b/paddle/cinn/hlir/pe/CMakeLists.txt
index 00b920e2994730..6ac7787749fd40 100755
--- a/paddle/cinn/hlir/pe/CMakeLists.txt
+++ b/paddle/cinn/hlir/pe/CMakeLists.txt
@@ -14,8 +14,11 @@ gather_srcs(
   schedule.cc
   ir_schedule_pe.cc
   transform.cc
-  vision.cc
-  map_expr_to_ir.cc)
+  vision.cc)
+
+if(NOT CINN_ONLY)
+  gather_srcs(cinnapi_src SRCS map_expr_to_ir.cc)
+endif()
 
 cinn_cc_test(test_cinn_pe_elementwise SRCS pe_elementwise_test.cc DEPS cinncore)
 cinn_cc_test(test_cinn_pe_broadcast SRCS pe_broadcast_test.cc DEPS cinncore)
diff --git a/paddle/cinn/hlir/pe/map_expr_to_ir.cc b/paddle/cinn/hlir/pe/map_expr_to_ir.cc
index 521f5a4ed46508..84dd3e7d302eee 100644
--- a/paddle/cinn/hlir/pe/map_expr_to_ir.cc
+++ b/paddle/cinn/hlir/pe/map_expr_to_ir.cc
@@ -22,11 +22,15 @@
 #include "paddle/cinn/adt/m_expr.h"
 #include "paddle/cinn/adt/map_expr_ctx.h"
 #include "paddle/cinn/adt/match.h"
+#include "paddle/cinn/adt/no_inline_translator.h"
 #include "paddle/cinn/adt/print.h"
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_base.h"
 #include "paddle/cinn/ir/ir_printer.h"
+#include "paddle/cinn/runtime/flags.h"
+
+PD_DECLARE_bool(cinn_enable_map_expr_inline);
 
 namespace cinn::adt {
 
@@ -34,7 +38,7 @@ namespace {
 
 using IteratorInt = std::int32_t;
 using Node2LoweredFuncs =
-    std::unordered_map<hlir::framework::Node*, std::vector<ir::LoweredFunc>>;
+    std::unordered_map<::pir::Operation*, std::vector<ir::LoweredFunc>>;
 
 using TensorIteratorExpr4TensorT =
     std::function<adt::List<adt::TensorIteratorExpr>(const adt::Tensor&)>;
@@ -65,9 +69,9 @@ class MapExprToIrTranslator {
   }
 
  private:
-  ir::Expr GetStoreExprForOp(const hlir::framework::Node* op) const {
+  ir::Expr GetStoreExprForOp(const ::pir::Operation* op) const {
     const auto& iter =
-        node2lowered_funcs_->find(const_cast<hlir::framework::Node*>(op));
+        node2lowered_funcs_->find(const_cast<::pir::Operation*>(op));
     CHECK(iter != node2lowered_funcs_->end());
     const auto& lowered_funcs = iter->second;
     CHECK_EQ(lowered_funcs.size(), 1);
@@ -81,9 +85,9 @@ class MapExprToIrTranslator {
   }
 
   ir::Expr GetStoreExprForOp(
-      tReduceInit<const hlir::framework::Node*> op) const {
-    const auto& iter = node2lowered_funcs_->find(
-        const_cast<hlir::framework::Node*>(op.value()));
+      const tReduceInit<const ::pir::Operation*>& op) const {
+    const auto& iter =
+        node2lowered_funcs_->find(const_cast<::pir::Operation*>(op.value()));
     CHECK(iter != node2lowered_funcs_->end());
     const auto& lowered_funcs = iter->second;
     CHECK_EQ(lowered_funcs.size(), 1);
@@ -96,9 +100,9 @@ class MapExprToIrTranslator {
   }
 
   ir::Expr GetStoreExprForOp(
-      tReduceAcc<const hlir::framework::Node*> op) const {
-    const auto& iter = node2lowered_funcs_->find(
-        const_cast<hlir::framework::Node*>(op.value()));
+      const tReduceAcc<const ::pir::Operation*>& op) const {
+    const auto& iter =
+        node2lowered_funcs_->find(const_cast<::pir::Operation*>(op.value()));
     CHECK(iter != node2lowered_funcs_->end());
     const auto& lowered_funcs = iter->second;
     CHECK_EQ(lowered_funcs.size(), 1);
@@ -209,7 +213,12 @@ class MapExprToIrTranslator {
   }
 
   InlineStmt ConvertToInlineStmt(const InternalStmt& internal_stmt) const {
-    return InlineTranslator<MapStmt, OpCall, Tensor>::Call(internal_stmt);
+    if (FLAGS_cinn_enable_map_expr_inline) {
+      return InlineTranslator<MapStmt, OpCall, Tensor>::Call(internal_stmt);
+    } else {
+      return NoInlineTranslator<MapStmt, OpCall, Tensor>::Call(internal_stmt);
+    }
+    LOG(FATAL) << "Dead code";
   }
 
   std::optional<ir::Expr> TranslateOpExprImpl(
@@ -242,8 +251,61 @@ class MapExprToIrTranslator {
         op_expr.variant());
   }
 
+  std::optional<ir::Expr> MakeLoadExpr(
+      const ir::Expr& input_expr,
+      const List<OpExpr>& op_expr_children,
+      const IterExprs4TensorT& IterExprs4Tensor) const {
+    ir::Expr store_rvalue = ir::ir_utils::IRCopy(input_expr);
+    CHECK_EQ(store_rvalue->operands.size(), 0);
+    CHECK_EQ(op_expr_children->size(), 1);
+    store_rvalue.As<ir::Load>()->indices =
+        TranslateTensorIndex(op_expr_children->at(0), IterExprs4Tensor);
+    return store_rvalue;
+  }
+
+  std::optional<ir::Expr> MakeCallExpr(
+      const ir::Expr& input_expr,
+      const List<OpExpr>& op_expr_children,
+      const IterExprs4TensorT& IterExprs4Tensor) const {
+    ir::Expr store_rvalue = ir::ir_utils::IRCopy(input_expr);
+    CHECK_EQ(store_rvalue->operands.size(), 0);
+    CHECK(!op_expr_children->empty());
+    CHECK_EQ((store_rvalue.As<ir::Call>()->read_args.size()),
+             (op_expr_children->size()));
+    for (int i = 0; i < op_expr_children->size(); ++i) {
+      const auto& opt_operant = TranslateOpExpr(
+          op_expr_children->at(i), std::nullopt, IterExprs4Tensor);
+      if (opt_operant.has_value()) {
+        store_rvalue.As<ir::Call>()->read_args.at(i) = opt_operant.value();
+      } else {
+        store_rvalue.As<ir::Call>()->read_args.at(i).As<ir::Load>()->indices =
+            TranslateTensorIndex(op_expr_children->at(i), IterExprs4Tensor);
+      }
+    }
+    return store_rvalue;
+  }
+
+  std::optional<ir::Expr> MakeGeneralExpr(
+      const ir::Expr& input_expr,
+      const List<OpExpr>& op_expr_children,
+      const IterExprs4TensorT& IterExprs4Tensor) const {
+    ir::Expr store_rvalue = ir::ir_utils::IRCopy(input_expr);
+    CHECK_EQ(store_rvalue->operands.size(), op_expr_children->size());
+    for (int i = 0; i < op_expr_children->size(); ++i) {
+      const auto& opt_operant = TranslateOpExpr(
+          op_expr_children->at(i), std::nullopt, IterExprs4Tensor);
+      if (opt_operant.has_value()) {
+        store_rvalue->operands.at(i) = opt_operant.value();
+      } else {
+        store_rvalue->operands.at(i).As<ir::Load>()->indices =
+            TranslateTensorIndex(op_expr_children->at(i), IterExprs4Tensor);
+      }
+    }
+    return store_rvalue;
+  }
+
   std::optional<ir::Expr> TranslateOpCallImpl(
-      const hlir::framework::Node* op,
+      const ::pir::Operation*,
       const OpCall<OpExpr>& op_expr,
       const std::optional<Tensor>& opt_output_tensor,
       const IterExprs4TensorT& IterExprs4Tensor) const {
@@ -253,23 +315,13 @@ class MapExprToIrTranslator {
     ir::Expr store_rvalue = store_expr.value().As<ir::Store>()->value;
 
     if (store_rvalue.As<ir::Load>()) {
-      CHECK_EQ(store_rvalue->operands.size(), 0);
-      CHECK_EQ(op_expr_children->size(), 1);
-      store_rvalue.As<ir::Load>()->indices =
-          TranslateTensorIndex(op_expr_children->at(0), IterExprs4Tensor);
+      return MakeLoadExpr(store_rvalue, op_expr_children, IterExprs4Tensor);
+    } else if (store_rvalue.As<ir::Call>()) {
+      return MakeCallExpr(store_rvalue, op_expr_children, IterExprs4Tensor);
     } else {
       if (!op_expr_children->empty()) {
-        CHECK_EQ(store_rvalue->operands.size(), op_expr_children->size());
-        for (int i = 0; i < op_expr_children->size(); ++i) {
-          const auto& opt_operant = TranslateOpExpr(
-              op_expr_children->at(i), std::nullopt, IterExprs4Tensor);
-          if (opt_operant.has_value()) {
-            store_rvalue->operands.at(i) = opt_operant.value();
-          } else {
-            store_rvalue->operands.at(i).As<ir::Load>()->indices =
-                TranslateTensorIndex(op_expr_children->at(i), IterExprs4Tensor);
-          }
-        }
+        return MakeGeneralExpr(
+            store_rvalue, op_expr_children, IterExprs4Tensor);
       } else {
         // Do nothing
       }
@@ -278,7 +330,7 @@ class MapExprToIrTranslator {
   }
 
   std::optional<ir::Expr> TranslateOpCallImpl(
-      const tReduceInit<const hlir::framework::Node*>& op,
+      const tReduceInit<const ::pir::Operation*>&,
       const OpCall<OpExpr>& op_expr,
       const std::optional<Tensor>& opt_output_tensor,
       const IterExprs4TensorT& IterExprs4Tensor) const {
@@ -294,7 +346,7 @@ class MapExprToIrTranslator {
   }
 
   std::optional<ir::Expr> TranslateOpCallImpl(
-      const tReduceAcc<const hlir::framework::Node*>& op,
+      const tReduceAcc<const ::pir::Operation*>&,
       const OpCall<OpExpr>& op_expr,
       const std::optional<Tensor>& opt_output_tensor,
       const IterExprs4TensorT& IterExprs4Tensor) const {
diff --git a/paddle/cinn/hlir/pe/map_expr_to_ir.h b/paddle/cinn/hlir/pe/map_expr_to_ir.h
index 5f7aff6dfc4350..f4d37ad05c3527 100644
--- a/paddle/cinn/hlir/pe/map_expr_to_ir.h
+++ b/paddle/cinn/hlir/pe/map_expr_to_ir.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/cinn/adt/m_expr.h"
+#include "paddle/cinn/adt/map_expr_ctx.h"
 #include "paddle/cinn/ir/ir.h"
 
 namespace cinn::common {
diff --git a/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.cc b/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.cc
index 8c2ae6a6799c9e..21ef03bd6d5b36 100644
--- a/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.cc
+++ b/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.cc
@@ -143,6 +143,13 @@ void StaticShapeGroupScheduler::Schedule() {
 #endif
 }
 
+void StaticShapeGroupScheduler::MapExprSchedule() {
+  DoComputeInline();
+#ifdef CINN_WITH_CUDA
+  AllocateStorage();
+#endif
+}
+
 std::vector<std::pair<SymbolicCondition, ir::Expr>>
 StaticShapeGroupScheduler::GetIRs() {
   return {{Expr(1), ir_sch_->GetModule().GetExprs()[0]}};
diff --git a/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h b/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h
index b2b89c392bdc05..81d71a853dbfd5 100644
--- a/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h
+++ b/paddle/cinn/ir/group_schedule/st_shape_group_scheduler.h
@@ -51,6 +51,8 @@ class StaticShapeGroupScheduler : public GroupScheduler {
 
   void Schedule() override;
 
+  void MapExprSchedule();
+
   std::vector<std::pair<SymbolicCondition, ir::Expr>> GetIRs() override;
 
  private:
diff --git a/paddle/cinn/pybind/frontend.cc b/paddle/cinn/pybind/frontend.cc
index 3fa0abc78f8fd2..aafa9bedf40d07 100644
--- a/paddle/cinn/pybind/frontend.cc
+++ b/paddle/cinn/pybind/frontend.cc
@@ -203,8 +203,6 @@ void BindFrontend(pybind11::module *m) {
 
             auto graph = Optimize(&self, fetch_ids, target, passes);
 
-            cinn::adt::TryGenerateMapExprFromGraph(graph);
-
             scope = hlir::framework::BuildScope(target, graph, scope);
             hlir::framework::CompilationContext context(graph, scope, target);
 
diff --git a/paddle/cinn/runtime/flags.cc b/paddle/cinn/runtime/flags.cc
index d2c141e3b79f52..afe6342c26a6f2 100644
--- a/paddle/cinn/runtime/flags.cc
+++ b/paddle/cinn/runtime/flags.cc
@@ -79,10 +79,13 @@ PD_DEFINE_bool(cinn_enable_map_expr,
                BoolFromEnv("FLAGS_cinn_enable_map_expr", false),
                "It controls whether to use cinn with map_expr");
 
-PD_DEFINE_bool(
-    cinn_map_expr_enable_schedule,
-    BoolFromEnv("FLAGS_cinn_map_expr_enable_schedule", false),
-    "It controls whether to use schedule and pass when enables map_expr");
+PD_DEFINE_bool(cinn_enable_map_expr_schedule,
+               BoolFromEnv("FLAGS_cinn_enable_map_expr_schedule", false),
+               "It controls whether to schedule by map_expr");
+
+PD_DEFINE_bool(cinn_enable_map_expr_inline,
+               BoolFromEnv("FLAGS_cinn_enable_map_expr_inline", false),
+               "It controls whether to inline by map_expr");
 
 PD_DEFINE_bool(
     cinn_use_custom_call,
diff --git a/paddle/fluid/framework/paddle2cinn/cinn_compiler.cc b/paddle/fluid/framework/paddle2cinn/cinn_compiler.cc
index 34e43ac5d71ee6..47bc076ca689aa 100644
--- a/paddle/fluid/framework/paddle2cinn/cinn_compiler.cc
+++ b/paddle/fluid/framework/paddle2cinn/cinn_compiler.cc
@@ -321,8 +321,6 @@ std::unique_ptr<CinnCompiledObject> CinnCompiler::CompileGraph(
           << target.arch_str() << "), and its related graph:\n"
           << cinn_graph->Visualize();
 
-  cinn::adt::TryGenerateMapExprFromGraph(cinn_graph);
-
   auto scope = BuildScope(target, cinn_graph);
   CompilationContext context(cinn_graph, scope, target);
   context.with_instantiate_variables = false;
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 381b29804e2844..00186357bfac5e 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -100,7 +100,7 @@ if(WITH_TESTING)
     return()
   endif()
   # The following unittests only run in PR-CI-CINN
-  if(WITH_CINN AND NOT CINN_ONMLY)
+  if(WITH_CINN AND NOT CINN_ONLY)
     add_subdirectory(ir/pir/cinn)
   endif()
 
diff --git a/test/cinn/CMakeLists.txt b/test/cinn/CMakeLists.txt
index 3a25d7e1ef182a..470af5313df9c2 100644
--- a/test/cinn/CMakeLists.txt
+++ b/test/cinn/CMakeLists.txt
@@ -310,27 +310,6 @@ if(WITH_GPU)
       WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
   endforeach()
 
-  # adt test
-  file(
-    GLOB CINN_ADT_TEST
-    RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
-    "adt/test_*.py")
-  set(EXCLUDE_ADT_TEST test_add_inline)
-
-  foreach(adt_test_name ${EXCLUDE_ADT_TEST})
-    list(REMOVE_ITEM CINN_ADT_TEST adt/${adt_test_name}.py)
-  endforeach()
-
-  foreach(adt_test_name ${CINN_ADT_TEST})
-    string(REGEX REPLACE ".py" "" adt_test_name ${adt_test_name})
-    add_test(
-      NAME ${adt_test_name}
-      COMMAND
-        ${CMAKE_COMMAND} -E env
-        PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/cinn:$ENV{PYTHONPATH}
-        FLAGS_cinn_enable_map_expr=True python3
-        ${CMAKE_CURRENT_SOURCE_DIR}/${adt_test_name}.py
-      WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
-  endforeach()
-
 endif()
+
+add_subdirectory(adt)
diff --git a/test/cinn/adt/CMakeLists.txt b/test/cinn/adt/CMakeLists.txt
new file mode 100644
index 00000000000000..1b80d84300b422
--- /dev/null
+++ b/test/cinn/adt/CMakeLists.txt
@@ -0,0 +1,21 @@
+if(WITH_GPU AND NOT CINN_ONLY)
+  file(
+    GLOB CINN_PIR_TEST
+    RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+    "test_*.py")
+
+  foreach(cinn_pir_test_name ${CINN_PIR_TEST})
+    string(REGEX REPLACE ".py" "" cinn_pir_test_name ${cinn_pir_test_name})
+    add_test(
+      NAME ${cinn_pir_test_name}
+      COMMAND
+        ${CMAKE_COMMAND} -E env
+        PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+        FLAGS_enable_pir_api=1 FLAGS_cinn_enable_map_expr=1 ${PYTHON_EXECUTABLE}
+        ${CMAKE_CURRENT_SOURCE_DIR}/${cinn_pir_test_name}.py
+      WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+    set_tests_properties(${cinn_pir_test_name} PROPERTIES LABELS
+                                                          "RUN_TYPE=CINN")
+  endforeach()
+
+endif()
diff --git a/test/cinn/adt/test_add_inline.py b/test/cinn/adt/test_add_inline.py
deleted file mode 100755
index df92e86677abcc..00000000000000
--- a/test/cinn/adt/test_add_inline.py
+++ /dev/null
@@ -1,60 +0,0 @@
-#!/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 TestMapExprAddFusion(unittest.TestCase):
-    def setUp(self):
-        self.inputs = {
-            "x": np.random.uniform(-1.0, 1.0, [1024, 1024]).astype("float32"),
-            "y": np.random.uniform(-1.0, 1.0, [1024, 1024]).astype("float32"),
-            "z": np.random.uniform(-1.0, 1.0, [1024, 1024]).astype("float32"),
-        }
-
-    def test_add_fusion(self):
-        builder = NetBuilder("TestMapExprAddFusion")
-        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
-        y = builder.create_input(Float(32), self.inputs["y"].shape, "y")
-        z = builder.create_input(Float(32), self.inputs["z"].shape, "z")
-
-        a = builder.elementwise_add(x, y)
-        out = builder.elementwise_add(a, z)
-
-        prog = builder.build()
-        target = DefaultNVGPUTarget()
-        result = prog.build_and_get_output(
-            target,
-            [x, y, z],
-            [self.inputs["x"], self.inputs["y"], self.inputs["z"]],
-            [out],
-            passes=[],
-            scope=None,
-        )
-
-        np.testing.assert_allclose(
-            result[0].numpy(target),
-            self.inputs["x"] + self.inputs["y"] + self.inputs["z"],
-            err_msg="TestMapExprAddFusion failed!",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/test/cinn/adt/test_broadcast_expr.py b/test/cinn/adt/test_broadcast_expr.py
deleted file mode 100755
index a9caea37c6d064..00000000000000
--- a/test/cinn/adt/test_broadcast_expr.py
+++ /dev/null
@@ -1,59 +0,0 @@
-#!/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 TestMapExprBroadcast(unittest.TestCase):
-    def setUp(self):
-        self.inputs = {
-            "x1": np.random.uniform(-1.0, 1.0, [4, 16]).astype("float32"),
-            "x2": np.random.uniform(-1.0, 1.0, [16]).astype("float32"),
-        }
-
-    def test_broadcast(self):
-        builder = NetBuilder("TestMapExprBroadcast")
-        x1 = builder.create_input(Float(32), self.inputs["x1"].shape, "x1")
-        x2 = builder.create_input(Float(32), self.inputs["x2"].shape, "x2")
-        z = builder.elementwise_add(x1, x2)
-        out = builder.relu(z)
-        prog = builder.build()
-
-        target = DefaultNVGPUTarget()
-
-        result = prog.build_and_get_output(
-            target,
-            [x1, x2],
-            [self.inputs["x1"], self.inputs["x2"]],
-            [out],
-            passes=[],
-            scope=None,
-        )
-
-        np.testing.assert_allclose(
-            result[0].numpy(target),
-            np.maximum((self.inputs["x1"] + self.inputs["x2"]), 0),
-            err_msg="TestMapExprBroadcast failed!",
-        )
-        print("Finish Test")
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/test/cinn/adt/test_cinn_sub_graph_map_expr.py b/test/cinn/adt/test_cinn_sub_graph_map_expr.py
new file mode 100644
index 00000000000000..71a2145b5203f0
--- /dev/null
+++ b/test/cinn/adt/test_cinn_sub_graph_map_expr.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2022 PaddlePaddle 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
+
+import paddle
+
+
+def apply_to_static(net, use_cinn):
+    build_strategy = paddle.static.BuildStrategy()
+    build_strategy.build_cinn_pass = use_cinn
+    return paddle.jit.to_static(
+        net, build_strategy=build_strategy, full_graph=True
+    )
+
+
+def exp_sub(x):
+    y = paddle.exp(x)
+    z = y - x
+    return z
+
+
+class CINNSubGraphNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.fn = exp_sub
+
+    def forward(self, x):
+        out = self.fn(x)
+        return out
+
+
+class TestCinnSubGraphBase(unittest.TestCase):
+    """
+    Test Pir API + @to_static + CINN.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.shape = [64, 128]
+        self.axis = -1
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x = paddle.randn(self.shape, dtype="float32")
+        self.x.stop_gradient = False
+
+    def train(self, use_cinn):
+        paddle.seed(2022)
+        net = CINNSubGraphNet()
+        net = apply_to_static(net, use_cinn)
+        net.eval()
+        out = net(self.x)
+        return out
+
+    def test_forward(self):
+        cinn_out = self.train(use_cinn=True)
+        dy_out = self.train(use_cinn=False)
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/cinn/adt/test_naive_add.py b/test/cinn/adt/test_naive_add.py
deleted file mode 100755
index 304978ae5ceab5..00000000000000
--- a/test/cinn/adt/test_naive_add.py
+++ /dev/null
@@ -1,57 +0,0 @@
-#!/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, [1024, 1024]).astype("float32"),
-            "y": np.random.uniform(-1.0, 1.0, [1024, 1024]).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")
-
-        out = builder.elementwise_add(x, y)
-        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),
-            self.inputs["x"] + self.inputs["y"],
-            err_msg="TestMapExprNaiveAdd failed!",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/test/cinn/adt/test_naive_reduce.py b/test/cinn/adt/test_naive_reduce.py
deleted file mode 100644
index 217ef7695f539e..00000000000000
--- a/test/cinn/adt/test_naive_reduce.py
+++ /dev/null
@@ -1,50 +0,0 @@
-#!/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 TestMapExprNaiveReduce(unittest.TestCase):
-    def setUp(self):
-        self.inputs = {
-            "x": np.random.uniform(-1.0, 1.0, [2, 1024, 1024]).astype("float32")
-        }
-
-    def test_naive_reduce(self):
-        builder = NetBuilder("TestMapExprNaiveReduce")
-        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
-
-        out = builder.reduce_sum(x, [0], False)
-        prog = builder.build()
-        target = DefaultNVGPUTarget()
-
-        result = prog.build_and_get_output(
-            target, [x], [self.inputs["x"]], [out], passes=[], scope=None
-        )
-
-        np.testing.assert_allclose(
-            result[0].numpy(target),
-            np.sum(self.inputs["x"], axis=0),
-            err_msg="TestMapExprNaiveReduce failed!",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/test/cinn/adt/test_reduce_fusion.py b/test/cinn/adt/test_reduce_fusion.py
deleted file mode 100644
index 2432e53e21bcc5..00000000000000
--- a/test/cinn/adt/test_reduce_fusion.py
+++ /dev/null
@@ -1,62 +0,0 @@
-#!/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 TestMapExprReduceFusion(unittest.TestCase):
-    def setUp(self):
-        self.inputs = {
-            "x": np.random.uniform(-1.0, 1.0, [2, 1024, 1024]).astype(
-                "float32"
-            ),
-            "y": np.random.uniform(-1.0, 1.0, [2, 1024, 1024]).astype(
-                "float32"
-            ),
-        }
-
-    def test_reduce_fusion(self):
-        builder = NetBuilder("TestMapExprReduceFusion")
-        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
-        y = builder.create_input(Float(32), self.inputs["y"].shape, "y")
-
-        t = builder.elementwise_add(x, y)
-        out = builder.reduce_sum(t, [0], False)
-        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.sum(self.inputs["x"] + self.inputs["y"], axis=0),
-            err_msg="TestMapExprReduceFusion failed!",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/test/cinn/adt/test_reduce_schedule_mesh.py b/test/cinn/adt/test_reduce_schedule_mesh.py
deleted file mode 100644
index 10bde8159932b5..00000000000000
--- a/test/cinn/adt/test_reduce_schedule_mesh.py
+++ /dev/null
@@ -1,59 +0,0 @@
-#!/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 TestMapExprReduceScheduleMesh(unittest.TestCase):
-    def setUp(self):
-        self.inputs = {
-            "x": np.random.uniform(-1.0, 1.0, [32, 2048]).astype("float32"),
-            "y": np.random.uniform(-1.0, 1.0, [32, 2048]).astype("float32"),
-        }
-
-    def test_schedule_mesh(self):
-        builder = NetBuilder("TestMapExprReduceScheduleMesh")
-        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
-        y = builder.create_input(Float(32), self.inputs["y"].shape, "y")
-
-        t = builder.elementwise_add(x, y)
-        out = builder.reduce_sum(t, [0], False)
-        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.sum(self.inputs["x"] + self.inputs["y"], axis=0),
-            err_msg="TestMapExprReduceScheduleMesh failed!",
-        )
-
-
-if __name__ == "__main__":
-    unittest.main()