PaddlePaddle
diff --git a/‎paddle/fluid/distributed/collective/reducer.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/distributed/collective/reducer.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+10-1Lines changed: 10 additions & 1 deletion b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+10-1Lines changed: 10 additions & 1 deletion
diff --git a/‎paddle/fluid/framework/new_executor/pir_interpreter.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/framework/new_executor/pir_interpreter.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_api.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_api.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+4Lines changed: 4 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
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+11-4Lines changed: 11 additions & 4 deletions b/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+11-4Lines changed: 11 additions & 4 deletions
@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/distributed/collective/reducer.h"
 #include "paddle/common/flags.h"
+#include "paddle/fluid/pir/dialect/operator/ir/ir_tensor.h"
 #include "paddle/phi/api/lib/data_transform.h"
 #include "paddle/phi/backends/device_guard.h"
 #include "paddle/phi/backends/device_manager.h"
@@ -108,6 +109,8 @@ std::vector<std::vector<size_t>> Eager_AssignGroupBySize(
     if (var.is_dense_tensor()) {
       var_size =
           std::dynamic_pointer_cast<phi::DenseTensor>(var.impl())->numel();
+    } else if (dialect::IrTensor::classof(var.impl().get())) {
+      var_size = var.numel();
     } else {
       VLOG(3) << "var " << var.name()
               << " is not tensor or selected_rows, so skip it";
 
@@ -238,6 +238,7 @@ const std::unordered_set<std::string> SpecialOps = {
     paddle::dialect::PyLayerOp::name(),
     paddle::dialect::WhileOp::name(),
     pir::StackCreateOp::name(),
+    paddle::dialect::ShareVarOp::name(),
 };
 
 Variable* CreateVar(pir::Value value,
@@ -685,14 +686,22 @@ void HandleForSpecialOp(pir::Operation* op,
     auto outlet_value = stack_create_op.outlet();
     value_exe_info->AddValue2VarName(inlet_value, stack_var_name);
     value_exe_info->AddValue2VarName(outlet_value, stack_var_name);
+  } else if (op_name == "pd_op.share_var") {
+    VLOG(6) << "Handle for pd_op.share_var";
+    auto first_name =
+        value_exe_info->GetValue2VarName().at(op->operand_source(0));
+    for (size_t idx = 1u; idx < op->num_operands(); ++idx) {
+      value_exe_info->UpdateValue2VarName(op->operand_source(idx), first_name);
+    }
   }
 }
 
 bool IsNeedVarInplace(pir::Operation* op,
                       pir::Value value,
                       std::string op_name) {
   return (value.type().isa<paddle::dialect::DenseTensorArrayType>() ||
-          op_name == "pd_op.assign_value_" || op_name == "pd_op.assign_out_");
+          op_name == "pd_op.assign_value_" || op_name == "pd_op.assign_out_" ||
+          op_name == "pd_op.coalesce_tensor_");
 }
 
 // NOTE(chenxi67): Here, we only perform inplace processing for variables that
 
@@ -924,7 +924,7 @@ void PirInterpreter::BuildInstruction() {
         continue;
       }
       VLOG(6) << "process " << op_name;
-
+      if (op_name == "pd_op.share_var") continue;
       if (op.isa<paddle::dialect::LegacyKernelOp>()) {  // NOLINT
         CREATE_INSTR(LegacyKernelInstruction);
       } else {
 
@@ -17,6 +17,7 @@
 #include <vector>
 
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_op.h"
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
 
@@ -151,7 +151,7 @@ void ProcessDistBlock(pir::Block* block) {
                      .chunk_id();
       op_item->erase_attribute(kAttrOpDistAttr);
     }
-    op_item->set_attribute("chunk_id", pir::Int64Attribute::get(ctx, chunk_id));
+    op_item->set_attribute("chunk_id", pir::Int32Attribute::get(ctx, chunk_id));
 
     // TODO(2024-Q2) Handle other special dist op in future.
   }
 
@@ -368,4 +368,8 @@ std::vector<pir::Value> tensorrt_engine(
   return out_split_op.outputs();
 }
 
+pir::Operation* share_var(const std::vector<pir::Value>& x) {
+  return ApiBuilder::Instance().GetBuilder()->Build<ShareVarOp>(x);
+}
+
 }  // namespace paddle::dialect
@@ -112,5 +112,7 @@ std::vector<pir::Value> tensorrt_engine(
     std::vector<std::vector<int64_t>> outputs_shape,
     std::vector<phi::DataType> outputs_dtype,
     const std::string& converter_debug_info);
+
+pir::Operation* share_var(const std::vector<pir::Value>& x);
 }  // namespace dialect
 }  // namespace paddle
