chore: import upstream snapshot with attribution

test/cpp/pir/pattern_rewrite/CMakeLists.txt

@@ -0,0 +1,24 @@
+paddle_test(pattern_rewrite_test SRCS pattern_rewrite_test.cc)
+
+paddle_test(drr_same_type_binding_test SRCS drr_same_type_binding_test.cc)
+
+paddle_test(drr_fuse_linear_test SRCS drr_fuse_linear_test.cc)
+
+paddle_test(drr_fuse_linear_param_grad_add_test SRCS
+            drr_fuse_linear_param_grad_add_test.cc)
+
+if(WITH_GPU)
+  paddle_test(drr_attention_fuse_test SRCS drr_attention_fuse_test.cc)
+endif()
+
+if(WITH_ONNXRUNTIME AND WIN32)
+  # Copy onnxruntime for some c++ test in Windows, since the test will
+  # be build only in CI, so suppose the generator in Windows is Ninja.
+  copy_onnx(pattern_rewrite_test)
+  copy_onnx(drr_same_type_binding_test)
+  copy_onnx(drr_fuse_linear_test)
+  copy_onnx(drr_fuse_linear_param_grad_add_test)
+  if(WITH_GPU)
+    copy_onnx(drr_attention_fuse_test)
+  endif()
+endif()

test/cpp/pir/pattern_rewrite/drr_attention_fuse_test.cc

@@ -0,0 +1,168 @@
+// 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 <gtest/gtest.h>
+
+#include <cstdint>
+#include <vector>
+
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/general/constant_folding_pass.h"
+#include "paddle/fluid/pir/transforms/general/dead_code_elimination_pass.h"
+#include "paddle/fluid/pir/transforms/gpu/multihead_matmul_fuse_pass.h"
+
+#include "paddle/phi/common/place.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/pass/pass_manager.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_DECLARE_KERNEL(multihead_matmul, GPU, ALL_LAYOUT);
+
+void BuildProgram(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp matmul_1_in_1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{1, 300, 256},
+                                             0.9,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+  // The first path to matmul with scale (q).
+  paddle::dialect::FullOp matmul_1_in_2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{256, 256},
+                                             1.1,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::MatmulOp matmul_1 = builder.Build<paddle::dialect::MatmulOp>(
+      matmul_1_in_1.out(), matmul_1_in_2.out(), false, false);
+
+  paddle::dialect::FullOp add_1_in_2 = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{256}, 1.5, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  paddle::dialect::AddOp add_1 =
+      builder.Build<paddle::dialect::AddOp>(matmul_1.out(), add_1_in_2.out());
+
+  paddle::dialect::ReshapeOp reshape_1 =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          add_1.out(), std::vector<int64_t>{0, 0, 8, 32});
+
+  paddle::dialect::TransposeOp transpose_1 =
+      builder.Build<paddle::dialect::TransposeOp>(reshape_1.out(),
+                                                  std::vector<int>{0, 2, 1, 3});
+
+  paddle::dialect::ScaleOp scale_op = builder.Build<paddle::dialect::ScaleOp>(
+      transpose_1.out(), 0.1767766922712326, 0.0, true);
+
+  // The second path to matmul (k).
+  paddle::dialect::FullOp matmul_2_in_2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{256, 256},
+                                             1.1,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::MatmulOp matmul_2 = builder.Build<paddle::dialect::MatmulOp>(
+      matmul_1_in_1.out(), matmul_2_in_2.out(), false, false);
+
+  paddle::dialect::FullOp add_2_in_2 = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{256}, 1.5, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::AddOp add_op2 =
+      builder.Build<paddle::dialect::AddOp>(matmul_2.out(), add_2_in_2.out());
+
+  paddle::dialect::ReshapeOp reshape_2 =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          add_op2.out(), std::vector<int64_t>{0, 0, 8, 32});
+
+  paddle::dialect::TransposeOp transpose_2 =
+      builder.Build<paddle::dialect::TransposeOp>(reshape_2.out(),
+                                                  std::vector<int>{0, 2, 1, 3});
+
+  // The third path to matmul (v).
+  paddle::dialect::FullOp matmul_3_in_2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{256, 256},
+                                             1.1,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+  paddle::dialect::MatmulOp matmul_3 = builder.Build<paddle::dialect::MatmulOp>(
+      matmul_1_in_1.out(), matmul_3_in_2.out(), false, false);
+
+  paddle::dialect::FullOp add_3_in_2 = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{256}, 1.5, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  paddle::dialect::AddOp add_3 =
+      builder.Build<paddle::dialect::AddOp>(matmul_3.out(), add_3_in_2.out());
+
+  paddle::dialect::ReshapeOp reshape_3 =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          add_3.out(), std::vector<int64_t>{0, 0, 8, 32});
+
+  paddle::dialect::TransposeOp transpose_3 =
+      builder.Build<paddle::dialect::TransposeOp>(reshape_3.out(),
+                                                  std::vector<int>{0, 2, 1, 3});
+
+  // softmax(qk)v
+  paddle::dialect::MatmulOp matmul_4 = builder.Build<paddle::dialect::MatmulOp>(
+      scale_op.out(), transpose_2.out(), false, true);
+
+  paddle::dialect::FullOp add_4_in_2 = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1, 8, 300, 300},
+      1.5,
+      phi::DataType::FLOAT32,
+      phi::CPUPlace());
+
+  paddle::dialect::AddOp add_4 =
+      builder.Build<paddle::dialect::AddOp>(matmul_4.out(), add_4_in_2.out());
+
+  paddle::dialect::SoftmaxOp softmax_op =
+      builder.Build<paddle::dialect::SoftmaxOp>(add_4.out(), -1);
+  paddle::dialect::MatmulOp matmul_5 = builder.Build<paddle::dialect::MatmulOp>(
+      softmax_op.out(), transpose_3.out(), false, false);
+
+  paddle::dialect::TransposeOp transpose_4 =
+      builder.Build<paddle::dialect::TransposeOp>(matmul_5.out(),
+                                                  std::vector<int>{0, 2, 1, 3});
+
+  paddle::dialect::ReshapeOp reshape_4 =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          transpose_4.out(), std::vector<int64_t>{0, 0, 256});
+
+  builder.Build<paddle::dialect::FetchOp>(reshape_4.out(), "out", 0);
+}
+
+TEST(DrrTest, AttentionFuse) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram(builder);
+  EXPECT_EQ(program.block()->size(), 33u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateMultiHeadMatmulFusePass());
+  std::unique_ptr<pir::Pass> constant_folding_pass =
+      pir::CreateConstantFoldingPass();
+  constant_folding_pass->Set(pir::Pass::kPlaceAttr,
+                             new phi::Place{phi::GPUPlace{}});
+  constant_folding_pass->Set(pir::Pass::kParamScopeAttr,
+                             new paddle::framework::Scope{});
+  pm.AddPass(std::move(constant_folding_pass));
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::Unavailable("pm fail to run program"));
+  EXPECT_EQ(program.block()->size(), 2u);
+}

test/cpp/pir/pattern_rewrite/drr_fuse_linear_param_grad_add_test.cc

