PaddlePaddle
diff --git a/‎paddle/fluid/framework/new_executor/instruction/instruction_util.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/fluid/framework/new_executor/instruction/instruction_util.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/fluid/ir_adaptor/translator/attribute_translator.cc
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions b/‎paddle/fluid/ir_adaptor/translator/attribute_translator.cc
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions
diff --git a/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+27-18Lines changed: 27 additions & 18 deletions b/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+27-18Lines changed: 27 additions & 18 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+86-2Lines changed: 86 additions & 2 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+86-2Lines changed: 86 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+35Lines changed: 35 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+35Lines changed: 35 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+6-2Lines changed: 6 additions & 2 deletions b/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+6-2Lines changed: 6 additions & 2 deletions
diff --git a/‎paddle/pir/core/builtin_op.cc
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions b/‎paddle/pir/core/builtin_op.cc
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions
diff --git a/‎paddle/pir/core/dialect.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/pir/core/dialect.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/pir/core/interface_support.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/pir/core/interface_support.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
@@ -181,7 +181,9 @@ OpFuncType AnalyseOpFuncType(pir::Operation* op, const platform::Place& place) {
       return OpFuncType::kGpuSync;
     }
 
-    if (platform::is_gpu_place(place) && op_name == "pd_op.memcpy_d2h") {
+    if (platform::is_gpu_place(place) &&
+        (op_name == "pd_op.memcpy_d2h" ||
+         op_name == "pd_op.memcpy_d2h_multi_io")) {
       return OpFuncType::kGpuSync;
     }
 
 
@@ -187,6 +187,21 @@ class Int64AttributeVisitor : public AttributeVisitor {
   }
 };
 
+class Int32ArrayAttributeVisitor : public AttributeVisitor {
+ public:
+  using AttributeVisitor::AttributeVisitor;
+
+  pir::Attribute operator()(const std::vector<int64_t>& i64s) override {
+    VLOG(10) << "translating vector<int64> size: " << i64s.size();
+    std::vector<pir::Attribute> attrs;
+    attrs.reserve(i64s.size());
+    for (const auto& v : i64s) {
+      attrs.push_back(pir::Int32Attribute::get(ctx, v));
+    }
+    return pir::ArrayAttribute::get(ctx, attrs);
+  }
+};
+
 class IntArrayAttributeVisitor : public AttributeVisitor {
  public:
   using AttributeVisitor::AttributeVisitor;
@@ -240,6 +255,8 @@ AttributeTranslator::AttributeTranslator() {
       new PlaceAttributeVisitor();
   special_visitors["pir::ArrayAttribute<pir::Int64Attribute>"] =
       new Int64ArrayAttributeVisitor();
+  special_visitors["pir::ArrayAttribute<pir::Int32Attribute>"] =
+      new Int32ArrayAttributeVisitor();
   special_visitors["pir::Int64Attribute"] = new Int64AttributeVisitor();
 }
 
 
@@ -19,6 +19,7 @@
 #include <numeric>
 #include <string>
 #include <tuple>
+#include <typeinfo>
 #include <unordered_map>
 #include <vector>
 
@@ -37,12 +38,16 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/phi/core/utils/data_type.h"
+#include "paddle/pir/core/attribute.h"
 #include "paddle/pir/core/builder.h"
+#include "paddle/pir/core/builtin_attribute.h"
 #include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/core/builtin_type.h"
 #include "paddle/pir/core/ir_context.h"
 #include "paddle/pir/core/operation.h"
+#include "paddle/pir/core/utils.h"
 #include "paddle/pir/core/value.h"
+#include "paddle/utils/blank.h"
 
 #ifdef PADDLE_WITH_DNNL
 #include "paddle/fluid/pir/dialect/operator/ir/pd_onednn_op.h"
@@ -982,25 +987,28 @@ struct AssignValueOpTranscriber : public OpTranscriber {
         ctx, phi::Place(phi::AllocationType::UNDEFINED));
     attribute_map["place"] = attr_place;
 
-    if (op_desc.HasAttr("bool_values")) {
-      legacy_attr = op_desc.GetAttr("bool_values");
-    } else if (op_desc.HasAttr("fp32_values")) {
-      legacy_attr = op_desc.GetAttr("fp32_values");
-    } else if (op_desc.HasAttr("int32_values")) {
-      legacy_attr = op_desc.GetAttr("int32_values");
-    } else if (op_desc.HasAttr("int64_values")) {
-      legacy_attr = op_desc.GetAttr("int64_values");
-    } else if (op_desc.HasAttr("values")) {
-      legacy_attr = op_desc.GetAttr("values");
-    } else {
-      IR_THROW(
-          "Op assign_value should have attribute `**_values` or `values` but "
-          "not find");
+    const std::vector<std::string> possible_attrs = {
+        "bool_values", "fp32_values", "int32_values", "int64_values", "values"};
+    for (const auto& attr_name : possible_attrs) {
+      if (!op_desc.HasAttr(attr_name)) {
+        continue;
+      }
+      legacy_attr = op_desc.GetAttr(attr_name);
+      pir::Attribute attr_values = attribute_translator(
+          attr_info_maps.at("values").type_name, legacy_attr);
+      if (attr_values && attr_values.isa<pir::ArrayAttribute>() &&
+          !attr_values.dyn_cast<pir::ArrayAttribute>().empty()) {
+        attribute_map["values"] = attr_values;
+        VLOG(10) << "[op assign_value][values]" << attr_name << " "
+                 << attr_values;
+        break;
+      }
     }
 
-    pir::Attribute attr_values = attribute_translator(
-        attr_info_maps.at("values").type_name, legacy_attr);
-    attribute_map["values"] = attr_values;
+    IR_ENFORCE(
+        attribute_map.find("values") != attribute_map.end(),
+        "Op assign_value should have attribute `**_values` or `values` but "
+        "not find");
 
     TranslateOpDistAttribute(op_desc, &attribute_map);
 
@@ -2062,7 +2070,8 @@ struct SelectInputOpTranscriber : public OpTranscriber {
         auto& attribute_translator = AttributeTranslator::instance();
         undefine_value.defining_op()->set_attribute(
             "shape",
-            attribute_translator(common::vectorize(undefined_var_type.dims())));
+            attribute_translator("pir::ArrayAttribute<pir::Int32Attribute>",
+                                 common::vectorize(undefined_var_type.dims())));
       }
       auto dim1 = input1.dyn_cast<paddle::dialect::DenseTensorType>().dims();
       auto dim2 = input2.dyn_cast<paddle::dialect::DenseTensorType>().dims();
 
@@ -83,6 +83,7 @@
     'generate_sequence_xpu',
     'layer_norm_act_xpu',
     'memcpy',
+    'memcpy_d2h_multi_io',
     'batch_norm_',
     'multi_encoder_xpu',
     'multihead_matmul',
 
@@ -23,7 +23,8 @@ paddle::dialect::AddNOp, paddle::dialect::AddN_Op,
     paddle::dialect::SliceArrayOp, paddle::dialect::SliceArrayDenseOp,
     paddle::dialect::AssignArray_Op, paddle::dialect::ArrayToTensorOp,
     paddle::dialect::SelectInputOp, paddle::dialect::IncrementOp,
-    paddle::dialect::Increment_Op, paddle::dialect::ShapeBroadcastOp
+    paddle::dialect::Increment_Op, paddle::dialect::ShapeBroadcastOp,
+    paddle::dialect::MemcpyD2hMultiIoOp
 #else
 
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
@@ -1686,7 +1687,7 @@ OpInfoTuple ArrayWrite_Op::GetOpInfo() {
                     "array_write",
                     {"array", "x", "i"},
                     {"array"},
-                    {},
+                    {"array"},
                     {{"out", "array"}},
                     {});
 
@@ -3034,6 +3035,88 @@ bool ShapeBroadcastOp::InferSymbolicShape(
   return true;
 }
 
+const char *MemcpyD2hMultiIoOp::attributes_name[1] = {"dst_place_type"};
+
+OpInfoTuple MemcpyD2hMultiIoOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo("x",
+                                   "paddle::dialect::DenseTensorArrayType",
+                                   false,
+                                   false,
+                                   false,
+                                   false)};
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo(
+          "dst_place_type", "pir::Int32Attribute", "")};
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      paddle::dialect::OpOutputInfo(
+          "out", "paddle::dialect::DenseTensorArrayType", false, false)};
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      paddle::dialect::OpRunTimeInfo("UnchangedInferMeta",
+                                     {"x"},
+                                     "memcpy_d2h_multi_io",
+                                     {"x", "dst_place_type"},
+                                     {},
+                                     {},
+                                     {},
+                                     {});
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "memcpy_d2h_multi_io");
+}
+
+void MemcpyD2hMultiIoOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: "
+             "MemcpyD2hMultiIoOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    IR_ENFORCE(input_size == 1u,
+               "The size %d of inputs must be equal to 1.",
+               input_size);
+
+    IR_ENFORCE((*this)
+                   ->operand_source(0)
+                   .type()
+                   .isa<paddle::dialect::DenseTensorArrayType>(),
+               "Type validation failed for the 0th input, got %s.",
+               (*this)->operand_source(0).type());
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    IR_ENFORCE(attributes.count("dst_place_type") > 0,
+               "dst_place_type does not exist.");
+    IR_ENFORCE(attributes.at("dst_place_type").isa<pir::Int32Attribute>(),
+               "Type of attribute: dst_place_type is not pir::Int32Attribute.");
+  }
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    IR_ENFORCE(output_size == 1u,
+               "The size %d of outputs must be equal to 1.",
+               output_size);
+    auto output_0_type = (*this)->result(0).type();
+
+    IR_ENFORCE(output_0_type.isa<paddle::dialect::DenseTensorArrayType>(),
+               "Type validation failed for the 0th output.");
+  }
+  VLOG(4) << "End Verifying for: MemcpyD2hMultiIoOp.";
+}
+
+void MemcpyD2hMultiIoOp::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::UnchangedArrayInferMeta);
+  fn(infer_meta);
+}
+
+phi::DataType MemcpyD2hMultiIoOp::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: MemcpyD2hMultiIoOp";
+
+  return expected_kernel_dtype;
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
@@ -3057,5 +3140,6 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::Increment_Op)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::MemcpyD2hMultiIoOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ShapeBroadcastOp)
 #endif
