chore: import upstream snapshot with attribution

test/cpp/pir/custom_engine/CMakeLists.txt

@@ -0,0 +1,12 @@
+if(WITH_CUSTOM_DEVICE)
+  paddle_test(test_custom_engine_loadlib SRCS test_custom_engine_loadlib.cc
+              custom_engine_op.cc)
+
+endif()
+paddle_test(test_custom_engine SRCS test_custom_engine_operation.cc
+            custom_engine_op.cc)
+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_custom_engine)
+endif()

test/cpp/pir/custom_engine/custom_engine_op.cc

@@ -0,0 +1,184 @@
+// Copyright (c) 2024 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/pir/custom_engine/custom_engine_op.h"
+#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
+
+namespace paddle {
+namespace dialect {
+
+const char *FakeEngineOp::attributes_name[2] = {"input_names", "output_names"};
+
+OpInfoTuple FakeEngineOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      OpInputInfo("x",
+                  "pir::VectorType<paddle::dialect::DenseTensorType>",
+                  false,
+                  false,
+                  false,
+                  false)};
+
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo(
+          "input_names", "pir::ArrayAttribute", ""),
+      paddle::dialect::OpAttributeInfo(
+          "output_names", "pir::ArrayAttribute", "")};
+
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      OpOutputInfo("out",
+                   "pir::VectorType<paddle::dialect::DenseTensorType>",
+                   false,
+                   false)};
+
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      OpRunTimeInfo("", {}, "", {}, {}, {}, {}, {});
+
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "fake_engine");
+}
+
+#define ADD_VEC_ATTRIBUTE(type, name)                                   \
+  std::vector<pir::Attribute> name##_tmp;                               \
+  name##_tmp.reserve(name.size());                                      \
+  for (const auto &v : name) {                                          \
+    name##_tmp.push_back(type::get(pir::IrContext::Instance(), v));     \
+  }                                                                     \
+  pir::Attribute attr_##name =                                          \
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), name##_tmp); \
+  argument.AddAttribute(#name, attr_##name)
+
+#define VERIFY_ATTRIBUTE(type, name)                              \
+  PADDLE_ENFORCE_GT(                                              \
+      attributes.count(#name),                                    \
+      0,                                                          \
+      common::errors::InvalidArgument(#name " does not exist.")); \
+  PADDLE_ENFORCE_EQ(attributes.at(#name).isa<type>(),             \
+                    true,                                         \
+                    common::errors::InvalidArgument(              \
+                        "Type of attribute: " #name " is not " #type))
+
+void FakeEngineOp::Build(pir::Builder &builder,             // NOLINT
+                         pir::OperationArgument &argument,  // NOLINT
+                         pir::Value x,
+                         std::vector<std::string> input_names,
+                         std::vector<std::string> output_names,
+                         std::vector<std::vector<int64_t>> outputs_shape,
+                         std::vector<phi::DataType> outputs_dtype) {
+  VLOG(4) << "Start build FakeEngineOp";
+
+  VLOG(4) << "Builder construction inputs";
+  std::vector<pir::Value> argument_inputs = {x};
+  argument.AddInputs(argument_inputs);
+
+  VLOG(4) << "Builder construction attributes";
+
+  ADD_VEC_ATTRIBUTE(pir::StrAttribute, input_names);
+  ADD_VEC_ATTRIBUTE(pir::StrAttribute, output_names);
+
+  VLOG(4) << "Builder construction outputs";
+
+  std::vector<pir::Type> argument_outputs;
+  std::vector<pir::Type> out_types;
+  for (size_t i = 0; i < static_cast<size_t>(outputs_shape.size()); i++) {
+    if (outputs_dtype[i] == phi::DataType::UNDEFINED) {
+      out_types.push_back(pir::Type());
+    } else {
+      out_types.push_back(pir::DenseTensorType::get(
+          pir::IrContext::Instance(),
+          TransToIrDataType(outputs_dtype[i]),
+          phi::DDim(outputs_shape[i].data(), outputs_shape[i].size()),
+          phi::DataLayout::kNCHW,
+          phi::LoD(),
+          0));
+    }
+  }
+  pir::Type out_vector_type =
+      pir::VectorType::get(pir::IrContext::Instance(), out_types);
+  argument_outputs.push_back(out_vector_type);
+
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  argument.AddRegion(nullptr);
+  ::pir::PassStopGradientsDefaultly(argument);
+}
+
+void FakeEngineOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: "
+             "FakeEngineOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    PADDLE_ENFORCE_EQ(input_size,
+                      1,
+                      common::errors::InvalidArgument(
+                          "The size of inputs must be equal to 1."));
+    PADDLE_ENFORCE_EQ((*this)->operand_source(0).type().isa<pir::VectorType>(),
+                      true,
+                      common::errors::InvalidArgument(
+                          "Type validation failed for the 0th input, got %s.",
+                          (*this)->operand_source(0).type()));
+    if (auto vec_type =
+            (*this)->operand_source(0).type().dyn_cast<pir::VectorType>()) {
+      for (size_t i = 0; i < vec_type.size(); ++i) {
+        PADDLE_ENFORCE_EQ(
+            vec_type[i].isa<pir::DenseTensorType>(),
+            true,
+            common::errors::InvalidArgument(
+                "Type validation failed for the 0th input, got %s.",
+                (*this)->operand_source(0).type()));
+      }
+    }
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    VERIFY_ATTRIBUTE(pir::ArrayAttribute, input_names);
+    VERIFY_ATTRIBUTE(pir::ArrayAttribute, output_names);
+  }
+
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    PADDLE_ENFORCE_EQ(output_size,
+                      1,
+                      common::errors::InvalidArgument(
+                          "The size of outputs must be equal to 1."));
+    auto output_type = (*this)->result(0).type();
+
+    PADDLE_ENFORCE_EQ(output_type.isa<pir::VectorType>(),
+                      true,
+                      common::errors::InvalidArgument(
+                          "Type validation failed for the 0th output."));
+  }
+  VLOG(4) << "End Verifying for: FakeEngineOp.";
+}
+
+pir::Block *FakeEngineOp::block() {
+  pir::Region &region = (*this)->region(0);
+  if (region.empty()) region.emplace_back();
+  return &region.front();
+}
+
+pir::Block *FakeEngineOp::block() const {
+  pir::Region &region = (*this)->region(0);
+  PADDLE_ENFORCE_EQ(region.empty(),
+                    false,
+                    ::common::errors::Unavailable(
+                        "Required CustomEngineOp's region must not be empty."));
+  return &region.front();
+}
+
+}  // namespace dialect
+}  // namespace paddle
+
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::FakeEngineOp)