@@ -0,0 +1,301 @@
+// 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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/gpu/fused_linear_param_grad_add_pass.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/pass/pass_manager.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
+
+void BuildProgram0(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+  paddle::dialect::FullOp full_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+  paddle::dialect::FullOp full_bias_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32}, 1.0);
+
+  paddle::dialect::MatmulOp matmul_op1 =
+      builder.Build<paddle::dialect::MatmulOp>(full_input_op1.out(),
+                                               full_weight_op1.out());
+  paddle::dialect::AddOp add_op1 = builder.Build<paddle::dialect::AddOp>(
+      matmul_op1.out(), full_bias_op1.out());
+
+  paddle::dialect::FullOp full_d_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::FullOp full_d_out_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::AddGradOp add_grad_op1 =
+      builder.Build<paddle::dialect::AddGradOp>(
+          matmul_op1.out(), full_bias_op1.out(), full_d_out_op1.out());
+
+  paddle::dialect::MatmulGradOp matmul_grad_op1 =
+      builder.Build<paddle::dialect::MatmulGradOp>(
+          full_input_op1.out(), full_weight_op1.out(), add_grad_op1.x_grad());
+
+  paddle::dialect::Add_Op add__op1 = builder.Build<paddle::dialect::Add_Op>(
+      full_d_weight_op1.out(), matmul_grad_op1.y_grad());
+
+  builder.Build<paddle::dialect::FetchOp>(add_op1.out(), "out", 0);
+  builder.Build<paddle::dialect::FetchOp>(add_grad_op1.y_grad(), "dbias", 1);
+  builder.Build<paddle::dialect::FetchOp>(add__op1.out(), "dweight", 2);
+}
+
+void BuildProgram1(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+  paddle::dialect::FullOp full_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::MatmulOp matmul_op1 =
+      builder.Build<paddle::dialect::MatmulOp>(full_input_op1.out(),
+                                               full_weight_op1.out());
+
+  paddle::dialect::FullOp full_d_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::FullOp full_d_out_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::MatmulGradOp matmul_grad_op1 =
+      builder.Build<paddle::dialect::MatmulGradOp>(
+          full_input_op1.out(), full_weight_op1.out(), full_d_out_op1.out());
+
+  paddle::dialect::Add_Op add__op1 = builder.Build<paddle::dialect::Add_Op>(
+      full_d_weight_op1.out(), matmul_grad_op1.y_grad());
+
+  builder.Build<paddle::dialect::FetchOp>(matmul_op1.out(), "out", 0);
+  builder.Build<paddle::dialect::FetchOp>(add__op1.out(), "dweight", 1);
+}
+
+void BuildProgram2(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+  paddle::dialect::FullOp full_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::MatmulOp matmul_op1 =
+      builder.Build<paddle::dialect::MatmulOp>(full_input_op1.out(),
+                                               full_weight_op1.out());
+
+  paddle::dialect::FullOp full_d_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::FullOp full_d_out_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::MatmulOp matmul_op2 =
+      builder.Build<paddle::dialect::MatmulOp>(
+          full_input_op1.out(), full_d_out_op1.out(), true, false);
+
+  paddle::dialect::Add_Op add__op1 = builder.Build<paddle::dialect::Add_Op>(
+      full_d_weight_op1.out(), matmul_op2.out());
+
+  builder.Build<paddle::dialect::FetchOp>(matmul_op1.out(), "out", 0);
+  builder.Build<paddle::dialect::FetchOp>(add__op1.out(), "dweight", 1);
+}
+
+void BuildProgram3(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+  paddle::dialect::FullOp full_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+  paddle::dialect::FullOp full_bias_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32}, 1.0);
+
+  paddle::dialect::MatmulOp matmul_op1 =
+      builder.Build<paddle::dialect::MatmulOp>(full_input_op1.out(),
+                                               full_weight_op1.out());
+  paddle::dialect::AddOp add_op1 = builder.Build<paddle::dialect::AddOp>(
+      matmul_op1.out(), full_bias_op1.out());
+
+  paddle::dialect::FullOp full_d_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::FullOp full_d_out_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{32, 32}, 1.5);
+
+  paddle::dialect::AddGradOp add_grad_op1 =
+      builder.Build<paddle::dialect::AddGradOp>(
+          matmul_op1.out(), full_bias_op1.out(), full_d_out_op1.out());
+
+  paddle::dialect::MatmulOp matmul_op2 =
+      builder.Build<paddle::dialect::MatmulOp>(
+          add_grad_op1.x_grad(), full_weight_op1.out(), false, true);
+
+  paddle::dialect::MatmulOp matmul_op3 =
+      builder.Build<paddle::dialect::MatmulOp>(
+          full_input_op1.out(), add_grad_op1.x_grad(), true, false);
+
+  paddle::dialect::Add_Op add__op1 = builder.Build<paddle::dialect::Add_Op>(
+      full_d_weight_op1.out(), matmul_op3.out());
+
+  builder.Build<paddle::dialect::FetchOp>(add_op1.out(), "out", 0);
+  builder.Build<paddle::dialect::FetchOp>(add_grad_op1.y_grad(), "dbias", 1);
+  builder.Build<paddle::dialect::FetchOp>(add__op1.out(), "dweight", 2);
+  builder.Build<paddle::dialect::FetchOp>(matmul_op2.out(), "dx", 3);
+}
+
+void BuildProgram4(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp input_x = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2048, 5120}, 1.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp input_dy = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2048, 6912}, 2.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp input_weight_grad =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{5120, 6912}, 1.0, phi::DataType::FLOAT16);
+
+  paddle::dialect::ReshapeOp reshape_x =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          input_x.out(), std::vector<int64_t>{4096, 5120});
+  paddle::dialect::ReshapeOp reshape_dy =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          input_dy.out(), std::vector<int64_t>{4096, 6912});
+  paddle::dialect::MatmulOp matmul_op =
+      builder.Build<paddle::dialect::MatmulOp>(
+          reshape_x.out(), reshape_dy.out(), true, false);
+  paddle::dialect::ReshapeOp output_dw =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          matmul_op.out(), std::vector<int64_t>{5120, 6912});
+  paddle::dialect::Add_Op add__op1 = builder.Build<paddle::dialect::Add_Op>(
+      input_weight_grad.out(), output_dw.out());
+  builder.Build<paddle::dialect::FetchOp>(add__op1.out(), "dw", 0);
+}
+
+bool verify_pass(const pir::Program &program) {
+  for (auto &op : *(program.block())) {
+    if (op.name() == paddle::dialect::FusedLinearParamGradAddOp::name()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+TEST(DrrTest, FusedLinearParamGradAdd0) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram0(builder);
+
+  EXPECT_EQ(program.block()->size(), 13u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Required pm.Run(&program) should be true"));
+  EXPECT_EQ(verify_pass(program), true);
+}
+
+TEST(DrrTest, FusedLinearParamGradAdd1) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram1(builder);
+
+  EXPECT_EQ(program.block()->size(), 9u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Required pm.Run(&program) should be true"));
+  EXPECT_EQ(verify_pass(program), true);
+}
+
+TEST(DrrTest, FusedLinearParamGradAdd2) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram2(builder);
+
+  EXPECT_EQ(program.block()->size(), 9u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Required pm.Run(&program) should be true"));
+  EXPECT_EQ(verify_pass(program), true);
+}
+
+TEST(DrrTest, FusedLinearParamGradAdd3) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram3(builder);
+
+  EXPECT_EQ(program.block()->size(), 15u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Required pm.Run(&program) should be true"));
+  EXPECT_EQ(verify_pass(program), true);
+}
+
+TEST(DrrTest, FusedMatmulReshapeMatmulAddPattern) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram4(builder);
+
+  EXPECT_EQ(program.block()->size(), 12u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Required pm.Run(&program) should be true"));
+  EXPECT_EQ(verify_pass(program), true);
+  EXPECT_EQ(program.block()->size(), 5u);
+}

test/cpp/pir/pattern_rewrite/drr_fuse_linear_test.cc

@@ -0,0 +1,152 @@
+// 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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/gpu/fused_gemm_epilogue_pass.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/pass/pass_manager.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
+
+void BuildProgram(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op1 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{1, 512, 64}, 1.5, phi::DataType::FLOAT16);
+  // linear 1
+  paddle::dialect::FullOp full_weight_op1 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{64, 64}, 1.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp full_bias_op1 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{64}, 1.0, phi::DataType::FLOAT16);
+  paddle::dialect::MatmulOp matmul_op1 =
+      builder.Build<paddle::dialect::MatmulOp>(full_input_op1.out(),
+                                               full_weight_op1.out());
+  paddle::dialect::AddOp add_op1 = builder.Build<paddle::dialect::AddOp>(
+      matmul_op1.out(), full_bias_op1.out());
+  // linear 2
+  paddle::dialect::FullOp full_weight_op2 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{64, 128}, 1.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp full_bias_op2 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{128}, 1.0, phi::DataType::FLOAT16);
+  paddle::dialect::MatmulOp matmul_op2 =
+      builder.Build<paddle::dialect::MatmulOp>(add_op1.out(),
+                                               full_weight_op2.out());
+  paddle::dialect::AddOp add_op2 = builder.Build<paddle::dialect::AddOp>(
+      matmul_op2.out(), full_bias_op2.out());
+  paddle::dialect::ReluOp relu_op =
+      builder.Build<paddle::dialect::ReluOp>(add_op2.out());
+  // linear 3
+  paddle::dialect::FullOp full_weight_op3 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{128, 64}, 1.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp full_bias_op3 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{64}, 1.0, phi::DataType::FLOAT16);
+  paddle::dialect::MatmulOp matmul_op3 =
+      builder.Build<paddle::dialect::MatmulOp>(relu_op.out(),
+                                               full_weight_op3.out());
+  paddle::dialect::AddOp add_op3 = builder.Build<paddle::dialect::AddOp>(
+      matmul_op3.out(), full_bias_op3.out());
+  paddle::dialect::GeluOp gelu_op1 =
+      builder.Build<paddle::dialect::GeluOp>(add_op3.out());
+  // linear 4
+  paddle::dialect::FullOp full_weight_op4 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{64, 64}, 1.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp full_bias_op4 =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{64}, 1.0, phi::DataType::FLOAT16);
+  paddle::dialect::MatmulOp matmul_op4 =
+      builder.Build<paddle::dialect::MatmulOp>(gelu_op1.out(),
+                                               full_weight_op4.out());
+  paddle::dialect::AddOp add_op4 = builder.Build<paddle::dialect::AddOp>(
+      matmul_op4.out(), full_bias_op4.out());
+  paddle::dialect::GeluOp gelu_op2 =
+      builder.Build<paddle::dialect::GeluOp>(add_op4.out());
+
+  // backward
+  paddle::dialect::FullOp full_grad_op = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1, 512, 64}, 1.0, phi::DataType::FLOAT16);
+
+  paddle::dialect::GeluGradOp gelu_op2_grad =
+      builder.Build<paddle::dialect::GeluGradOp>(
+          add_op4.out(), full_grad_op.out(), false);
+  // backward linear 4
+  paddle::dialect::AddGradOp add_op4_grad =
+      builder.Build<paddle::dialect::AddGradOp>(
+          matmul_op4.out(), full_bias_op4.out(), gelu_op2_grad.x_grad());
+  paddle::dialect::MatmulGradOp matmul_op4_grad =
+      builder.Build<paddle::dialect::MatmulGradOp>(
+          gelu_op1.out(), full_weight_op4.out(), add_op4_grad.x_grad());
+
+  paddle::dialect::GeluGradOp gelu_op1_grad =
+      builder.Build<paddle::dialect::GeluGradOp>(
+          add_op3.out(), matmul_op4_grad.x_grad(), false);
+  // backward linear 3
+  paddle::dialect::AddGradOp add_op3_grad =
+      builder.Build<paddle::dialect::AddGradOp>(
+          matmul_op3.out(), full_bias_op3.out(), gelu_op1_grad.x_grad());
+  paddle::dialect::MatmulGradOp matmul_op3_grad =
+      builder.Build<paddle::dialect::MatmulGradOp>(
+          relu_op.out(), full_weight_op3.out(), add_op3_grad.x_grad());
+
+  paddle::dialect::ReluGradOp relu_op_grad =
+      builder.Build<paddle::dialect::ReluGradOp>(relu_op.out(),
+                                                 matmul_op3_grad.x_grad());
+  // backward linear 2
+  paddle::dialect::AddGradOp add_op2_grad =
+      builder.Build<paddle::dialect::AddGradOp>(
+          matmul_op2.out(), full_bias_op2.out(), relu_op_grad.x_grad());
+  paddle::dialect::MatmulGradOp matmul_op2_grad =
+      builder.Build<paddle::dialect::MatmulGradOp>(
+          add_op1.out(), full_weight_op2.out(), add_op2_grad.x_grad());
+  // backward linear 1
+  paddle::dialect::AddGradOp add_op1_grad =
+      builder.Build<paddle::dialect::AddGradOp>(
+          matmul_op1.out(), full_bias_op1.out(), matmul_op2_grad.x_grad());
+  paddle::dialect::MatmulGradOp matmul_op1_grad =
+      builder.Build<paddle::dialect::MatmulGradOp>(
+          full_input_op1.out(), full_weight_op1.out(), add_op1_grad.x_grad());
+
+  builder.Build<paddle::dialect::FetchOp>(gelu_op2.out(), "out", 0);
+  builder.Build<paddle::dialect::FetchOp>(matmul_op1_grad.x_grad(), "dx", 1);
+}
+
+TEST(DrrTest, FusedLinear) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram(builder);
+
+  EXPECT_EQ(program.block()->size(), 34u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedGemmEpiloguePass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::Unavailable("pm fail to run program"));
+  EXPECT_EQ(program.block()->size(), 22u);
+}

test/cpp/pir/pattern_rewrite/drr_same_type_binding_test.cc

@@ -0,0 +1,324 @@
+// 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 <glog/logging.h>
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
+#include "paddle/fluid/pir/transforms/general/dead_code_elimination_pass.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/pass/pass.h"
+#include "paddle/pir/include/pass/pass_manager.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
+
+/* Source pattern:
+                                       input1
+                                    /  |  \  \  \
+                                  /    |   \   \    \
+             full               /      |    |    \     \           full_tmp
+            /  |        transpose1      | trans2 trans3    \        /   |
+           /   |         /    |        |    |      |        \      /    |
+    softmax1   |        /     |        |    |      |          \   /     |
+         \     |      /    softmax2    |    |      |          add1      |
+           \   |    /             \    |     \    /             |       |
+           layernorm             matmul2     matmul1             \      |
+             / | \                   |         |                  \     |
+           /   |   \                  \       /                     \   |
+         /     |     \                 matmul3                        add2
+        |      |      |                /  |  \                          |
+        |      |      |              /    |    \                        |
+        |      |      |            /      |      \                      |
+        |      |      |         trans4  trans5  trans6                  |
+        |      |      |           |       |        |                    |
+        |      |      |         relu1  softmax3 softmax4              relu2
+        |      |      |           |       |        |                    |
+    output0 output1 output2    output3  output4  output5             output6
+*/
+
+// This class is for test cases of the same type of OP.
+// (without considering the computational logic between OPs,
+// only focusing on the process of matching and replacing)
+class SameTypeBindingTestPattern : public paddle::drr::DrrPatternBase {
+ public:
+  std::string name() const override { return "SameTypeBindingTestPattern"; }
+
+  void operator()(paddle::drr::DrrPatternContext *ctx) const override {
+    paddle::drr::SourcePattern src = ctx->SourcePattern();
+
+    // path 1
+    const auto &transpose_1 =
+        src.Op("pd_op.transpose", {{"perm", src.Attr("perm_1")}});
+    src.Tensor("transpose_1_out") = transpose_1(src.Tensor("input_1"));
+    const auto &softmax_2 =
+        src.Op("pd_op.softmax", {{"axis", src.Attr("softmax_2_axis")}});
+    src.Tensor("softmax_2_out") = softmax_2(src.Tensor("transpose_1_out"));
+    const auto &matmul_2 =
+        src.Op("pd_op.matmul",
+               {{"transpose_x", src.Attr("matmul_2_transpose_x")},
+                {"transpose_y", src.Attr("matmul_2_transpose_y")}});
+    src.Tensor("matmul_2_out") =
+        matmul_2(src.Tensor("softmax_2_out"), src.Tensor("input_1"));
+
+    // path 2
+    const auto &full_1 = src.Op("pd_op.full",
+                                {{"shape", src.Attr("shape_1")},
+                                 {"value", src.Attr("value_1")},
+                                 {"dtype", src.Attr("dtype_1")},
+                                 {"place", src.Attr("place_1")}});
+    src.Tensor("full_1_out") = full_1();
+    const auto &softmax_1 =
+        src.Op("pd_op.softmax", {{"axis", src.Attr("softmax_1_axis")}});
+    src.Tensor("softmax_1_out") = softmax_1(src.Tensor("full_1_out"));
+    const auto &layernorm_1 =
+        src.Op("pd_op.layer_norm",
+               {{"epsilon", src.Attr("layernorm_epsilon")},
+                {"begin_norm_axis", src.Attr("layernorm_begin_norm_axis")}});
+    layernorm_1({&src.Tensor("transpose_1_out"),
+                 &src.Tensor("full_1_out"),
+                 &src.Tensor("softmax_1_out")},
+                {&src.Tensor("output0"),
+                 &src.Tensor("output1"),
+                 &src.Tensor("output2")});
+
+    // path 3
+    const auto &transpose_2 =
+        src.Op("pd_op.transpose", {{"perm", src.Attr("perm_2")}});
+    const auto &transpose_3 =
+        src.Op("pd_op.transpose", {{"perm", src.Attr("perm_3")}});
+    const auto &matmul_1 =
+        src.Op("pd_op.matmul",
+               {{"transpose_x", src.Attr("matmul_1_transpose_x")},
+                {"transpose_y", src.Attr("matmul_1_transpose_y")}});
+    src.Tensor("matmul_1_out") = matmul_1(transpose_2(src.Tensor("input_1")),
+                                          transpose_3(src.Tensor("input_1")));
+    const auto &matmul_3 =
+        src.Op("pd_op.matmul",
+               {{"transpose_x", src.Attr("matmul_3_transpose_x")},
+                {"transpose_y", src.Attr("matmul_3_transpose_y")}});
+    src.Tensor("matmul_3_out") =
+        matmul_3(src.Tensor("matmul_2_out"), src.Tensor("matmul_1_out"));
+    const auto &transpose_4 =
+        src.Op("pd_op.transpose", {{"perm", src.Attr("perm_4")}});
+    const auto &transpose_5 =
+        src.Op("pd_op.transpose", {{"perm", src.Attr("perm_5")}});
+    const auto &transpose_6 =
+        src.Op("pd_op.transpose", {{"perm", src.Attr("perm_6")}});
+    const auto &relu_1 = src.Op("pd_op.relu");
+    const auto &softmax_3 =
+        src.Op("pd_op.softmax", {{"axis", src.Attr("softmax_3_axis")}});
+    const auto &softmax_4 =
+        src.Op("pd_op.softmax", {{"axis", src.Attr("softmax_4_axis")}});
+    src.Tensor("output3") = relu_1(transpose_4(src.Tensor("matmul_3_out")));
+    src.Tensor("output4") = softmax_3(transpose_5(src.Tensor("matmul_3_out")));
+    src.Tensor("output5") = softmax_4(transpose_6(src.Tensor("matmul_3_out")));
+
+    // path 4
+    const auto &full_tmp = src.Op("pd_op.full",
+                                  {{"shape", src.Attr("shape_tmp")},
+                                   {"value", src.Attr("value_tmp")},
+                                   {"dtype", src.Attr("dtype_tmp")},
+                                   {"place", src.Attr("place_tmp")}});
+    src.Tensor("full_tmp_out") = full_tmp();
+    const auto &add_1 = src.Op("pd_op.add");
+    src.Tensor("add_1_out") =
+        add_1(src.Tensor("input_1"), src.Tensor("full_tmp_out"));
+    const auto &add_2 = src.Op("pd_op.add");
+    src.Tensor("add_2_out") =
+        add_2(src.Tensor("add_1_out"), src.Tensor("full_tmp_out"));
+    const auto &relu_2 = src.Op("pd_op.relu");
+    src.Tensor("output6") = relu_2(src.Tensor("add_2_out"));
+
+    paddle::drr::ResultPattern res = src.ResultPattern();
+    const auto &transpose_7 =
+        res.Op("pd_op.transpose", {{"perm", src.Attr("perm_4")}});
+    res.Tensor("output0") = transpose_7(res.Tensor("input_1"));
+    const auto &transpose_8 =
+        res.Op("pd_op.transpose", {{"perm", src.Attr("perm_5")}});
+    res.Tensor("output1") = transpose_8(res.Tensor("input_1"));
+    const auto &full_2 = res.Op("pd_op.full",
+                                {{"shape", src.Attr("shape_tmp")},
+                                 {"value", src.Attr("value_tmp")},
+                                 {"dtype", src.Attr("dtype_tmp")},
+                                 {"place", src.Attr("place_tmp")}});
+    const auto &full_3 = res.Op("pd_op.full",
+                                {{"shape", src.Attr("shape_tmp")},
+                                 {"value", src.Attr("value_tmp")},
+                                 {"dtype", src.Attr("dtype_tmp")},
+                                 {"place", src.Attr("place_tmp")}});
+    const auto &full_4 = res.Op("pd_op.full",
+                                {{"shape", src.Attr("shape_tmp")},
+                                 {"value", src.Attr("value_tmp")},
+                                 {"dtype", src.Attr("dtype_tmp")},
+                                 {"place", src.Attr("place_tmp")}});
+    const auto &full_5 = res.Op("pd_op.full",
+                                {{"shape", src.Attr("shape_tmp")},
+                                 {"value", src.Attr("value_tmp")},
+                                 {"dtype", src.Attr("dtype_tmp")},
+                                 {"place", src.Attr("place_tmp")}});
+    const auto &full_6 = res.Op("pd_op.full",
+                                {{"shape", src.Attr("shape_tmp")},
+                                 {"value", src.Attr("value_tmp")},
+                                 {"dtype", src.Attr("dtype_tmp")},
+                                 {"place", src.Attr("place_tmp")}});
+    res.Tensor("output2") = full_2();
+    res.Tensor("output3") = full_3();
+    res.Tensor("output4") = full_4();
+    res.Tensor("output5") = full_5();
+    res.Tensor("output6") = full_6();
+  }
+};
+
+void BuildProgram(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{4, 3, 16},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  // path 1
+  paddle::dialect::TransposeOp transpose_op1 =
+      builder.Build<paddle::dialect::TransposeOp>(full_input_op1.out(),
+                                                  std::vector<int>{0, 1, 2});
+
+  paddle::dialect::SoftmaxOp softmax_op2 =
+      builder.Build<paddle::dialect::SoftmaxOp>(transpose_op1.out(), -1);
+
+  paddle::dialect::TransposeOp matmul_rhs_trans =
+      builder.Build<paddle::dialect::TransposeOp>(full_input_op1.out(),
+                                                  std::vector<int>{0, 2, 1});
+
+  paddle::dialect::MatmulOp matmul_op2 =
+      builder.Build<paddle::dialect::MatmulOp>(softmax_op2.out(),
+                                               matmul_rhs_trans.out());
+
+  // path 2
+  paddle::dialect::FullOp full_op_scale =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{16},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+  paddle::dialect::SoftmaxOp softmax_op_bias =
+      builder.Build<paddle::dialect::SoftmaxOp>(full_op_scale.out(), -1);
+  paddle::dialect::LayerNormOp layernorm_op1 =
+      builder.Build<paddle::dialect::LayerNormOp>(
+          transpose_op1.out(), full_op_scale.out(), softmax_op_bias.out());
+
+  // path 3
+  paddle::dialect::TransposeOp transpose_op2 =
+      builder.Build<paddle::dialect::TransposeOp>(full_input_op1.out(),
+                                                  std::vector<int>{0, 1, 2});
+
+  paddle::dialect::TransposeOp transpose_op3 =
+      builder.Build<paddle::dialect::TransposeOp>(full_input_op1.out(),
+                                                  std::vector<int>{0, 2, 1});
+
+  paddle::dialect::MatmulOp matmul_op1 =
+      builder.Build<paddle::dialect::MatmulOp>(transpose_op2.out(),
+                                               transpose_op3.out());
+
+  paddle::dialect::MatmulOp matmul_op3 =
+      builder.Build<paddle::dialect::MatmulOp>(matmul_op2.out(),
+                                               matmul_op1.out());
+
+  paddle::dialect::TransposeOp transpose_op4 =
+      builder.Build<paddle::dialect::TransposeOp>(matmul_op3.out(),
+                                                  std::vector<int>{0, 1, 2});
+
+  paddle::dialect::ReluOp relu_op1 =
+      builder.Build<paddle::dialect::ReluOp>(transpose_op4.out());
+
+  paddle::dialect::TransposeOp transpose_op5 =
+      builder.Build<paddle::dialect::TransposeOp>(matmul_op3.out(),
+                                                  std::vector<int>{0, 1, 2});
+
+  paddle::dialect::SoftmaxOp softmax_op3 =
+      builder.Build<paddle::dialect::SoftmaxOp>(transpose_op5.out(), -1);
+
+  paddle::dialect::TransposeOp transpose_op6 =
+      builder.Build<paddle::dialect::TransposeOp>(matmul_op3.out(),
+                                                  std::vector<int>{0, 1, 2});
+
+  paddle::dialect::SoftmaxOp softmax_op4 =
+      builder.Build<paddle::dialect::SoftmaxOp>(transpose_op6.out(), -1);
+
+  // path 4
+  paddle::dialect::FullOp full_input_op2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{4, 3, 16},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::AddOp add_op1 = builder.Build<paddle::dialect::AddOp>(
+      full_input_op1.out(), full_input_op2.out());
+
+  paddle::dialect::AddOp add_op2 = builder.Build<paddle::dialect::AddOp>(
+      add_op1.out(), full_input_op2.out());
+
+  paddle::dialect::ReluOp relu_op2 =
+      builder.Build<paddle::dialect::ReluOp>(add_op2.out());
+
+  // tail
+  paddle::dialect::AddOp safe_op1 = builder.Build<paddle::dialect::AddOp>(
+      layernorm_op1.variance(), layernorm_op1.mean());
+
+  paddle::dialect::MatmulOp matmul_op5 =
+      builder.Build<paddle::dialect::MatmulOp>(relu_op1.out(),
+                                               softmax_op3.out());
+
+  paddle::dialect::MatmulOp matmul_op6 =
+      builder.Build<paddle::dialect::MatmulOp>(softmax_op4.out(),
+                                               relu_op2.out());
+
+  builder.Build<paddle::dialect::FetchOp>(safe_op1.out(), "out1", 0);
+  builder.Build<paddle::dialect::FetchOp>(matmul_op5.out(), "out2", 1);
+  builder.Build<paddle::dialect::FetchOp>(matmul_op6.out(), "out3", 2);
+}
+
+class DrrPatternRewritePass : public pir::PatternRewritePass {
+ public:
+  DrrPatternRewritePass()
+      : pir::PatternRewritePass("drr_pattern_rewrite_pass", 1) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add(paddle::drr::Create<SameTypeBindingTestPattern>(context));
+
+    return ps;
+  }
+};
+
+TEST(DrrTest, drr_demo) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram(builder);
+
+  EXPECT_EQ(program.block()->size(), 28u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(std::make_unique<DrrPatternRewritePass>());
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::Unavailable("pm fail to run program"));
+  EXPECT_EQ(program.block()->size(), 28u);
+}