@@ -555,6 +555,40 @@ class Increment_Op
       const std::vector<std::vector<bool>> &stop_gradients);
 };
 
+class MemcpyD2hMultiIoOp
+    : public pir::Op<MemcpyD2hMultiIoOp,
+                     paddle::dialect::OpYamlInfoInterface,
+                     paddle::dialect::InferMetaInterface,
+                     paddle::dialect::GetKernelTypeForVarInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.memcpy_d2h_multi_io"; }
+  static const char *attributes_name[1];
+  static constexpr uint32_t attributes_num = 1;
+  static OpInfoTuple GetOpInfo();
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x_,
+                    int dst_place_type);
+
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x_,
+                    pir::AttributeMap attributes);
+
+  void VerifySig();
+
+  static phi::DataType GetKernelTypeForVar(
+      const std::string &var_name,
+      const phi::DataType &tensor_dtype,
+      const phi::DataType &expected_kernel_dtype);
+
+  pir::Value x() { return operand_source(0); }
+  pir::OpResult out() { return result(0); }
+
+  static void InferMeta(phi::InferMetaContext *infer_meta);
+};
+
 class IR_API ShapeBroadcastOp
     : public pir::Op<ShapeBroadcastOp,
                      paddle::dialect::InferSymbolicShapeInterface> {
@@ -601,3 +635,4 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::Increment_Op)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ShapeBroadcastOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::MemcpyD2hMultiIoOp)
@@ -325,9 +325,13 @@ static pir::OpResult AddPlaceTransferOp(pir::Value in,
   } else if ((src_place.GetType() == phi::AllocationType::GPU) &&
              (dst_place.GetType() == phi::AllocationType::CPU)) {
     copy_kernel_key.set_backend(phi::Backend::GPU);
+    std::string copy_kernel_name = "memcpy_d2h";
+    if (in.type().isa<AllocatedDenseTensorArrayType>()) {
+      copy_kernel_name = "memcpy_d2h_multi_io";
+    }
     op_attribute = {
-        {"op_name", pir::StrAttribute::get(ctx, "pd_op.memcpy_d2h")},
-        {"kernel_name", pir::StrAttribute::get(ctx, "memcpy_d2h")},
+        {"op_name", pir::StrAttribute::get(ctx, "pd_op." + copy_kernel_name)},
+        {"kernel_name", pir::StrAttribute::get(ctx, copy_kernel_name)},
         {"kernel_key", KernelAttribute::get(ctx, copy_kernel_key)},
         {"dst_place_type", pir::Int32Attribute::get(ctx, 0)}};
   } else {
 
@@ -22,7 +22,7 @@ namespace pir {
 const char *ModuleOp::attributes_name[attributes_num] = {"program"};  // NOLINT
 
 void PassStopGradientsDefaultly(OperationArgument &argument) {  // NOLINT
-  VLOG(4) << "Builder construction stop gradient for OpResults.";
+  VLOG(10) << "Builder construction stop gradient for OpResults.";
   bool stop_gradient = true;
   for (auto value : argument.inputs) {
     auto attr = value.attribute<BoolAttribute>(kStopGradientAttrName);
@@ -90,7 +90,7 @@ void ModuleOp::Destroy() {
 }
 
 void ModuleOp::VerifySig() const {
-  VLOG(4) << "Verifying inputs, outputs and attributes for: ModuleOp.";
+  VLOG(10) << "Verifying inputs, outputs and attributes for: ModuleOp.";
   // Verify inputs:
   IR_ENFORCE(num_operands() == 0u, "The size of inputs must be equal to 0.");
 
@@ -129,7 +129,7 @@ std::string ParameterOp::param_name() const {
   return attribute<StrAttribute>("parameter_name").AsString();
 }
 void ParameterOp::VerifySig() const {
-  VLOG(4) << "Verifying inputs, outputs and attributes for: ParameterOp.";
+  VLOG(10) << "Verifying inputs, outputs and attributes for: ParameterOp.";
   // Verify inputs:
   IR_ENFORCE(num_operands() == 0u, "The size of inputs must be equal to 0.");
 
@@ -155,7 +155,7 @@ void SetParameterOp::Build(Builder &builder,             // NOLINT
                         pir::StrAttribute::get(builder.ir_context(), name));
 }
 void SetParameterOp::VerifySig() const {
-  VLOG(4) << "Verifying inputs, outputs and attributes for: SetParameterOp.";
+  VLOG(10) << "Verifying inputs, outputs and attributes for: SetParameterOp.";
   // Verify inputs:
   IR_ENFORCE(num_operands() == 1, "The size of outputs must be equal to 1.");
 
@@ -181,7 +181,7 @@ void ShadowOutputOp::Build(Builder &builder,             // NOLINT
                         pir::StrAttribute::get(builder.ir_context(), name));
 }
 void ShadowOutputOp::VerifySig() const {
-  VLOG(4) << "Verifying inputs, outputs and attributes for: ShadowOutputOp.";
+  VLOG(10) << "Verifying inputs, outputs and attributes for: ShadowOutputOp.";
   // Verify inputs:
   IR_ENFORCE(num_operands() == 1, "The size of outputs must be equal to 1.");
 
 
@@ -21,7 +21,7 @@ Dialect::Dialect(std::string name, pir::IrContext *context, pir::TypeId id)
 Dialect::~Dialect() = default;
 
 void Dialect::RegisterInterface(std::unique_ptr<DialectInterface> interface) {
-  VLOG(4) << "Register interface into dialect" << std::endl;
+  VLOG(9) << "Register interface into dialect" << std::endl;
   registered_interfaces_.emplace(interface->interface_id(),
                                  std::move(interface));
 }
 
@@ -45,14 +45,14 @@ class ConstructInterfacesOrTraits {
     IR_ENFORCE(suceess,
                "Interface: id[%u] is already registered. inset failed",
                TypeId::get<T>());
-    VLOG(6) << "New a interface: id[" << TypeId::get<T>() << "].";
+    VLOG(10) << "New a interface: id[" << TypeId::get<T>() << "].";
   }
 
   /// Placement new trait.
   template <typename T>
   static void PlacementConstrctTrait(pir::TypeId *&p_trait) {  // NOLINT
     *p_trait = TypeId::get<T>();
-    VLOG(6) << "New a trait: id[" << *p_trait << "].";
+    VLOG(10) << "New a trait: id[" << *p_trait << "].";
     ++p_trait;
   }
 };
Original file line number	Diff line number	Diff line change
`@@ -181,7 +181,9 @@ OpFuncType AnalyseOpFuncType(pir::Operation* op, const platform::Place& place) {`
`181`	`181`	`return OpFuncType::kGpuSync;`
`182`	`182`	`}`
`183`	`183`
`184`		`- if (platform::is_gpu_place(place) && op_name == "pd_op.memcpy_d2h") {`
	`184`	`+ if (platform::is_gpu_place(place) &&`
	`185`	`+ (op_name == "pd_op.memcpy_d2h" \|\|`
	`186`	`+ op_name == "pd_op.memcpy_d2h_multi_io")) {`
`185`	`187`	`return OpFuncType::kGpuSync;`
`186`	`188`	`}`
`187`	`189`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ Dialect::Dialect(std::string name, pir::IrContext *context, pir::TypeId id)`
`21`	`21`	`Dialect::~Dialect() = default;`
`22`	`22`
`23`	`23`	`void Dialect::RegisterInterface(std::unique_ptr<DialectInterface> interface) {`
`24`		`- VLOG(4) << "Register interface into dialect" << std::endl;`
	`24`	`+ VLOG(9) << "Register interface into dialect" << std::endl;`
`25`	`25`	`registered_interfaces_.emplace(interface->interface_id(),`
`26`	`26`	`std::move(interface));`
`27`	`27`	`}`