test/cpp/pir/custom_engine/custom_engine_op.h

@@ -0,0 +1,81 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <glog/logging.h>
+#include <vector>
+
+#include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/pir/include/core/builder.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/builtin_op.h"
+#include "paddle/pir/include/core/builtin_type.h"
+#include "paddle/pir/include/core/op_base.h"
+#include "paddle/pir/include/core/op_trait.h"
+#include "paddle/pir/include/core/operation_utils.h"
+#include "test/cpp/pir/tools/macros_utils.h"
+
+#if defined(_WIN32)
+#ifndef EXPORT_API
+#define EXPORT_API __declspec(dllexport)
+#endif  // EXPORT_API
+#else
+#define EXPORT_API
+#endif  // _WIN32
+
+#define IR_DECLARE_EXPLICIT_PLUGIN_TYPE_ID(TYPE_CLASS) \
+  namespace pir {                                      \
+  namespace detail {                                   \
+  template <>                                          \
+  class EXPORT_API TypeIdResolver<TYPE_CLASS> {        \
+   public:                                             \
+    static TypeId Resolve() { return id_; }            \
+    static UniqueingId id_;                            \
+  };                                                   \
+  }                                                    \
+  }  // namespace pir
+
+namespace paddle {
+namespace dialect {
+
+class FakeEngineOp
+    : public pir::Op<FakeEngineOp, paddle::dialect::OpYamlInfoInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "custom_engine.fake_engine"; }
+  static const char *attributes_name[2];
+  static constexpr uint32_t attributes_num = 2;
+  static OpInfoTuple GetOpInfo();
+
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x,
+                    std::vector<std::string> input_names,
+                    std::vector<std::string> output_names,
+                    std::vector<std::vector<int64_t>> outputs_shape,
+                    std::vector<phi::DataType> outputs_dtype);
+
+  void VerifySig();
+  pir::Block *block();
+  pir::Block *block() const;
+  pir::Value x() { return operand_source(0); }
+  pir::Value out() { return result(0); }
+};
+
+}  // namespace dialect
+}  // namespace paddle
+
+IR_DECLARE_EXPLICIT_PLUGIN_TYPE_ID(paddle::dialect::FakeEngineOp)

test/cpp/pir/custom_engine/fake_cpu_engine.h

