PaddlePaddle
diff --git a/‎paddle/fluid/framework/new_executor/feed_fetch_utils.cc
Copy file name to clipboard
+23-9Lines changed: 23 additions & 9 deletions b/‎paddle/fluid/framework/new_executor/feed_fetch_utils.cc
Copy file name to clipboard
+23-9Lines changed: 23 additions & 9 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion
diff --git a/‎paddle/fluid/operators/ops_signature/fetch_sig.cc
Copy file name to clipboard
+31Lines changed: 31 additions & 0 deletions b/‎paddle/fluid/operators/ops_signature/fetch_sig.cc
Copy file name to clipboard
+31Lines changed: 31 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+232-6Lines changed: 232 additions & 6 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+232-6Lines changed: 232 additions & 6 deletions
@@ -109,15 +109,29 @@ void FetchTensors(const std::vector<std::string>& job_fetch_names,
     int col = find(fetch_var_names.begin(), fetch_var_names.end(), var_name) -
               fetch_var_names.begin();
     auto* var = scope->FindVar(var_name);
-    auto& src = var->Get<phi::DenseTensor>();
-    auto* dst =
-        &(PADDLE_GET(phi::DenseTensor, fetch_list->at(micro_batch_id)[col]));
-    if (src.IsInitialized()) {
-      TensorCopy(src, phi::CPUPlace(), dst);
-      dst->set_lod(src.lod());
-    } else {
-      VLOG(6) << "Found " << var_name
-              << " is not initialized and skip TensorCopy.";
+    if (var->IsType<phi::DenseTensor>()) {
+      auto& src = var->Get<phi::DenseTensor>();
+      auto* dst =
+          &(PADDLE_GET(phi::DenseTensor, fetch_list->at(micro_batch_id)[col]));
+      if (src.IsInitialized()) {
+        TensorCopy(src, platform::CPUPlace(), dst);
+        dst->set_lod(src.lod());
+      } else {
+        VLOG(6) << "Found " << var_name
+                << " is not initialized and skip TensorCopy.";
+      }
+    } else if (var->IsType<phi::TensorArray>()) {
+      auto& src = var->Get<phi::TensorArray>();
+      fetch_list->at(micro_batch_id)[col] =
+          phi::TensorArray();  // default DenseTensor, we replace it with
+                               // TensorArray.
+      auto* dst =
+          &(PADDLE_GET(phi::TensorArray, fetch_list->at(micro_batch_id)[col]));
+      dst->resize(src.size());
+      for (size_t i = 0; i < src.size(); ++i) {
+        TensorCopy(src[i], platform::CPUPlace(), &dst->at(i));
+        dst->at(i).set_lod(src[i].lod());
+      }
     }
   }
 }
 
