special deal with the inplace case in auto parallel pass.

Author: winter-wang

paddle/fluid/pir/dialect/op_generator/ops_api_gen.py

@@ -158,7 +158,6 @@
     'c_allreduce_avg',
     'c_allreduce_max',
     'c_allreduce_min',
-    'c_allreduce_sum',
     'c_allreduce_prod',
     'c_embedding',
     'c_identity',

paddle/fluid/pir/dialect/operator/ir/manual_op.cc

@@ -3866,6 +3866,7 @@ void AssignOut_Op::Build(pir::Builder &builder,
   std::vector<pir::Type> argument_outputs =
       AssignOut_Op::InferMeta(argument_inputs, &argument_attributes);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  argument.AddAttributes(argument_attributes);
   constexpr char kStopGradientAttrName[] = "stop_gradient";
   auto stop_gradient0 =
       argument.inputs[0].attribute<pir::BoolAttribute>(kStopGradientAttrName);
@@ -3970,6 +3971,22 @@ std::vector<pir::Type> AssignOut_Op::InferMeta(
       dense_out.layout(),
       dense_out.lod(),
       dense_out.offset());
+#ifdef PADDLE_WITH_DISTRIBUTE
+  // Auto Parallel condition
+  if (auto dist_type = input_values[1].type().dyn_cast<DistTypeInterface>()) {
+    ProcessMeshAttribute op_mesh = dist_type.process_mesh_attr();
+    auto ctx = pir::IrContext::Instance();
+    std::vector<pir::Attribute> dist_operand_attrs{
+        dist_type.tensor_dist_attr(),
+        dist_type.tensor_dist_attr(),
+    },
+        dist_result_attrs{dist_type.tensor_dist_attr()};
+    argument_outputs.push_back(dist_type);
+    (*p_attributes)[kAttrOpDistAttr] = OperationDistAttribute::get(
+        ctx, op_mesh, dist_operand_attrs, dist_result_attrs);
+    return argument_outputs;
+  }
+#endif
   argument_outputs.push_back(out_dense_tensor_type);
 
   return argument_outputs;

python/paddle/distributed/auto_parallel/static/pir_pass.py

@@ -69,6 +69,50 @@ def apply_partition_pass(program):
         assert len(op.operands()) == len(
             op.dist_attr.operands()
         ), f"The number of operands and the number of op_dist_attr's operands are not equal in op: {op}"
+        assert len(op.results()) == len(
+            op.dist_attr.results()
+        ), f"The number of results and the number of op_dist_attr's results are not equal in op: {op}"
+        # deal with inplace value
+        for out_idx, in_idx in paddle.core.pir.get_op_inplace_info(op).items():
+            operand = op.operand(in_idx)
+            operand_attr = op.dist_attr.operand(in_idx)
+            prev_var = operand.source()
+            if not prev_var.is_dist() or operand_attr == prev_var.dist_attr():
+                continue
+            assert (
+                not prev_var.is_combine()
+            ), f"The current partition pass not support inplace value of {op} is tensor list."
+            operand_attr = operand_attr.as_tensor_dist_attr()
+            # reshard input
+            paddle.pir.set_insertion_point(op)
+            reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
+            operand.set_source(reshard_var)
+
+            result = op.result(out_idx)
+            result_attr = op.dist_attr.result(out_idx).as_tensor_dist_attr()
+            assert (
+                operand_attr == result_attr
+            ), f"For inplace value, The operend dist attr should be equal to result dist attr , please check your infer_spmd func of {op}"
+
+            # reshard output
+            paddle.pir.set_insertion_point_after(op)
+            old_dist_attr = result.dist_attr()
+            result.update_dist_attr(result_attr)
+
+            # reshard output to assign out input
+            reshard_var_1 = paddle._C_ops.reshard_v2(
+                result, prev_var.dist_attr()
+            )
+            paddle.assign(reshard_var_1, prev_var)
+
+            if old_dist_attr == result.dist_attr():
+                continue
+            reshard_var_2 = reshard_var_1
+            if old_dist_attr != reshard_var_1.dist_attr():
+                reshard_var_2 = paddle._C_ops.reshard_v2(result, old_dist_attr)
+            result.replace_all_uses_with(reshard_var_1)
+            reshard_var_1.get_defining_op().operand(0).set_source(result)
+            reshard_var_2.get_defining_op().operand(0).set_source(result)
 
         for operand, attr in zip(op.operands(), op.dist_attr.operands()):
             prev_var = operand.source()

python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_r_reshard_func.py

@@ -48,7 +48,7 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
             reduce_mean = True
 
         group = new_process_group(sorted(src_mesh.process_ids))
-        reduced_value = paddle._C_ops.c_allreduce_sum_(
+        reduced_value = paddle._C_ops.c_allreduce_sum(
             src_value, group.id, True, False
         )