@@ -0,0 +1,263 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <unordered_map>
+
+#include "paddle/fluid/custom_engine/custom_engine_ext.h"
+#include "paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.h"
+#include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
+#include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
+#include "test/cpp/pir/custom_engine/custom_engine_op.h"
+#include "test/cpp/pir/custom_engine/fake_cpu_engine_base.h"
+
+C_Status RegisterCustomEngineOp() {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  pir::Dialect* custom_engine_dialect =
+      ctx->GetOrRegisterDialect<paddle::dialect::CustomEngineDialect>();
+  EXPECT_EQ(custom_engine_dialect != nullptr, true);
+  ctx->RegisterOpInfo(custom_engine_dialect,
+                      pir::TypeId::get<paddle::dialect::FakeEngineOp>(),
+                      paddle::dialect::FakeEngineOp::name(),
+                      paddle::dialect::FakeEngineOp::interface_set(),
+                      paddle::dialect::FakeEngineOp::GetTraitSet(),
+                      paddle::dialect::FakeEngineOp::attributes_num,
+                      paddle::dialect::FakeEngineOp::attributes_name,
+                      paddle::dialect::FakeEngineOp::VerifySigInvariants,
+                      paddle::dialect::FakeEngineOp::VerifyRegionInvariants);
+  return C_SUCCESS;
+}
+
+C_Status CustomEngineOpLower(C_CustomEngineLowerParams* lower_param) {
+  // get lower params
+  pir::IrContext* ctx =
+      reinterpret_cast<pir::IrContext*>(lower_param->ir_context);
+  pir::Operation* op_item =
+      reinterpret_cast<pir::Operation*>(lower_param->operation);
+  phi::KernelKey* kernel_key =
+      reinterpret_cast<phi::KernelKey*>(lower_param->kernel_key);
+  phi::Place* place = reinterpret_cast<phi::Place*>(lower_param->place);
+  std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair =
+      reinterpret_cast<std::unordered_map<pir::Operation*, pir::Operation*>*>(
+          lower_param->map_op_pair);
+  std::unordered_map<pir::Value, pir::Value>* map_value_pair =
+      reinterpret_cast<std::unordered_map<pir::Value, pir::Value>*>(
+          lower_param->map_value_pair);
+  pir::Block* block = reinterpret_cast<pir::Block*>(lower_param->block);
+
+  // Prepare output types
+  std::vector<pir::Type> op_output_types;
+  for (size_t i = 0; i < op_item->num_results(); ++i) {
+    PushBackOutputTypes(ctx,
+                        op_item,
+                        op_item->result(i).type(),
+                        *place,
+                        *kernel_key,
+                        &op_output_types);
+  }
+  // Prepare input
+  std::vector<pir::Value> vec_inputs;
+
+  for (size_t i = 0; i < op_item->num_operands(); ++i) {
+    auto cur_in = op_item->operand_source(i);
+    PADDLE_ENFORCE_EQ(
+        map_value_pair->count(cur_in),
+        true,
+        common::errors::PreconditionNotMet(
+            "[%d]'s input of [%s] op MUST in map pair", i, op_item->name()));
+
+    auto new_in = map_value_pair->at(cur_in);
+
+    vec_inputs.push_back(new_in);
+  }
+  // Prepare attr
+  std::unordered_map<std::string, pir::Attribute> op_attribute;
+  auto op_attr_map = op_item->attributes();
+  for (auto& map_item : op_attr_map) {
+    op_attribute.emplace(map_item.first, map_item.second);
+  }
+
+  pir::OpInfo custom_engine_op_info =
+      ctx->GetRegisteredOpInfo(paddle::dialect::FakeEngineOp::name());
+  pir::Operation* op = pir::Operation::Create(vec_inputs,
+                                              op_attribute,
+                                              op_output_types,
+                                              custom_engine_op_info,
+                                              1,
+                                              {},
+                                              true);
+  op->set_attribute("origin_id", pir::Int64Attribute::get(ctx, op->id()));
+  op->set_attribute("op_name", pir::StrAttribute::get(ctx, op->name()));
+
+  VLOG(3) << "CustomEngineOpLower get op_item subgraph block.";
+  pir::Region& op_item_region = op_item->region(0);
+  PADDLE_ENFORCE_EQ(op_item_region.empty(),
+                    false,
+                    ::common::errors::Unavailable(
+                        "Required CustomEngineOp's region must not be empty."));
+  pir::Block* sub_graph_block = &(op_item_region.front());
+
+  VLOG(3) << "CustomEngineOpLower set new op subgraph block.";
+  pir::Region& region = op->region(0);
+  if (region.empty()) {
+    region.emplace_back();
+  }
+  pir::Block* op_block = &(region.front());
+
+  // process subgraph block
+  pir::ProcessBlock(
+      *place, sub_graph_block, op_block, ctx, map_op_pair, map_value_pair);
+
+  if (VLOG_IS_ON(3)) {
+    std::stringstream ss;
+    ss << "CustomEngineOpLower new op:";
+    op->Print(ss);
+    VLOG(3) << ss.str();
+  }
+
+  (*map_op_pair)[op_item] = op;
+  // only deal with single output
+  if (op_item->num_results() > 0) {
+    for (size_t i = 0; i < op_item->num_results(); ++i) {
+      (*map_value_pair)[op_item->result(i)] = op->result(i);
+    }
+  }
+  block->push_back(op);
+  return C_SUCCESS;
+}
+
+class CustomEngine {
+ public:
+  CustomEngine(std::vector<phi::DenseTensor*> tensor_args,
+               std::vector<phi::DenseTensor*> return_tensor)
+      : tensor_args_(tensor_args), return_tensor_(return_tensor) {}
+  ~CustomEngine() {}
+
+  void Run(const phi::DeviceContext& device_ctx, const phi::Place& place) {
+    PADDLE_ENFORCE_EQ(
+        tensor_args_.size(),
+        2u,
+        common::errors::PreconditionNotMet("tensor_args.size != 2"));
+    PADDLE_ENFORCE_EQ(
+        return_tensor_.size(),
+        1u,
+        common::errors::PreconditionNotMet("return_tensor.size != 1"));
+    // phi::AddKernel<float, phi::DeviceContext>(device_ctx, *(tensor_args_[0]),
+    // *(tensor_args_[1]),return_tensor_[0]);
+    phi::Copy(device_ctx, *(tensor_args_[0]), place, true, return_tensor_[0]);
+    return;
+  }
+
+ private:
+  std::vector<phi::DenseTensor*> tensor_args_;
+  std::vector<phi::DenseTensor*> return_tensor_;
+  std::vector<phi::DenseTensor*> template_tensor_;
+};
+
+C_Status GraphEngineExecute(C_CustomEngineInstruction instruction) {
+  paddle::framework::CustomEngineInstruction* instruction_ =
+      reinterpret_cast<paddle::framework::CustomEngineInstruction*>(
+          instruction);
+  CustomEngine* customengine =
+      reinterpret_cast<CustomEngine*>(instruction_->CustomEngine());
+
+  customengine->Run(instruction_->DeviceContext(),
+                    instruction_->DeviceContext().GetPlace());
+  return C_SUCCESS;
+}
+
+C_Status GraphEngineBuild(C_CustomEngineInstruction instruction) {
+  paddle::framework::CustomEngineInstruction* instruction_ =
+      reinterpret_cast<paddle::framework::CustomEngineInstruction*>(
+          instruction);
+  pir::Operation* op = instruction_->Operation();
+  const paddle::framework::ValueExecutionInfo* value_exec_info =
+      instruction_->GetValueExecutionInfo();
+  // prepare input tensors
+  std::vector<phi::DenseTensor*> tensor_args;
+  PADDLE_ENFORCE_EQ(op->num_operands(),
+                    1u,
+                    common::errors::PreconditionNotMet(
+                        "custom engine op should has 1 operand"));
+  auto vec_in = op->operand_source(0).defining_op()->operands_source();
+  for (auto in : vec_in) {
+    auto var_name = value_exec_info->GetVarName(in);
+    auto tensor = value_exec_info->GetScope()
+                      ->FindVar(var_name)
+                      ->GetMutable<phi::DenseTensor>();
+    tensor_args.push_back(tensor);
+  }
+
+  // prepare output tensors
+  std::vector<phi::DenseTensor*> return_tensor;
+  PADDLE_ENFORCE_EQ(op->num_results(),
+                    1u,
+                    common::errors::PreconditionNotMet(
+                        "custom engine op should has 1 result"));
+  pir::Value vec_result = op->result(0);
+  PADDLE_ENFORCE_EQ(vec_result.type().isa<pir::VectorType>(),
+                    true,
+                    common::errors::PreconditionNotMet(
+                        "custom engine op result should be vectortype"));
+  auto vec_out = op->result(0).first_use().owner()->results();
+
+  for (auto out : vec_out) {
+    bool check =
+        out && out.type() && out.type().isa<paddle::dialect::DenseTensorType>();
+    PADDLE_ENFORCE_EQ(
+        check,
+        true,
+        common::errors::PreconditionNotMet(
+            "customEngine instruction only support DenseTensorType"));
+    auto var_name = value_exec_info->GetVarName(out);
+
+    auto tensor = value_exec_info->GetScope()
+                      ->Var(var_name)
+                      ->GetMutable<phi::DenseTensor>();
+
+    return_tensor.push_back(tensor);
+    auto alloc_tensor_type =
+        out.type().dyn_cast<paddle::dialect::AllocatedDenseTensorType>();
+    tensor->set_type(
+        paddle::dialect::TransToPhiDataType(alloc_tensor_type.dtype()));
+    tensor->Resize(alloc_tensor_type.dims());
+  }
+
+  CustomEngine* fake_engine = new CustomEngine(tensor_args, return_tensor);
+  auto customEngineDeleter = [](void* ptr) {
+    CustomEngine* customEngine = static_cast<CustomEngine*>(ptr);
+
+    if (customEngine != nullptr) {
+      delete customEngine;
+    } else {
+      PADDLE_THROW(
+          common::errors::PreconditionNotMet("customEngine is nullptr"));
+    }
+  };
+  instruction_->SetCustomEngine(reinterpret_cast<void*>(fake_engine));
+  instruction_->SetCustomEngineDeleter(customEngineDeleter);
+
+  return C_SUCCESS;
+}
+
+void InitPluginCustomEngine(CustomEngineParams* params) {
+  memset(reinterpret_cast<void*>(params->interface),
+         0,
+         sizeof(C_CustomEngineInterface));
+
+  params->interface->register_custom_engine_op = RegisterCustomEngineOp;
+  params->interface->graph_engine_build = GraphEngineBuild;
+  params->interface->graph_engine_execute = GraphEngineExecute;
+  params->interface->custom_engine_op_lower = CustomEngineOpLower;
+}