@@ -468,7 +468,13 @@ void HandleForSpecialOp(pir::Operation* op,
     auto fetch_var_name = fetch_src_name + "@fetch";
     auto* var = const_cast<Scope*>(value_exe_info->GetScope()->root())
                     ->Var(fetch_var_name);
-    var->GetMutable<phi::DenseTensor>();
+    if (op->operand_source(0)
+            .type()
+            .isa<paddle::dialect::DenseTensorArrayType>()) {
+      var->GetMutable<phi::TensorArray>();
+    } else {
+      var->GetMutable<phi::DenseTensor>();
+    }
     auto value = op->result(0);
 
     value_exe_info->Add(value, fetch_var_name);
 
@@ -0,0 +1,31 @@
+// 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
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/compat/op_utils.h"
+
+#include "glog/logging.h"
+
+namespace phi {
+
+KernelSignature FetchOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  if (ctx.IsDenseTensorVectorInput("X")) {
+    return KernelSignature("fetch_array", {"X"}, {"dst_place_type"}, {"Out"});
+  }
+
+  return KernelSignature("fetch", {"X"}, {"dst_place_type"}, {"Out"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(fetch, phi::FetchOpArgumentMapping);
@@ -188,6 +188,13 @@ pir::Value array_read(pir::Value array, pir::Value i) {
   return array_read_op.out();
 }
 
+pir::Value fetch(pir::Value value, std::string name, int col) {
+  auto fetch_op =
+      ApiBuilder::Instance().GetBuilder()->Build<paddle::dialect::FetchOp>(
+          value, name, col);
+  return fetch_op.out();
+}
+
 pir::Value array_write_(pir::Value array, pir::Value x, pir::Value i) {
   auto array_write_op =
       ApiBuilder::Instance()
 
@@ -74,6 +74,8 @@ pir::Value array_length(pir::Value x);
 
 pir::Value array_read(pir::Value array, pir::Value i);
 
+pir::Value fetch(pir::Value value, std::string name, int col);
+
 pir::Value array_write_(pir::Value array, pir::Value x, pir::Value i);
 
 std::tuple<pir::Value, pir::Value> array_to_tensor(pir::Value x,
 
@@ -19,12 +19,12 @@ paddle::dialect::AddN_Op, paddle::dialect::AddNArrayOp,
     paddle::dialect::ExpandOp, paddle::dialect::CreateArrayOp,
     paddle::dialect::CreateArrayLikeOp, paddle::dialect::ArrayLengthOp,
     paddle::dialect::ArrayReadOp, paddle::dialect::ArrayWrite_Op,
-    paddle::dialect::SliceArrayOp, paddle::dialect::SliceArrayDenseOp,
-    paddle::dialect::AssignArrayOp, paddle::dialect::AssignArray_Op,
-    paddle::dialect::ArrayToTensorOp, paddle::dialect::TensorToArrayOp,
-    paddle::dialect::IncrementOp, paddle::dialect::Increment_Op,
-    paddle::dialect::ShapeBroadcastOp, paddle::dialect::MemcpyD2hMultiIoOp,
-    paddle::dialect::ArrayPopOp
+    paddle::dialect::FetchOp, paddle::dialect::SliceArrayOp,
+    paddle::dialect::SliceArrayDenseOp, paddle::dialect::AssignArrayOp,
+    paddle::dialect::AssignArray_Op, paddle::dialect::ArrayToTensorOp,
+    paddle::dialect::TensorToArrayOp, paddle::dialect::IncrementOp,
+    paddle::dialect::Increment_Op, paddle::dialect::ShapeBroadcastOp,
+    paddle::dialect::MemcpyD2hMultiIoOp, paddle::dialect::ArrayPopOp
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
@@ -57,6 +57,231 @@ paddle::dialect::AddN_Op, paddle::dialect::AddNArrayOp,
 #endif
 
 namespace paddle::dialect {
+void FetchOp::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::UnchangedInferMetaIncludingTensorArray);
+  fn(infer_meta);
+}
+
+std::vector<pir::Type> FetchOp::InferMeta(
+    const std::vector<pir::Value> &input_values,
+    pir::AttributeMap *p_attributes) {
+  PADDLE_ENFORCE_NOT_NULL(
+      p_attributes,
+      common::errors::Fatal(
+          "AttrtibueMap pointer in InferMeta function is nullptr."));
+  auto &attributes = *p_attributes;
+  (void)attributes;
+
+  PADDLE_ENFORCE_EQ(input_values.size() == 1,
+                    true,
+                    phi::errors::InvalidArgument(
+                        "Num of inputs is expected to be 1 but got %d.",
+                        input_values.size()));
+
+  pir::Value x_ = input_values[0];
+  (void)x_;
+  VLOG(4) << "Builder construction outputs";
+  bool is_from_tensor = false;
+  (void)is_from_tensor;
+
+  std::vector<pir::Type> argument_outputs;
+  pir::Type out_type = x_.type();
+
+  // Auto Parallel condition
+#ifdef PADDLE_WITH_DISTRIBUTE
+  ProcessMeshAttribute op_mesh;
+  if (HasDistInput(input_values, &op_mesh)) {
+    CvtAllInputsToDist(input_values, op_mesh);
+    auto ctx = pir::IrContext::Instance();
+    std::vector<pir::Attribute> dist_operand_attrs, dist_result_attrs;
+    auto dist_meta_x =
+        CvtToDistMetaTensor(x_.type().dyn_cast<DistDenseTensorType>());
+    auto spmd_info =
+        phi::distributed::VariadicReplicatedInferSpmdDynamic(dist_meta_x);
+    PADDLE_ENFORCE_EQ(
+        spmd_info.first.size(),
+        1u,
+        common::errors::Unavailable(
+            "Size of spmd_info.first for op[FetchOp]is unexpected."));
+    for (auto &arg_dist : spmd_info.first) {
+      dist_operand_attrs.push_back(CvtToPirAttr(arg_dist));
+    }
+
+    auto dist_attr_out = CreateReplicatedDistAttr(out_type, op_mesh);
+
+    dist_result_attrs.push_back(dist_attr_out);
+    argument_outputs.push_back(CvtToPirDistType(out_type, dist_attr_out));
+
+    attributes[kAttrOpDistAttr] = OperationDistAttribute::get(
+        ctx, op_mesh, dist_operand_attrs, dist_result_attrs);
+    return argument_outputs;
+  }
+#endif
+
+  argument_outputs.push_back(out_type);
+
+  return argument_outputs;
+}
+
+const char *FetchOp::attributes_name[2] = {"name", "col"};
+
+OpInfoTuple FetchOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo(
+          "x", "paddle::dialect::DenseTensorType", false, false, false, false)};
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo("name", "pir::StrAttribute", ""),
+      paddle::dialect::OpAttributeInfo("col", "pir::Int32Attribute", "")};
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      paddle::dialect::OpOutputInfo(
+          "out", "paddle::dialect::DenseTensorType", false, false)};
+  paddle::dialect::OpRunTimeInfo run_time_info = paddle::dialect::OpRunTimeInfo(
+      "UnchangedInferMeta", {"x"}, "fetch", {"x"}, {}, {}, {}, {});
+  return std::make_tuple(inputs, attributes, outputs, run_time_info, "fetch");
+}
+
+void FetchOp::Build(pir::Builder &builder,
+                    pir::OperationArgument &argument,
+                    pir::Value x_,
+                    const std::string &name,
+                    int col) {
+  VLOG(4) << "Start build FetchOp";
+
+  VLOG(4) << "Builder construction inputs";
+  std::vector<pir::Value> argument_inputs = {x_};
+  argument.AddInputs(argument_inputs);
+
+  VLOG(4) << "Builder construction attributes";
+  pir::AttributeMap argument_attributes = {};
+  pir::Attribute attr_name =
+      pir::StrAttribute::get(pir::IrContext::Instance(), name);
+  argument_attributes.insert({"name", attr_name});
+  pir::Attribute attr_col =
+      pir::Int32Attribute::get(pir::IrContext::Instance(), col);
+  argument_attributes.insert({"col", attr_col});
+
+  std::vector<pir::Type> argument_outputs =
+      FetchOp::InferMeta(argument_inputs, &argument_attributes);
+  argument.AddAttributes(argument_attributes);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
+}
+
+void FetchOp::Build(pir::Builder &builder,
+                    pir::OperationArgument &argument,
+                    pir::Value x_,
+                    pir::AttributeMap attributes) {
+  VLOG(4) << "Start build FetchOp";
+
+  PADDLE_ENFORCE_NE(attributes.find("name"),
+                    attributes.end(),
+                    phi::errors::InvalidArgument(
+                        "'name' Attribute is expected for FetchOp. "));
+  std::string name =
+      attributes.at("name").dyn_cast<pir::StrAttribute>().AsString();
+
+  PADDLE_ENFORCE_NE(attributes.find("col"),
+                    attributes.end(),
+                    phi::errors::InvalidArgument(
+                        "'col' Attribute is expected for FetchOp. "));
+  int col = attributes.at("col").dyn_cast<pir::Int32Attribute>().data();
+
+  VLOG(4) << "Builder construction inputs";
+  std::vector<pir::Value> argument_inputs = {x_};
+  argument.AddInputs(argument_inputs);
+
+  VLOG(4) << "Builder construction attributes";
+  pir::AttributeMap argument_attributes = {};
+  pir::Attribute attr_name =
+      pir::StrAttribute::get(pir::IrContext::Instance(), name);
+  argument_attributes.insert({"name", attr_name});
+  pir::Attribute attr_col =
+      pir::Int32Attribute::get(pir::IrContext::Instance(), col);
+  argument_attributes.insert({"col", attr_col});
+
+  std::vector<pir::Type> argument_outputs =
+      FetchOp::InferMeta(argument_inputs, &argument_attributes);
+  argument.AddAttributes(argument_attributes);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
+}
+
+void FetchOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: FetchOp.";
+  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<paddle::dialect::DenseTensorType>() ||
+                          (*this)
+                              ->operand_source(0)
+                              .type()
+                              .isa<paddle::dialect::DenseTensorArrayType>(),
+                      true,
+                      phi::errors::InvalidArgument(
+                          "Type validation failed for the 0th input, got %s.",
+                          (*this)->operand_source(0).type()));
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    PADDLE_ENFORCE_GT(attributes.count("name"),
+                      0,
+                      phi::errors::InvalidArgument("name does not exist."));
+    PADDLE_ENFORCE_EQ(attributes.at("name").isa<pir::StrAttribute>(),
+                      true,
+                      phi::errors::InvalidArgument(
+                          "Type of attribute: name is not pir::StrAttribute."));
+
+    PADDLE_ENFORCE_GT(attributes.count("col"),
+                      0,
+                      phi::errors::InvalidArgument("col does not exist."));
+    PADDLE_ENFORCE_EQ(
+        attributes.at("col").isa<pir::Int32Attribute>(),
+        true,
+        phi::errors::InvalidArgument(
+            "Type of attribute: col is not pir::Int32Attribute."));
+  }
+  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."));
+    PADDLE_ENFORCE_EQ(
+        (*this)->result(0).type().isa<paddle::dialect::DenseTensorType>() ||
+            (*this)
+                ->result(0)
+                .type()
+                .isa<paddle::dialect::DenseTensorArrayType>(),
+        true,
+        phi::errors::InvalidArgument(
+            "Type validation failed for the 0th output."));
+  }
+  VLOG(4) << "End Verifying for: FetchOp.";
+}
+
+phi::DataType FetchOp::GetKernelTypeForVar(
+    const std::string &var_name,
+    const phi::DataType &tensor_dtype,
+    const phi::DataType &expected_kernel_dtype) {
+  VLOG(4) << "Get KernelType for Var of op: FetchOp";
+
+  return expected_kernel_dtype;
+}
+
+bool FetchOp::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  VLOG(4) << "Infer symbolic shape for op: FetchOp";
+  return FetchOpInferSymbolicShape(this->operation(), infer_context);
+}
+
 OpInfoTuple AddN_Op::GetOpInfo() {
   std::vector<paddle::dialect::OpInputInfo> inputs = {
       paddle::dialect::OpInputInfo(
@@ -4486,6 +4711,7 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AssignOut_Op)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::FusedGemmEpilogueOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::FusedGemmEpilogueGradOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::CreateArrayOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::FetchOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::CreateArrayLikeOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayLengthOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayReadOp)