add call to spmd in expand infermeta when use pir (#71214)

jeff41404 · web-flow · commit 496947d136cb · 2025-02-21T11:15:12.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -3695,6 +3695,45 @@ std::vector<pir::Type> ExpandOp::InferMeta(
       dense_out.layout(),
       dense_out.lod(),
       dense_out.offset());
+
+  // 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>());
+    // Todo(jeff41404): When expand adds spmd rules, synchronous modifications
+    // are required here.
+    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[ExpandOp]is unexpected."));
+    for (auto &arg_dist : spmd_info.first) {
+      dist_operand_attrs.push_back(CvtToPirAttr(arg_dist));
+    }
+
+    for (int i = 1; i < 2; ++i) {
+      dist_operand_attrs.push_back(GetTensorDistAttr(input_values[i].type()));
+    }
+
+    auto dist_attr_out =
+        CreateReplicatedDistAttr(out_dense_tensor_type, op_mesh);
+
+    dist_result_attrs.push_back(dist_attr_out);
+    argument_outputs.push_back(
+        CvtToPirDistType(out_dense_tensor_type, dist_attr_out));
+
+    (*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;
 }