test/cpp/pir/custom_engine/fake_cpu_engine_base.h

@@ -0,0 +1,173 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <unordered_map>
+
+#include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/kernel_factory.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/pir/include/core/block.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/ir_context.h"
+#include "paddle/pir/include/core/operation.h"
+
+using paddle::dialect::AllocatedDenseTensorArrayType;
+using paddle::dialect::AllocatedDenseTensorType;
+using paddle::dialect::AllocatedSelectedRowsType;
+using paddle::dialect::AllocatedSparseCooTensorType;
+using paddle::dialect::AllocatedSparseCsrTensorType;
+using paddle::dialect::DenseTensorArrayType;
+using paddle::dialect::DenseTensorType;
+using paddle::dialect::SelectedRowsType;
+using paddle::dialect::SparseCooTensorType;
+using paddle::dialect::SparseCsrTensorType;
+
+template <class IrType1, class IrType2>
+static pir::Type create_sparse_coo_tensor_type(pir::Type type,
+                                               const phi::Place& place,
+                                               pir::Type out_dtype,
+                                               pir::IrContext* ctx) {
+  auto input_type = type.dyn_cast<IrType1>();
+  return IrType2::get(ctx,
+                      place,
+                      out_dtype,
+                      input_type.dims(),
+                      input_type.non_zero_dims(),
+                      input_type.data_layout(),
+                      input_type.non_zero_indices(),
+                      input_type.non_zero_elements(),
+                      input_type.coalesced());
+}
+
+template <class IrType1, class IrType2>
+static pir::Type create_sparse_csr_tensor_type(pir::Type type,
+                                               const phi::Place& place,
+                                               pir::Type out_dtype,
+                                               pir::IrContext* ctx) {
+  auto input_type = type.dyn_cast<IrType1>();
+  return IrType2::get(ctx,
+                      place,
+                      out_dtype,
+                      input_type.dims(),
+                      input_type.data_layout(),
+                      input_type.non_zero_crows(),
+                      input_type.non_zero_cols(),
+                      input_type.non_zero_elements());
+}
+
+template <class IrType1, class IrType2>
+static pir::Type create_type(pir::Type type,
+                             const phi::Place& place,
+                             pir::Type out_dtype,
+                             pir::IrContext* ctx) {
+  auto input_type = type.dyn_cast<IrType1>();
+  return IrType2::get(ctx,
+                      place,
+                      out_dtype,
+                      input_type.dims(),
+                      input_type.data_layout(),
+                      input_type.lod(),
+                      input_type.offset());
+}
+
+static pir::Type BuildOutputType(pir::Type type,
+                                 const phi::Place& place,
+                                 pir::IrContext* ctx) {
+  if (type.isa<DenseTensorType>()) {
+    auto out_dtype = type.dyn_cast<DenseTensorType>().dtype();
+    return create_type<DenseTensorType, AllocatedDenseTensorType>(
+        type, place, out_dtype, ctx);
+  } else if (type.isa<SelectedRowsType>()) {
+    auto out_dtype = type.dyn_cast<SelectedRowsType>().dtype();
+    return create_type<SelectedRowsType, AllocatedSelectedRowsType>(
+        type, place, out_dtype, ctx);
+  } else if (type.isa<DenseTensorArrayType>()) {
+    auto array_type = type.dyn_cast<DenseTensorArrayType>();
+    return AllocatedDenseTensorArrayType::get(ctx,
+                                              place,
+                                              array_type.dtype(),
+                                              array_type.dims(),
+                                              array_type.data_layout());
+  } else if (type.isa<SparseCooTensorType>()) {
+    auto out_dtype = type.dyn_cast<SparseCooTensorType>().dtype();
+    return create_sparse_coo_tensor_type<SparseCooTensorType,
+                                         AllocatedSparseCooTensorType>(
+        type, place, out_dtype, ctx);
+  } else if (type.isa<SparseCsrTensorType>()) {
+    auto out_dtype = type.dyn_cast<SparseCsrTensorType>().dtype();
+    return create_sparse_csr_tensor_type<SparseCsrTensorType,
+                                         AllocatedSparseCsrTensorType>(
+        type, place, out_dtype, ctx);
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "BuildOutputType only support DenseTensorType, SelectedRowsType, "
+        "SparseCooTensorType and SparseCsrTensorType"));
+  }
+}
+
+void PushBackOutputTypes(pir::IrContext* ctx,
+                         pir::Operation* op_item,
+                         const pir::Type& origin_type,
+                         const phi::Place& out_place,
+                         const phi::KernelKey& kernel_key,
+                         std::vector<pir::Type>* op_output_types) {
+  auto result_type = origin_type;
+  if (!result_type) {
+    op_output_types->push_back(result_type);
+  } else if (result_type.isa<paddle::dialect::DenseTensorType>() ||
+             result_type.isa<paddle::dialect::SelectedRowsType>() ||
+             result_type.isa<paddle::dialect::DenseTensorArrayType>() ||
+             result_type.isa<paddle::dialect::SparseCooTensorType>() ||
+             result_type.isa<paddle::dialect::SparseCsrTensorType>()) {
+  } else if (result_type.isa<pir::VectorType>()) {
+    std::vector<pir::Type> vec_inner_types;
+    auto base_types = result_type.dyn_cast<pir::VectorType>().data();
+    for (auto& base_type : base_types) {
+      if (base_type) {
+        if (base_type.isa<paddle::dialect::DenseTensorType>() ||
+            base_type.isa<paddle::dialect::SelectedRowsType>()) {
+          vec_inner_types.push_back(BuildOutputType(base_type, out_place, ctx));
+        } else {
+          PADDLE_THROW(common::errors::Unimplemented(
+              "only support dense tensor and selected rows in vector type "
+              "for now"));
+        }
+      } else {
+        pir::Type fp32_dtype = pir::Float32Type::get(ctx);
+        phi::DDim dims = {};
+        phi::DataLayout data_layout = phi::DataLayout::NCHW;
+        phi::LegacyLoD lod = {{}};
+        size_t offset = 0;
+        auto dense_tensor_dtype = paddle::dialect::DenseTensorType::get(
+            ctx, fp32_dtype, dims, data_layout, lod, offset);
+        auto allocated_dense_tensor_dtype =
+            paddle::dialect::AllocatedDenseTensorType::get(
+                ctx, out_place, dense_tensor_dtype);
+        vec_inner_types.push_back(allocated_dense_tensor_dtype);
+      }
+    }
+
+    pir::Type t1 = pir::VectorType::get(ctx, vec_inner_types);
+    op_output_types->push_back(t1);
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Result type only support DenseTensorType, SelectedRowType, "
+        "SparseCooTensorType, SparseCsrTensorType and "
+        "VectorType"));
+  }
+}

test/cpp/pir/custom_engine/test_custom_engine_loadlib.cc

@@ -0,0 +1,152 @@
+// 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 <gtest/gtest.h>
+
+#include <array>
+#include <string>
+
+#include "paddle/fluid/custom_engine/custom_engine_manager.h"
+#include "paddle/fluid/platform/init.h"
+#include "paddle/phi/backends/custom/fake_cpu_device.h"
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator_facade.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+#include "paddle/fluid/framework/new_executor/interpretercore.h"
+#include "paddle/fluid/framework/new_executor/pir_interpreter.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/pd_op_to_kernel_pass.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/ir_context.h"
+#include "paddle/pir/include/core/program.h"
+#include "paddle/pir/include/core/type_id.h"
+#include "test/cpp/pir/custom_engine/fake_cpu_engine.h"
+
+#define OUT_NAME "program_out"
+
+void RegisterDevice() {
+  CustomRuntimeParams runtime_params;
+  runtime_params.size = sizeof(CustomRuntimeParams);
+  auto device_interface = std::make_unique<C_DeviceInterface>();
+  runtime_params.interface = device_interface.get();
+  std::memset(runtime_params.interface, 0, sizeof(C_DeviceInterface));
+  runtime_params.interface->size = sizeof(C_DeviceInterface);
+
+  InitFakeCPUDevice(&runtime_params);
+  phi::LoadCustomRuntimeLib(
+      runtime_params, std::move(device_interface), "", nullptr);
+}
+
+void RegisterEngine() {
+  CustomEngineParams engine_params;
+  std::memset(&engine_params, 0, sizeof(CustomEngineParams));
+  engine_params.size = sizeof(CustomEngineParams);
+  auto engine_interface = new (C_CustomEngineInterface);
+  engine_params.interface = engine_interface;
+  std::memset(engine_params.interface, 0, sizeof(C_CustomEngineInterface));
+  engine_params.interface->size = sizeof(C_CustomEngineInterface);
+
+  InitPluginCustomEngine(&engine_params);
+  paddle::custom_engine::LoadCustomEngineLib("", &engine_params);
+}
+
+void InitCustom() {
+  RegisterDevice();
+  EXPECT_GT(static_cast<int>(phi::DeviceManager::GetAllDeviceTypes().size()),
+            0);
+  auto place = phi::CustomPlace(DEVICE_TYPE, 0);
+  auto device = phi::DeviceManager::GetDeviceWithPlace(place);
+  EXPECT_NE(device, nullptr);
+
+  std::vector<phi::Place> places;
+  auto device_types = phi::DeviceManager::GetAllDeviceTypes();
+  for (auto dev_type : device_types) {
+    auto devices = phi::DeviceManager::GetDeviceList(dev_type);
+    for (auto dev_id : devices) {
+      places.push_back(phi::PlaceHelper::CreatePlace(dev_type, dev_id));
+    }
+  }
+  EXPECT_GT(static_cast<int>(places.size()), 0);
+
+  places.emplace_back(phi::CPUPlace());
+  phi::DeviceContextPool::Init(places);
+  RegisterEngine();
+}
+
+void CreateProgram(pir::Program *program) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+
+  pir::Block *block = program->block();
+  pir::Builder builder(ctx, block);
+
+  auto full_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{2, 2}, 100);
+  auto full_op2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{2, 2}, 10);
+
+  auto buildin_combine_op = builder.Build<pir::CombineOp>(
+      std::vector<pir::Value>{full_op1.result(0), full_op2.result(0)});
+
+  auto engine_op = builder.Build<paddle::dialect::FakeEngineOp>(
+      buildin_combine_op.result(0),
+      std::vector<std::string>{"input_0", "input_1"},
+      std::vector<std::string>{"output_0"},
+      std::vector<std::vector<int64_t>>{{2, 2}},
+      std::vector<phi::DataType>{phi::DataType::FLOAT32});
+  engine_op->region(0).emplace_back();
+
+  auto output = builder.Build<pir::SplitOp>(engine_op.result(0)).outputs()[0];
+  builder.Build<pir::ShadowOutputOp>(output, OUT_NAME);
+
+  return;
+}
+
+TEST(CustomDevice, Tensor) {
+  paddle::framework::InitMemoryMethod();
+  InitCustom();
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  pir::Program *program = new pir::Program(ctx);
+  CreateProgram(program);
+
+  EXPECT_EQ(program->block()->size(), 6u);
+
+  auto kernel_program = pir::PdOpLowerToKernelPass(program);
+
+  auto place = phi::CustomPlace(DEVICE_TYPE, 0);
+  paddle::framework::Scope scope;
+  paddle::framework::InterpreterCore test_core(
+      place, {}, kernel_program->block(), &scope);
+  test_core.SetSkipGcVars({OUT_NAME});
+  test_core.Run({});
+
+  auto out_tensor =
+      test_core.local_scope() == nullptr
+          ? scope.FindVar(OUT_NAME)->Get<phi::DenseTensor>()
+          : test_core.local_scope()->FindVar(OUT_NAME)->Get<phi::DenseTensor>();
+  bool res0 = out_tensor.data<float>()[0] == 100;
+  EXPECT_EQ(res0, true);
+}
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