test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc

@@ -0,0 +1,644 @@
+// 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 <gtest/gtest.h>
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <numeric>
+#include <sstream>
+#include <vector>
+
+#include "paddle/common/enforce.h"
+#include "paddle/common/errors.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/general/constant_folding_pass.h"
+#include "paddle/fluid/pir/transforms/general/dead_code_elimination_pass.h"
+#include "paddle/fluid/pir/transforms/gpu/conv2d_add_act_fuse_pass.h"
+#include "paddle/fluid/pir/transforms/gpu/conv2d_add_fuse_pass.h"
+#include "paddle/fluid/pir/transforms/gpu/conv2d_bn_fuse_pass.h"
+#include "paddle/fluid/pir/utils/general_functions.h"
+#include "paddle/pir/include/core/builder.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/core/builtin_op.h"
+#include "paddle/pir/include/core/cast_utils.h"
+#include "paddle/pir/include/core/dialect.h"
+#include "paddle/pir/include/core/ir_context.h"
+#include "paddle/pir/include/core/op_info.h"
+#include "paddle/pir/include/core/parameter.h"
+#include "paddle/pir/include/core/program.h"
+#include "paddle/pir/include/core/value.h"
+#include "paddle/pir/include/pass/pass.h"
+#include "paddle/pir/include/pass/pass_manager.h"
+#include "paddle/pir/include/pattern_rewrite/frozen_rewrite_pattern_set.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_applicator.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_match.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "test/cpp/pir/tools/macros_utils.h"
+
+PD_DECLARE_KERNEL(full, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(add, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(sqrt, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(divide, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(multiply, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(subtract, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(full_int_array, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(reshape, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(fetch, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(conv2d, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(transpose, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(cummax, CPU, ALL_LAYOUT);
+
+// Define op1.
+class Operation1 : public pir::Op<Operation1> {
+ public:
+  using Op::Op;
+  static const char *name() { return "test.Operation1"; }
+  static constexpr uint32_t attributes_num = 2;
+  static const char *attributes_name[attributes_num];  // NOLINT
+  void VerifySig();
+  static void InferShape() { VLOG(2) << "This is op2's InferShape interface."; }
+};
+
+void Operation1::VerifySig() {
+  auto &attributes = this->attributes();
+  if (attributes.count("op2_attr1") == 0 ||
+      (!attributes.at("op2_attr1").isa<pir::StrAttribute>())) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Type of attribute: parameter_name is not right."));
+  }
+  if (attributes.count("op2_attr2") == 0 ||
+      (!attributes.at("op2_attr2").isa<pir::StrAttribute>())) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Type of attribute: parameter_name is not right."));
+  }
+}
+const char *Operation1::attributes_name[attributes_num] = {  // NOLINT
+    "op2_attr1",
+    "op2_attr2"};
+IR_DECLARE_EXPLICIT_TEST_TYPE_ID(Operation1)
+IR_DEFINE_EXPLICIT_TYPE_ID(Operation1)
+
+// Define a dialect, op1 and op2 will be registered by this dialect.
+class TestDialect : public pir::Dialect {
+ public:
+  explicit TestDialect(pir::IrContext *context)
+      : pir::Dialect(name(), context, pir::TypeId::get<TestDialect>()) {
+    initialize();
+  }
+  static const char *name() { return "test"; }
+
+ private:
+  void initialize() { RegisterOps<Operation1>(); }
+};
+IR_DECLARE_EXPLICIT_TEST_TYPE_ID(TestDialect)
+IR_DEFINE_EXPLICIT_TYPE_ID(TestDialect)
+
+// TODO(wilber): Add logical when ir support erase, replace or update.
+class TestPatternRewrite : public pir::OpRewritePattern<Operation1> {
+ public:
+  using pir::OpRewritePattern<Operation1>::OpRewritePattern;
+
+  void Rewrite(Operation1 op, pir::PatternRewriter &rewriter) const override {}
+  bool Match(Operation1 op) const override { return false; }
+};
+
+class TestPatternRewrite2 : public pir::OpRewritePattern<Operation1> {
+ public:
+  using pir::OpRewritePattern<Operation1>::OpRewritePattern;
+  bool MatchAndRewrite(
+      Operation1 op,
+      pir::PatternRewriter &rewriter) const override {  // NOLINT
+    return false;
+  }
+};
+
+TEST(PatternRewrite, PatternBenefit) {
+  pir::PatternBenefit benefit1(1);
+  EXPECT_EQ(benefit1.benefit(), 1U);
+  pir::PatternBenefit benefit2(2);
+  EXPECT_EQ(benefit2.benefit(), 2U);
+
+  EXPECT_TRUE(benefit2 > benefit1);
+  EXPECT_TRUE(benefit2 >= benefit1);
+  EXPECT_TRUE(benefit1 < benefit2);
+  EXPECT_TRUE(benefit1 <= benefit2);
+  EXPECT_TRUE(benefit1 != benefit2);
+  pir::PatternBenefit benefit3(2);
+  EXPECT_TRUE(benefit2 == benefit3);
+}
+
+TEST(RewritePattern, RewritePatternSet) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  auto *test_dialect = ctx->GetOrRegisterDialect<TestDialect>();
+  test_dialect->RegisterOp<Operation1>();
+
+  pir::RewritePatternSet ps(ctx);
+  ps.Add<TestPatternRewrite>(ctx, 1);
+  EXPECT_EQ(ps.native_patterns().size(), 1U);
+  EXPECT_TRUE(ps.native_patterns().back()->debug_labels().empty());
+  EXPECT_EQ(ps.native_patterns().back()->benefit(), 1U);
+  ps.AddWithLabel<TestPatternRewrite2>({"TestPatternRewrite2"}, ctx, 2);
+  EXPECT_EQ(ps.native_patterns().size(), 2U);
+  EXPECT_EQ(ps.native_patterns().back()->debug_labels()[0],
+            "TestPatternRewrite2");
+  EXPECT_EQ(ps.native_patterns().back()->benefit(), 2U);
+
+  ps.Clear();
+  ps.Add<TestPatternRewrite, TestPatternRewrite2>(ctx, 2);
+  EXPECT_EQ(ps.native_patterns().size(), 2U);
+  EXPECT_EQ(ps.native_patterns()[0]->benefit(), 2U);
+  EXPECT_EQ(ps.native_patterns()[1]->benefit(), 2U);
+}
+
+// TODO(wilber): Add actual case.
+// TEST(PatternRewrite, PatternApplicator) {
+//   pir::IrContext *ctx = pir::IrContext::Instance();
+//   ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+//   auto *test_dialect = ctx->GetOrRegisterDialect<TestDialect>();
+//   test_dialect->RegisterOp<Operation1>();
+//   pir::RewritePatternSet ps(ctx);
+//   ps.Add<TestPatternRewrite, TestPatternRewrite2>(ctx, 2);
+//   pir::FrozenRewritePatternSet frozen_set(std::move(ps));
+//   pir::PatternApplicator applicator(frozen_set);
+//   applicator.ApplyDefaultCostModel();
+// }
+
+// // TODO(wilber): Add actual case.
+TEST(PatternRewrite, FrozenRewritePatternSet) {
+  pir::FrozenRewritePatternSet frozen_set;
+  EXPECT_TRUE(frozen_set.match_any_op_native_patterns().empty());
+  EXPECT_TRUE(frozen_set.op_specific_native_patterns().empty());
+
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  auto *test_dialect = ctx->GetOrRegisterDialect<TestDialect>();
+  test_dialect->RegisterOp<Operation1>();
+  pir::RewritePatternSet ps(ctx);
+  ps.Add<TestPatternRewrite, TestPatternRewrite2>(ctx, 2);
+
+  pir::FrozenRewritePatternSet frozen_set2(std::move(ps));
+  EXPECT_TRUE(frozen_set2.match_any_op_native_patterns().empty());
+  const auto &pattern_maps = frozen_set2.op_specific_native_patterns();
+  EXPECT_EQ(pattern_maps.size(), 1U);
+  EXPECT_EQ(pattern_maps.at(ctx->GetRegisteredOpInfo("test.Operation1")).size(),
+            2U);
+}
+
+class RedundantTransposeFusePattern
+    : public pir::OpRewritePattern<paddle::dialect::TransposeOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::TransposeOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::TransposeOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    auto prev_op = pir::GetDefiningOpForInput(op, 0);
+    std::vector<int> axis_last = GetAxis(op);
+    auto prev_trans_op = prev_op->dyn_cast<paddle::dialect::TransposeOp>();
+    if (prev_trans_op) {
+      std::vector<int> axis_first = GetAxis(prev_trans_op);
+      PADDLE_ENFORCE_EQ(axis_first.size(),
+                        axis_last.size(),
+                        common::errors::InvalidArgument(
+                            "transpose op's perm rank should be same."));
+      auto new_perm = GetPerm(axis_first, axis_last);
+      rewriter.set_insertion_point(op);
+      auto new_transpose_op = rewriter.Build<paddle::dialect::TransposeOp>(
+          pir::GetDefiningOpForInput(prev_trans_op, 0)->result(0), new_perm);
+      rewriter.ReplaceOp(op, {new_transpose_op.out()});
+      return true;
+    }
+
+    return false;
+  }
+
+ private:
+  std::vector<int> GetAxis(paddle::dialect::TransposeOp op) const {
+    auto array_attr = op.attribute<pir::ArrayAttribute>("perm").AsVector();
+    std::vector<int> axis(array_attr.size());
+    for (size_t i = 0; i < array_attr.size(); ++i) {
+      axis[i] = array_attr[i].dyn_cast<pir::Int32Attribute>().data();
+    }
+    return axis;
+  }
+
+  std::vector<int> GetPerm(const std::vector<int> &perm1,
+                           const std::vector<int> &perm2) const {
+    int n = static_cast<int>(perm1.size());
+    std::vector<int> axis(n), axis1(n), axis2(n);
+    std::iota(axis.begin(), axis.end(), 0);
+    for (int i = 0; i < n; ++i) {
+      axis1[i] = axis[perm1[i]];
+    }
+    for (int i = 0; i < n; ++i) {
+      axis2[i] = axis1[perm2[i]];
+    }
+    return axis2;
+  }
+};
+
+class TestPass : public pir::PatternRewritePass {
+ public:
+  TestPass() : pir::PatternRewritePass("test_pass", 1) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add<RedundantTransposeFusePattern>(context);
+    return ps;
+  }
+};
+
+void BuildProgram(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp full_input_op =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{4, 3, 16, 16},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_filter_op =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 3, 3, 3},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_input_op_1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{4, 3, 16, 16},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_filter_op_1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 3, 3, 3},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp add_op_y =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{1, 64, 1, 1},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_input_op_2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{4, 3, 16, 16},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_filter_op_2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 3, 3, 3},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp add_op_y_1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{1, 64, 1, 1},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp add_op_2_x = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{4, 64, 14, 14},
+      1.5,
+      phi::DataType::FLOAT32,
+      phi::CPUPlace());
+
+  paddle::dialect::FullOp full_mean_op = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{64}, 1.5, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  paddle::dialect::FullOp full_variance_op =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_scale_op =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64},
+                                             1.5,
+                                             phi::DataType::FLOAT32,
+                                             phi::CPUPlace());
+
+  paddle::dialect::FullOp full_bias_op = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{64}, 1.5, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  paddle::dialect::Conv2dOp conv2d_op =
+      builder.Build<paddle::dialect::Conv2dOp>(full_input_op.out(),
+                                               full_filter_op.out());
+
+  paddle::dialect::BatchNorm_Op batch_norm_op =
+      builder.Build<paddle::dialect::BatchNorm_Op>(conv2d_op.out(),
+                                                   full_mean_op.out(),
+                                                   full_variance_op.out(),
+                                                   full_scale_op.out(),
+                                                   full_bias_op.out(),
+                                                   true,
+                                                   0.9,
+                                                   1e-6,
+                                                   "NCHW",
+                                                   false,
+                                                   false);
+  paddle::dialect::Conv2dOp conv2d_op_1 =
+      builder.Build<paddle::dialect::Conv2dOp>(full_input_op_1.out(),
+                                               full_filter_op_1.out());
+  paddle::dialect::AddOp add_op =
+      builder.Build<paddle::dialect::AddOp>(conv2d_op_1.out(), add_op_y.out());
+  paddle::dialect::ReluOp relu_op =
+      builder.Build<paddle::dialect::ReluOp>(add_op.out());
+
+  paddle::dialect::Conv2dOp conv2d_op_2 =
+      builder.Build<paddle::dialect::Conv2dOp>(full_input_op_2.out(),
+                                               full_filter_op_2.out());
+  paddle::dialect::AddOp add_op_1 = builder.Build<paddle::dialect::AddOp>(
+      conv2d_op_2.out(), add_op_y_1.out());
+  paddle::dialect::AddOp add_op_2 =
+      builder.Build<paddle::dialect::AddOp>(add_op_2_x.out(), add_op_1.out());
+  paddle::dialect::ReluOp relu_op_1 =
+      builder.Build<paddle::dialect::ReluOp>(add_op_2.out());
+
+  auto transpose1_op = builder.Build<paddle::dialect::TransposeOp>(
+      batch_norm_op.out(), std::vector<int>{0, 2, 3, 1});
+
+  auto transpose2_op = builder.Build<paddle::dialect::TransposeOp>(
+      transpose1_op.out(), std::vector<int>{0, 3, 1, 2});
+  auto add_out = builder.Build<paddle::dialect::AddOp>(transpose2_op.out(),
+                                                       relu_op_1.out());
+  auto add_out_1 =
+      builder.Build<paddle::dialect::AddOp>(add_out.out(), relu_op.out());
+  builder.Build<paddle::dialect::FetchOp>(add_out_1.out(), "out", 0);
+}
+
+TEST(pattern_rewrite, Patterns) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram(builder);
+
+  EXPECT_EQ(program.block()->size(), 27u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(std::make_unique<TestPass>());
+  pm.AddPass(pir::CreateConv2dBnFusePass());
+  pm.AddPass(pir::CreateConv2dAddActFusePass());
+  pm.AddPass(pir::CreateConv2dAddFusePass());
+  std::unique_ptr<pir::Pass> constant_folding_pass =
+      pir::CreateConstantFoldingPass();
+  phi::Place place = phi::CPUPlace();
+  constant_folding_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place);
+  constant_folding_pass->Set(pir::Pass::kParamScopeAttr,
+                             new paddle::framework::Scope());
+  pm.AddPass(std::move(constant_folding_pass));
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  // pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+  // pm.EnableIRPrinting(std::make_unique<pir::PassManager::IRPrinterOption>(
+  //     [](pir::Pass *pass, pir::Operation *op) {
+  //       return pass->name() == "constant_folding_pass";
+  //     },
+  //     [](pir::Pass *pass, pir::Operation *op) {
+  //       return pass->name() == "constant_folding_pass";
+  //     },
+  //     true,
+  //     true));
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::ExecutionTimeout(
+                        "Sorry, the for program test is timeout."));
+  EXPECT_EQ(program.block()->size(), 17u);
+}
+
+void BuildConstantFoldingProgram(pir::Program *program,
+                                 pir::IrContext *ctx,
+                                 paddle::framework::Scope *scope) {
+  pir::Builder builder = pir::Builder(ctx, program->block());
+
+  pir::Type fp32_dtype = pir::Float32Type::get(ctx);
+  phi::DDim dims = {2, 2};
+  phi::DataLayout data_layout = phi::DataLayout::NCHW;
+  phi::LegacyLoD lod = {{0, 1, 2}};
+  size_t offset = 0;
+  pir::Type dense_tensor_dtype = paddle::dialect::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+
+  phi::DenseTensorMeta meta(
+      phi::DataType::FLOAT32, dims, data_layout, lod, offset);
+  phi::DeviceContext *dev_ctx =
+      phi::DeviceContextPool::Instance().Get(phi::CPUPlace());
+
+  auto op1 = builder.Build<pir::ConstantTensorOp>("a", dense_tensor_dtype);
+  auto op2 = builder.Build<pir::ConstantTensorOp>("b", dense_tensor_dtype);
+
+  auto op3 =
+      builder.Build<paddle::dialect::AddOp>(op1->result(0), op2->result(0));
+
+  auto op4 = builder.Build<pir::ParameterOp>("c", dense_tensor_dtype);
+
+  auto op5 =
+      builder.Build<paddle::dialect::AddOp>(op3->result(0), op4->result(0));
+  builder.Build<paddle::dialect::FetchOp>(op5.out(), "out", 0);
+
+  auto *tensor_a = scope->Var("a")->GetMutable<phi::DenseTensor>();
+  auto *tensor_b = scope->Var("b")->GetMutable<phi::DenseTensor>();
+  auto *tensor_c = scope->Var("c")->GetMutable<phi::DenseTensor>();
+
+  tensor_a->set_meta(meta);
+  tensor_b->set_meta(meta);
+  tensor_c->set_meta(meta);
+
+  dev_ctx->Alloc(tensor_a, phi::DataType::FLOAT32);
+  dev_ctx->Alloc(tensor_b, phi::DataType::FLOAT32);
+  dev_ctx->Alloc(tensor_c, phi::DataType::FLOAT32);
+}
+
+TEST(constant_folding, ConstantFolding) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+
+  pir::Program program(ctx);
+  paddle::framework::Scope scope;
+  BuildConstantFoldingProgram(&program, ctx, &scope);
+
+  pir::PassManager pm(ctx);
+  std::unique_ptr<pir::Pass> constant_folding_pass =
+      pir::CreateConstantFoldingPass();
+  phi::Place place = phi::CPUPlace();
+  constant_folding_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place);
+  constant_folding_pass->SetNotOwned(pir::Pass::kParamScopeAttr, &scope);
+  pm.AddPass(std::move(constant_folding_pass));
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::ExecutionTimeout(
+                        "Sorry, the test for program is timeout."));
+  EXPECT_EQ(program.block()->size(), 2u);
+}
+
+TEST(constant_folding, ConstantFolding_Train) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+
+  pir::Program program(ctx);
+  paddle::framework::Scope scope;
+  BuildConstantFoldingProgram(&program, ctx, &scope);
+
+  pir::PassManager pm(ctx);
+  std::unique_ptr<pir::Pass> constant_folding_pass =
+      pir::CreateConstantFoldingPass();
+  phi::Place place = phi::CPUPlace();
+  constant_folding_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place);
+  constant_folding_pass->SetNotOwned(pir::Pass::kParamScopeAttr, &scope);
+  constant_folding_pass->Set("train_mode", new bool(true));
+
+  pm.AddPass(std::move(constant_folding_pass));
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::ExecutionTimeout(
+                        "Sorry, the test for program is timeout."));
+  EXPECT_EQ(program.block()->size(), 4u);
+}
+
+void BuildConcatProgram(pir::Program *program, pir::IrContext *ctx) {
+  pir::Builder builder = pir::Builder(ctx, program->block());
+  auto x = builder
+               .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16, 16}),
+                                               2.0,
+                                               phi::DataType::FLOAT32,
+                                               phi::GPUPlace())
+               .result(0);
+
+  auto y = builder
+               .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16, 16}),
+                                               2.0,
+                                               phi::DataType::FLOAT32,
+                                               phi::GPUPlace())
+               .result(0);
+
+  auto t1 =
+      builder.Build<pir::CombineOp>(std::vector<pir::Value>({x, y})).result(0);
+
+  auto out1 = builder.Build<paddle::dialect::ConcatOp>(t1, 1).result(0);
+
+  auto z = builder
+               .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16, 16}),
+                                               2.0,
+                                               phi::DataType::FLOAT32,
+                                               phi::GPUPlace())
+               .result(0);
+
+  auto w = builder
+               .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16, 16}),
+                                               2.0,
+                                               phi::DataType::FLOAT32,
+                                               phi::GPUPlace())
+               .result(0);
+
+  auto t2 =
+      builder.Build<pir::CombineOp>(std::vector<pir::Value>({z, w})).result(0);
+
+  auto out2 = builder.Build<paddle::dialect::ConcatOp>(t2, 1).result(0);
+
+  auto out = builder.Build<paddle::dialect::AddOp>(out1, out2).result(0);
+
+  builder.Build<paddle::dialect::FetchOp>(out, "out", 0);
+}
+
+TEST(constant_folding, ConstantFolding_Combine) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+
+  pir::Program program(ctx);
+  BuildConcatProgram(&program, ctx);
+
+  pir::PassManager pm(ctx);
+  std::unique_ptr<pir::Pass> constant_folding_pass =
+      pir::CreateConstantFoldingPass();
+  phi::Place place = phi::CPUPlace();
+  constant_folding_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place);
+  constant_folding_pass->Set(pir::Pass::kParamScopeAttr,
+                             new paddle::framework::Scope());
+  pm.AddPass(std::move(constant_folding_pass));
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::ExecutionTimeout(
+                        "Sorry, the test for program is timeout."));
+  EXPECT_EQ(program.block()->size(), 2u);
+}
+
+void BuildMultiOutputProgram(pir::Program *program, pir::IrContext *ctx) {
+  pir::Builder builder = pir::Builder(ctx, program->block());
+  auto x = builder
+               .Build<paddle::dialect::FullOp>(std::vector<int64_t>({2, 2}),
+                                               2.0,
+                                               phi::DataType::FLOAT32,
+                                               phi::GPUPlace())
+               .result(0);
+
+  auto cummax_op = builder.Build<paddle::dialect::CummaxOp>(x, 0);
+
+  auto out = cummax_op.out();
+  auto indices = cummax_op.indices();
+
+  builder.Build<paddle::dialect::FetchOp>(out, "out", 0);
+  builder.Build<paddle::dialect::FetchOp>(indices, "indices", 1);
+}
+
+TEST(constant_folding, ConstantFolding_MultiOutput) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+
+  pir::Program program(ctx);
+  BuildMultiOutputProgram(&program, ctx);
+
+  pir::PassManager pm(ctx);
+  std::unique_ptr<pir::Pass> constant_folding_pass =
+      pir::CreateConstantFoldingPass();
+  phi::Place place = phi::CPUPlace();
+  constant_folding_pass->SetNotOwned(pir::Pass::kPlaceAttr, &place);
+  constant_folding_pass->Set(pir::Pass::kParamScopeAttr,
+                             new paddle::framework::Scope());
+  pm.AddPass(std::move(constant_folding_pass));
+  pm.AddPass(pir::CreateDeadCodeEliminationPass());
+  pm.EnableIRPrinting();
+
+  PADDLE_ENFORCE_EQ(pm.Run(&program),
+                    true,
+                    common::errors::ExecutionTimeout(
+                        "Sorry, the test for program is timeout."));
+  EXPECT_EQ(program.block()->size(), 4u);
+}