@@ -24,7 +24,8 @@ paddle::dialect::AddN_Op, paddle::dialect::AddNArrayOp,
     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::MemcpyD2hMultiIoOp, paddle::dialect::ArrayPopOp,
+    paddle::dialect::ShareVarOp
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
@@ -4724,4 +4725,5 @@ 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)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayPopOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ShareVarOp)
 #endif
@@ -842,6 +842,20 @@ class TEST_API ArrayPopOp
       pir::AttributeMap *p_attributes);
 };
 
+class ShareVarOp : public pir::Op<ShareVarOp> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.share_var"; }
+  static constexpr const char **attributes_name = nullptr;
+  static constexpr uint32_t attributes_num = 0;
+  TEST_API static void Build(pir::Builder &builder,             // NOLINT
+                             pir::OperationArgument &argument,  // NOLINT
+                             const std::vector<pir::Value> &inputs) {
+    argument.AddInputs(inputs);
+  }
+  void VerifySig() {}
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
@@ -869,3 +883,4 @@ 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)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayPopOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ShareVarOp)
@@ -3086,9 +3086,12 @@ void RemoveRedundantMemcpyAfterShadowFeed(pir::Block* block,
             all_use_is_scalar = false;
             break;
           }
-
-          auto op_info = ctx->GetRegisteredOpInfo(
-              use_op->dyn_cast<PhiKernelOp>().op_name());
+          auto op_name = use_op->dyn_cast<PhiKernelOp>().op_name();
+          if (op_name == "pd_op.share_var") {
+            all_use_is_scalar = false;
+            break;
+          }
+          auto op_info = ctx->GetRegisteredOpInfo(op_name);
 
           if (!op_info) {
             all_use_is_scalar = false;
@@ -3327,14 +3330,18 @@ void ProcessBlock(
       new_arg.set_attribute(name, attr);
     }
     if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
+      auto place_attr = arg.attribute<PlaceAttribute>("palce");
+      auto var_place = place_attr ? place_attr.data() : phi::Place();
       new_arg.set_type(
-          AllocatedDenseTensorType::get(ctx, phi::Place(), dense_tensor_type));
+          AllocatedDenseTensorType::get(ctx, var_place, dense_tensor_type));
     }
     (*map_value_pair)[arg] = new_arg;
   }
   if (phi::is_accelerat_place(place)) {
     for (auto& [keyword, arg] : block->kwargs()) {
       if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
+        auto place_attr = arg.attribute<PlaceAttribute>("palce");
+        if (place_attr && place_attr.data() == place) continue;
         auto dtype = dense_tensor_type.dtype();
         phi::KernelKey shadow_key{paddle::experimental::get_accelerat_backend(),
                                   phi::DataLayout::ANY,
Original file line number	Diff line number	Diff line change
`@@ -924,7 +924,7 @@ void PirInterpreter::BuildInstruction() {`
`924`	`924`	`continue;`
`925`	`925`	`}`
`926`	`926`	`VLOG(6) << "process " << op_name;`
`927`		`-`
	`927`	`+ if (op_name == "pd_op.share_var") continue;`
`928`	`928`	`if (op.isa<paddle::dialect::LegacyKernelOp>()) { // NOLINT`
`929`	`929`	`CREATE_INSTR(LegacyKernelInstruction);`
`930`	`930`	`} else {`
Original file line number	Diff line number	Diff line change
`@@ -151,7 +151,7 @@ void ProcessDistBlock(pir::Block* block) {`
`151`	`151`	`.chunk_id();`
`152`	`152`	`op_item->erase_attribute(kAttrOpDistAttr);`
`153`	`153`	`}`
`154`		`- op_item->set_attribute("chunk_id", pir::Int64Attribute::get(ctx, chunk_id));`
	`154`	`+ op_item->set_attribute("chunk_id", pir::Int32Attribute::get(ctx, chunk_id));`
`155`	`155`
`156`	`156`	`// TODO(2024-Q2) Handle other special dist op in future.`
`157`	`157`	`}`
Original file line number	Diff line number	Diff line change
`@@ -368,4 +368,8 @@ std::vector<pir::Value> tensorrt_engine(`
`368`	`368`	`return out_split_op.outputs();`
`369`	`369`	`}`
`370`	`370`
	`371`	`+pir::Operation* share_var(const std::vector<pir::Value>& x) {`
	`372`	`+ return ApiBuilder::Instance().GetBuilder()->Build<ShareVarOp>(x);`
	`373`	`+}`
	`374`	`+`
`371`	`375`	`} // namespace paddle::dialect`