test/cpp/pir/custom_engine/test_custom_engine_operation.cc

@@ -0,0 +1,104 @@
+// Copyright (c) 2024 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 <sstream>
+
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.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/ir_context.h"
+#include "paddle/pir/include/core/program.h"
+#include "paddle/pir/include/core/type_id.h"
+#include "test/cpp/pir/custom_engine/custom_engine_op.h"
+
+TEST(op_test, region_test) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Dialect *custom_engine_dialect =
+      ctx->GetOrRegisterDialect<paddle::dialect::CustomEngineDialect>();
+  EXPECT_EQ(custom_engine_dialect != nullptr, true);
+  ctx->RegisterOpInfo(custom_engine_dialect,
+                      pir::TypeId::get<paddle::dialect::FakeEngineOp>(),
+                      paddle::dialect::FakeEngineOp::name(),
+                      paddle::dialect::FakeEngineOp::interface_set(),
+                      paddle::dialect::FakeEngineOp::GetTraitSet(),
+                      paddle::dialect::FakeEngineOp::attributes_num,
+                      paddle::dialect::FakeEngineOp::attributes_name,
+                      paddle::dialect::FakeEngineOp::VerifySigInvariants,
+                      paddle::dialect::FakeEngineOp::VerifyRegionInvariants);
+
+  pir::Program program(ctx);
+  pir::Block *block = program.block();
+  pir::Builder builder(ctx, block);
+
+  auto full_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{2, 2}, 100);
+  auto full_op2 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{2, 2}, 10);
+
+  auto buildin_combine_op = builder.Build<pir::CombineOp>(
+      std::vector<pir::Value>{full_op1.result(0), full_op2.result(0)});
+
+  pir::OpInfo fake_engine_op_info =
+      ctx->GetRegisteredOpInfo(paddle::dialect::FakeEngineOp::name());
+
+  std::vector<pir::Type> out_types;
+  out_types.push_back(
+      pir::DenseTensorType::get(pir::IrContext::Instance(),
+                                pir::Float32Type::get(ctx),
+                                phi::DDim(std::vector<int64_t>{2, 2}.data(), 2),
+                                phi::DataLayout::kNCHW,
+                                phi::LoD(),
+                                0));
+  pir::Type out_vector_type =
+      pir::VectorType::get(pir::IrContext::Instance(), out_types);
+  std::vector<pir::Type> output_types = {out_vector_type};
+
+  pir::AttributeMap attribute_map;
+  std::vector<pir::Attribute> val;
+  val.push_back(pir::StrAttribute::get(ctx, "input_0"));
+  val.push_back(pir::StrAttribute::get(ctx, "input_1"));
+  attribute_map.insert({"input_names", pir::ArrayAttribute::get(ctx, val)});
+  std::vector<pir::Attribute> out_val;
+  out_val.push_back(pir::StrAttribute::get(ctx, "output_0"));
+  out_val.push_back(pir::StrAttribute::get(ctx, "output_1"));
+  attribute_map.insert(
+      {"output_names", pir::ArrayAttribute::get(ctx, out_val)});
+
+  pir::Operation *op1 = pir::Operation::Create({buildin_combine_op.result(0)},
+                                               attribute_map,
+                                               output_types,
+                                               fake_engine_op_info);
+
+  // (3) Test custom operation printer
+  std::stringstream ss1;
+  op1->Print(ss1);
+
+  builder.Insert(op1);
+
+  auto op2 = builder.Build<paddle::dialect::FakeEngineOp>(
+      buildin_combine_op.result(0),
+      std::vector<std::string>{"input_0", "input_1"},
+      std::vector<std::string>{"output_0"},
+      std::vector<std::vector<int64_t>>{{2, 2}},
+      std::vector<phi::DataType>{phi::DataType::FLOAT32});
+
+  std::stringstream ss2;
+  op2->Print(ss2);
+
+  EXPECT_EQ(block->size(), 5u);
+}