chore: import upstream snapshot with attribution

test/cpp/prim/CMakeLists.txt

@@ -0,0 +1,39 @@
+set(prim_eager_deps
+    phi
+    common
+    hook_utils
+    utils
+    global_utils
+    backward
+    tracer
+    layer
+    autograd_meta
+    eager_nan_inf_utils
+    grad_node_info
+    grad_tensor_holder
+    custom_operator_node)
+
+set(prim_generated_deps final_dygraph_function final_dygraph_node
+                        dygraph_function dygraph_node)
+
+if(NOT (NOT WITH_PYTHON AND ON_INFER))
+  if(WITH_CINN)
+    set(prim_eager_deps ${prim_eager_deps})
+  endif()
+  cc_library(init_env_utils SRCS init_env_utils.cc)
+  target_compile_definitions(init_env_utils PUBLIC PADDLE_DLL_EXPORT)
+
+  paddle_test(test_comp_eager SRCS test_eager_prim.cc init_env_utils.cc)
+endif()
+
+# skip win32 since wget is not installed by default on windows machine.
+
+if(NOT WIN32)
+  paddle_test(test_vjp_pir SRCS test_vjp.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(test_comp_eager)
+endif()

test/cpp/prim/init_env_utils.cc

@@ -0,0 +1,31 @@
+// 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 "test/cpp/prim/init_env_utils.h"
+#include "paddle/fluid/prim/utils/eager/eager_tensor_operants.h"
+#include "paddle/fluid/prim/utils/static/static_tensor_operants.h"
+#include "paddle/phi/api/include/operants_manager.h"
+
+namespace paddle {
+namespace prim {
+
+void InitTensorOperants() {
+  paddle::OperantsManager::Instance().eager_operants =
+      std::make_unique<paddle::prim::EagerTensorOperants>();
+  paddle::OperantsManager::Instance().static_operants =
+      std::make_unique<paddle::prim::StaticTensorOperants>();
+}
+
+}  // namespace prim
+}  // namespace paddle

test/cpp/prim/init_env_utils.h

@@ -0,0 +1,23 @@
+// 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
+
+namespace paddle {
+namespace prim {
+
+void InitTensorOperants();
+
+}  // namespace prim
+}  // namespace paddle

test/cpp/prim/test_eager_prim.cc

@@ -0,0 +1,182 @@
+// 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.
+
+#include <sstream>
+
+#include "glog/logging.h"
+#include "gtest/gtest.h"
+#include "paddle/common/flags.h"
+#include "paddle/fluid/eager/api/generated/eager_generated/forwards/dygraph_functions.h"
+#include "paddle/fluid/eager/api/utils/hook_utils.h"
+#include "paddle/fluid/eager/backward.h"
+#include "paddle/fluid/prim/utils/utils.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "test/cpp/eager/test_utils.h"
+#include "test/cpp/prim/init_env_utils.h"
+
+COMMON_DECLARE_string(tensor_operants_mode);
+
+namespace paddle {
+namespace prim {
+
+TEST(EagerPrim, TanhBackwardTest) {
+  // 1. Initialized
+  eager_test::InitEnv(phi::CPUPlace());
+  FLAGS_tensor_operants_mode = "eager";
+  paddle::prim::InitTensorOperants();
+  // 2. pre
+  phi::DDim ddim = common::make_ddim({4, 16, 16, 32});
+  paddle::Tensor tensor0 =
+      eager_test::CreateTensorWithValue(ddim,
+                                        phi::CPUPlace(),
+                                        phi::DataType::FLOAT32,
+                                        phi::DataLayout::NCHW,
+                                        5.0 /*value*/,
+                                        true /*is_leaf*/);
+  ::egr::egr_utils_api::RetainGradForTensor(tensor0);
+  paddle::Tensor tensor1 =
+      eager_test::CreateTensorWithValue(ddim,
+                                        phi::CPUPlace(),
+                                        phi::DataType::FLOAT32,
+                                        phi::DataLayout::NCHW,
+                                        5.0 /*value*/,
+                                        true /*is_leaf*/);
+  ::egr::egr_utils_api::RetainGradForTensor(tensor1);
+  // 3. Run Forward once
+  paddle::Tensor out0 = tanh_ad_func(tensor0);
+  std::vector<paddle::Tensor> outs0 = {out0};
+  // Disable prim
+  PrimCommonUtils::SetBwdPrimEnabled(false);
+  ASSERT_FALSE(PrimCommonUtils::IsBwdPrimEnabled());
+  // 4. Run Backward
+  egr::Backward(outs0, {}, false);
+
+  paddle::Tensor out1 = tanh_ad_func(tensor1);
+  std::vector<paddle::Tensor> outs1 = {out1};
+  // Enable prim
+  PrimCommonUtils::SetBwdPrimEnabled(true);
+  ASSERT_TRUE(PrimCommonUtils::IsBwdPrimEnabled());
+  // 4. Run Backward
+  ::egr::Backward(outs1, {}, false);
+  VLOG(7)
+      << "Target Grad is: "
+      << std::static_pointer_cast<phi::DenseTensor>(
+             ::egr::EagerUtils::unsafe_autograd_meta(tensor0)->Grad().impl())
+             ->data<float>()[0];
+  VLOG(7)
+      << "Result Grad is: "
+      << std::static_pointer_cast<phi::DenseTensor>(
+             ::egr::EagerUtils::unsafe_autograd_meta(tensor1)->Grad().impl())
+             ->data<float>()[0];
+  // Examine Backward Grad
+  eager_test::CompareGradTensorWithValue<float>(
+      tensor1,
+      std::static_pointer_cast<phi::DenseTensor>(
+          ::egr::EagerUtils::unsafe_autograd_meta(tensor0)->Grad().impl())
+          ->data<float>()[0]);
+}
+
+TEST(EagerPrim, LogicalOperantsTest) {
+  // 1. Initialized
+  eager_test::InitEnv(phi::CPUPlace());
+  FLAGS_tensor_operants_mode = "eager";
+  paddle::prim::InitTensorOperants();
+  // 2. pre
+  phi::DDim ddim = common::make_ddim({4, 16, 16, 32});
+  paddle::Tensor tensor0 =
+      eager_test::CreateTensorWithValue(ddim,
+                                        phi::CPUPlace(),
+                                        phi::DataType::INT32,
+                                        phi::DataLayout::NCHW,
+                                        1 /*value*/,
+                                        true /*is_leaf*/);
+  ::egr::egr_utils_api::RetainGradForTensor(tensor0);
+  paddle::Tensor tensor1 =
+      eager_test::CreateTensorWithValue(ddim,
+                                        phi::CPUPlace(),
+                                        phi::DataType::INT32,
+                                        phi::DataLayout::NCHW,
+                                        0 /*value*/,
+                                        true /*is_leaf*/);
+  ::egr::egr_utils_api::RetainGradForTensor(tensor1);
+  // 3. Run Forward once
+  paddle::Tensor out0 = tensor0 & tensor1;
+  paddle::Tensor out1 = bitwise_and_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<int>()[0], out1.data<int>()[0]);
+  out0 = tensor0 | tensor1;
+  out1 = bitwise_or_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<int>()[0], out1.data<int>()[0]);
+  out0 = tensor0 ^ tensor1;
+  out1 = bitwise_xor_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<int>()[0], out1.data<int>()[0]);
+  out0 = ~tensor0;
+  out1 = bitwise_not_ad_func(tensor0);
+  EXPECT_EQ(out0.data<int>()[0], out1.data<int>()[0]);
+}
+
+TEST(EagerPrim, CompareOperantsTest) {
+  // 1. Initialized
+  eager_test::InitEnv(phi::CPUPlace());
+  FLAGS_tensor_operants_mode = "eager";
+  paddle::prim::InitTensorOperants();
+  // 2. pre
+  phi::DDim ddim = common::make_ddim({4, 16, 16, 32});
+  paddle::Tensor tensor0 =
+      eager_test::CreateTensorWithValue(ddim,
+                                        phi::CPUPlace(),
+                                        phi::DataType::INT32,
+                                        phi::DataLayout::NCHW,
+                                        1 /*value*/,
+                                        true /*is_leaf*/);
+  ::egr::egr_utils_api::RetainGradForTensor(tensor0);
+  paddle::Tensor tensor1 =
+      eager_test::CreateTensorWithValue(ddim,
+                                        phi::CPUPlace(),
+                                        phi::DataType::INT32,
+                                        phi::DataLayout::NCHW,
+                                        0 /*value*/,
+                                        true /*is_leaf*/);
+  ::egr::egr_utils_api::RetainGradForTensor(tensor1);
+  // 3. Run Forward once
+  paddle::Tensor out0 = (tensor0 < tensor1);
+  paddle::Tensor out1 = less_than_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<bool>()[0], out1.data<bool>()[0]);
+  out0 = (tensor0 <= tensor1);
+  out1 = less_equal_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<bool>()[0], out1.data<bool>()[0]);
+  out0 = (tensor0 == tensor1);
+  out1 = equal_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<bool>()[0], out1.data<bool>()[0]);
+  out0 = (tensor0 != tensor1);
+  out1 = not_equal_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<bool>()[0], out1.data<bool>()[0]);
+  out0 = (tensor0 > tensor1);
+  out1 = greater_than_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<bool>()[0], out1.data<bool>()[0]);
+  out0 = (tensor0 >= tensor1);
+  out1 = greater_equal_ad_func(tensor0, tensor1);
+  EXPECT_EQ(out0.data<bool>()[0], out1.data<bool>()[0]);
+}
+
+TEST(EagerPrim, TestFlags) {
+  PrimCommonUtils::SetBwdPrimEnabled(true);
+  ASSERT_TRUE(PrimCommonUtils::IsBwdPrimEnabled());
+  PrimCommonUtils::SetBwdPrimEnabled(false);
+  ASSERT_FALSE(PrimCommonUtils::IsBwdPrimEnabled());
+}
+
+}  // namespace prim
+}  // namespace paddle

test/cpp/prim/test_vjp.cc

@@ -0,0 +1,505 @@
+// 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 "paddle/common/macros.h"
+#include "paddle/fluid/framework/new_executor/pir_interpreter.h"
+#include "paddle/fluid/framework/new_executor/standalone_executor.h"
+#include "paddle/fluid/pir/dialect/operator/ir/api_builder.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/dialect/operator/utils/utils.h"
+#include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
+#include "paddle/pir/include/core/block.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/builtin_op.h"
+#include "paddle/pir/include/core/ir_context.h"
+#include "paddle/pir/include/core/program.h"
+#include "paddle/pir/include/core/utils.h"
+
+DECLARE_FILE_SYMBOLS(kernel_dialect);
+
+PD_DECLARE_KERNEL(full, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(tanh, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(mean, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(tanh_grad, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(mean_grad, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(add_grad, CPU, ALL_LAYOUT);
+
+namespace paddle {
+namespace framework {
+
+pir::Operation* GetOpFromProgram(const std::string& op_name,
+                                 const pir::Program& program) {
+  for (auto& op : *(program.block())) {
+    if (op.name() == op_name) {
+      return &op;
+    }
+  }
+  return nullptr;
+}
+
+TEST(VJP, TanhBackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::TanhOp op2 =
+      builder->Build<paddle::dialect::TanhOp>(op1.out());
+
+  paddle::dialect::FullOp op3 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  std::vector<std::vector<bool>> stop_gradients{{false}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()}};
+  std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
+  std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};
+
+  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.tanh");
+  auto tanh_vjp_interface_impl =
+      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+  tanh_vjp_interface_impl->vjp_(
+      op2.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  builder->Build<pir::ShadowOutputOp>(op2->result(0), "tanh_out");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.tanh_grad", program)->result(0), "tanh_grad_out");
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+  test_core.SetSkipGcVars({"tanh_out", "tanh_grad_out"});
+  test_core.Run({});
+  auto out_tensor = test_core.local_scope() == nullptr
+                        ? scope.FindVar("tanh_out")->Get<phi::DenseTensor>()
+                        : test_core.local_scope()
+                              ->FindVar("tanh_out")
+                              ->Get<phi::DenseTensor>();
+  auto grad_out_tensor =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar("tanh_grad_out")->Get<phi::DenseTensor>()
+          : test_core.local_scope()
+                ->FindVar("tanh_grad_out")
+                ->Get<phi::DenseTensor>();
+
+  ASSERT_NEAR(out_tensor.data<float>()[0], 0.76159, 1e-5);
+  ASSERT_NEAR(grad_out_tensor.data<float>()[0], 0.83995, 1e-5);
+}
+
+TEST(VJP, Tanh_BackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::Tanh_Op op2 =
+      builder->Build<paddle::dialect::Tanh_Op>(op1.out());
+
+  paddle::dialect::FullOp op3 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  std::vector<std::vector<bool>> stop_gradients{{false}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()}};
+  std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
+  std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};
+
+  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.tanh_");
+  auto tanh_vjp_interface_impl =
+      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+  tanh_vjp_interface_impl->vjp_(
+      op2.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  std::string tanh_out = "tanh_out";
+  std::string tanh_grad_out = "tanh_grad_out";
+  builder->Build<pir::ShadowOutputOp>(op2->result(0), tanh_out);
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.tanh_grad", program)->result(0), tanh_grad_out);
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+  test_core.SetSkipGcVars({tanh_out, tanh_grad_out});
+  test_core.Run({});
+  auto out_tensor =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar(tanh_out)->Get<phi::DenseTensor>()
+          : test_core.local_scope()->FindVar(tanh_out)->Get<phi::DenseTensor>();
+  auto grad_out_tensor =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar(tanh_grad_out)->Get<phi::DenseTensor>()
+          : test_core.local_scope()
+                ->FindVar(tanh_grad_out)
+                ->Get<phi::DenseTensor>();
+
+  ASSERT_NEAR(out_tensor.data<float>()[0], 0.76159, 1e-5);
+  ASSERT_NEAR(grad_out_tensor.data<float>()[0], 0.83995, 1e-5);
+}
+
+TEST(VJP, MeanBackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::MeanOp op2 =
+      builder->Build<paddle::dialect::MeanOp>(op1.out());
+
+  paddle::dialect::FullOp op3 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  std::vector<std::vector<bool>> stop_gradients{{false}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()},
+                                              {op2.operand_source(1)}};
+  std::vector<std::vector<pir::Value>> outputs{{op2.out()}};
+  std::vector<std::vector<pir::Value>> out_grads{{op3.out()}};
+
+  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.mean");
+  auto mean_vjp_interface_impl =
+      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+  mean_vjp_interface_impl->vjp_(
+      op2.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  builder->Build<pir::ShadowOutputOp>(op2->result(0), "mean_out");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.mean_grad", program)->result(0), "mean_grad_out");
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  ProgramDesc prog_desc;
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+  test_core.SetSkipGcVars({"mean_out", "mean_grad_out"});
+  test_core.Run({});
+  auto out_tensor = test_core.local_scope() == nullptr
+                        ? scope.FindVar("mean_out")->Get<phi::DenseTensor>()
+                        : test_core.local_scope()
+                              ->FindVar("mean_out")
+                              ->Get<phi::DenseTensor>();
+  auto grad_out_tensor =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar("mean_grad_out")->Get<phi::DenseTensor>()
+          : test_core.local_scope()
+                ->FindVar("mean_grad_out")
+                ->Get<phi::DenseTensor>();
+  ASSERT_EQ(out_tensor.data<float>()[0], 2.0);
+  ASSERT_EQ(grad_out_tensor.data<float>()[0], 0.25);
+  ASSERT_EQ(grad_out_tensor.data<float>()[1], 0.25);
+  ASSERT_EQ(grad_out_tensor.data<float>()[2], 0.25);
+  ASSERT_EQ(grad_out_tensor.data<float>()[3], 0.25);
+}
+
+TEST(VJP, ConcatBackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1, 2}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  std::vector<pir::Value> combine_input{{op1.out(), op1.out()}};
+  pir::CombineOp op2 = builder->Build<pir::CombineOp>(combine_input);
+  paddle::dialect::ConcatOp op3 =
+      builder->Build<paddle::dialect::ConcatOp>(op2.out(), 0);
+
+  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  std::vector<std::vector<bool>> stop_gradients{{false, false}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out(), op1.out()},
+                                              {op3.axis()}};
+  std::vector<std::vector<pir::Value>> outputs{{op3.out()}};
+  std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
+  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.concat");
+  auto concat_vjp_interface_impl =
+      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+  concat_vjp_interface_impl->vjp_(
+      op3.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  builder->Build<pir::ShadowOutputOp>(op3->result(0), "concat_out");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("builtin.split", program)->result(0), "split_out_0");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("builtin.split", program)->result(1), "split_out_1");
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  ProgramDesc prog_desc;
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+  test_core.SetSkipGcVars({"concat_out", "split_out_0", "split_out_1"});
+  test_core.Run({});
+  auto out_tensor = test_core.local_scope() == nullptr
+                        ? scope.FindVar("concat_out")->Get<phi::DenseTensor>()
+                        : test_core.local_scope()
+                              ->FindVar("concat_out")
+                              ->Get<phi::DenseTensor>();
+  auto grad_out_tensor_0 =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar("split_out_0")->Get<phi::DenseTensor>()
+          : test_core.local_scope()
+                ->FindVar("split_out_0")
+                ->Get<phi::DenseTensor>();
+  auto grad_out_tensor_1 =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar("split_out_1")->Get<phi::DenseTensor>()
+          : test_core.local_scope()
+                ->FindVar("split_out_1")
+                ->Get<phi::DenseTensor>();
+  ASSERT_EQ(out_tensor.data<float>()[0], 2.0);
+  ASSERT_EQ(grad_out_tensor_0.data<float>()[0], 1.0);
+  ASSERT_EQ(grad_out_tensor_0.data<float>()[1], 1.0);
+  ASSERT_EQ(grad_out_tensor_1.data<float>()[0], 1.0);
+  ASSERT_EQ(grad_out_tensor_1.data<float>()[1], 1.0);
+}
+
+TEST(VJP, AddBackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::FullOp op2 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::AddOp op3 =
+      builder->Build<paddle::dialect::AddOp>(op1.out(), op2.out());
+
+  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  std::vector<std::vector<bool>> stop_gradients{{false}, {false}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()}, {op2.out()}};
+  std::vector<std::vector<pir::Value>> outputs{{op3.out()}};
+  std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
+
+  pir::OpInfo op3_info = ctx->GetRegisteredOpInfo("pd_op.add");
+  auto add_vjp_interface_impl =
+      op3_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+  add_vjp_interface_impl->vjp_(
+      op3.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  builder->Build<pir::ShadowOutputOp>(op3->result(0), "add_out");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.add_grad", program)->result(0), "add_grad_out_0");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.add_grad", program)->result(1), "add_grad_out_1");
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  ProgramDesc prog_desc;
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+  test_core.SetSkipGcVars({"add_out", "add_grad_out_0", "add_grad_out_1"});
+  test_core.Run({});
+  auto out_tensor = test_core.local_scope() == nullptr
+                        ? scope.FindVar("add_out")->Get<phi::DenseTensor>()
+                        : test_core.local_scope()
+                              ->FindVar("add_out")
+                              ->Get<phi::DenseTensor>();
+  auto dx = test_core.local_scope() == nullptr
+                ? scope.FindVar("add_grad_out_0")->Get<phi::DenseTensor>()
+                : test_core.local_scope()
+                      ->FindVar("add_grad_out_0")
+                      ->Get<phi::DenseTensor>();
+
+  auto dy = test_core.local_scope() == nullptr
+                ? scope.FindVar("add_grad_out_1")->Get<phi::DenseTensor>()
+                : test_core.local_scope()
+                      ->FindVar("add_grad_out_1")
+                      ->Get<phi::DenseTensor>();
+  ASSERT_EQ(out_tensor.data<float>()[0], 4.0);
+  ASSERT_EQ(dx.data<float>()[0], 1.0);
+  ASSERT_EQ(dy.data<float>()[0], 1.0);
+}
+
+TEST(VJP, Add_BackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::FullOp op2 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+  paddle::dialect::Add_Op op3 =
+      builder->Build<paddle::dialect::Add_Op>(op1.out(), op2.out());
+
+  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  std::vector<std::vector<bool>> stop_gradients{{false}, {false}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()}, {op2.out()}};
+  std::vector<std::vector<pir::Value>> outputs{{op3.out()}};
+  std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
+
+  pir::OpInfo op3_info = ctx->GetRegisteredOpInfo("pd_op.add_");
+  auto add_inplace_vjp_interface_impl =
+      op3_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+  add_inplace_vjp_interface_impl->vjp_(
+      op3.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  builder->Build<pir::ShadowOutputOp>(op1->result(0), "full_op1_out");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.add_grad", program)->result(0), "add_grad_out_0");
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.add_grad", program)->result(1), "add_grad_out_1");
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  ProgramDesc prog_desc;
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+
+  test_core.SetSkipGcVars({"full_op1_out", "add_grad_out_0", "add_grad_out_1"});
+  test_core.Run({});
+  auto out_tensor = test_core.local_scope() == nullptr
+                        ? scope.FindVar("full_op1_out")->Get<phi::DenseTensor>()
+                        : test_core.local_scope()
+                              ->FindVar("full_op1_out")
+                              ->Get<phi::DenseTensor>();
+  auto dx = test_core.local_scope() == nullptr
+                ? scope.FindVar("add_grad_out_0")->Get<phi::DenseTensor>()
+                : test_core.local_scope()
+                      ->FindVar("add_grad_out_0")
+                      ->Get<phi::DenseTensor>();
+
+  auto dy = test_core.local_scope() == nullptr
+                ? scope.FindVar("add_grad_out_1")->Get<phi::DenseTensor>()
+                : test_core.local_scope()
+                      ->FindVar("add_grad_out_1")
+                      ->Get<phi::DenseTensor>();
+  ASSERT_EQ(out_tensor.data<float>()[0], 4.0);
+  ASSERT_EQ(dx.data<float>()[0], 1.0);
+  ASSERT_EQ(dy.data<float>()[0], 1.0);
+}
+
+TEST(VJP, SplitBackwardTest) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Program program((ctx));
+  paddle::dialect::ApiBuilder::Instance().SetProgram(&program);
+
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  paddle::dialect::FullOp op1 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2}, 2.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  paddle::dialect::SplitOp op2 = builder->Build<paddle::dialect::SplitOp>(
+      op1.out(), std::vector<int64_t>{1, 1}, 0);
+
+  pir::SplitOp op3 = builder->Build<pir::SplitOp>(op2.out());
+
+  paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1, 2}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
+
+  std::vector<std::vector<bool>> stop_gradients{{false}};
+  std::vector<std::vector<pir::Value>> inputs{
+      {op2.x()}, {op2.sections()}, {op2.axis()}};
+  std::vector<std::vector<pir::Value>> outputs(1);
+  std::vector<pir::Value> res;
+  for (uint32_t i = 0; i < op3.outputs().size(); i++) {
+    res.push_back(op3.outputs()[i]);
+  }
+  outputs[0] = res;
+  std::vector<std::vector<pir::Value>> out_grads{{op3.result(0), op4.out()}};
+  pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.split");
+
+  auto concat_vjp_interface_impl =
+      op2_info.GetInterfaceImpl<paddle::dialect::VjpInterface>();
+
+  concat_vjp_interface_impl->vjp_(
+      op2.operation(), inputs, outputs, out_grads, stop_gradients);
+
+  std::string split_out1 = "split_out1";
+  std::string split_out2 = "split_out2";
+  std::string concat_out = "concat_out";
+
+  builder->Build<pir::ShadowOutputOp>(op3->result(0), split_out1);
+  builder->Build<pir::ShadowOutputOp>(op3->result(1), split_out2);
+  builder->Build<pir::ShadowOutputOp>(
+      GetOpFromProgram("pd_op.concat", program)->result(0), concat_out);
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(&program);
+
+  auto place = phi::CPUPlace();
+  Scope scope;
+
+  InterpreterCore test_core(place, {}, kernel_program->block(), &scope);
+
+  test_core.SetSkipGcVars({split_out1, split_out2, concat_out});
+  test_core.Run({});
+  auto out_tensor_0 = test_core.local_scope() == nullptr
+                          ? scope.FindVar(split_out1)->Get<phi::DenseTensor>()
+                          : test_core.local_scope()
+                                ->FindVar(split_out1)
+                                ->Get<phi::DenseTensor>();
+  auto out_tensor_1 = test_core.local_scope() == nullptr
+                          ? scope.FindVar(split_out2)->Get<phi::DenseTensor>()
+                          : test_core.local_scope()
+                                ->FindVar(split_out2)
+                                ->Get<phi::DenseTensor>();
+  auto grad_out_tensor_0 =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar(concat_out)->Get<phi::DenseTensor>()
+          : test_core.local_scope()
+                ->FindVar(concat_out)
+                ->Get<phi::DenseTensor>();
+  ASSERT_EQ(out_tensor_0.data<float>()[0], 2.0);
+  ASSERT_EQ(out_tensor_0.data<float>()[1], 2.0);
+  ASSERT_EQ(out_tensor_1.data<float>()[0], 2.0);
+  ASSERT_EQ(out_tensor_1.data<float>()[1], 2.0);
+  ASSERT_EQ(grad_out_tensor_0.data<float>()[0], 2.0);
+  ASSERT_EQ(grad_out_tensor_0.data<float>()[1], 2.0);
+  ASSERT_EQ(grad_out_tensor_0.data<float>()[2], 1.0);
+  ASSERT_EQ(grad_out_tensor_0.data<float>()[3], 1.0);
+}
+
+}  // namespace framework
+}  // namespace paddle