PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_interface.h
Copy file name to clipboard
+26-3Lines changed: 26 additions & 3 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_interface.h
Copy file name to clipboard
+26-3Lines changed: 26 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_type.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_type.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_build_gen.py
Copy file name to clipboard
+15-20Lines changed: 15 additions & 20 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_build_gen.py
Copy file name to clipboard
+15-20Lines changed: 15 additions & 20 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+6-30Lines changed: 6 additions & 30 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+6-30Lines changed: 6 additions & 30 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
Copy file name to clipboard
+30-11Lines changed: 30 additions & 11 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
Copy file name to clipboard
+30-11Lines changed: 30 additions & 11 deletions
@@ -79,12 +79,12 @@ class TensorDistAttribute : public pir::AttrBase<TensorDistAttribute,
       pir::IrContext* ctx,
       ProcessMeshAttribute mesh,
       const std::vector<int64_t>& dims_mapping,
-      const flat_hash_map<int64_t, phi::ReduceType>& partial_status);
+      const flat_hash_map<int64_t, phi::ReduceType>& partial_status = {});
   static TensorDistAttribute get(
       pir::IrContext* ctx,
       const phi::distributed::ProcessMesh& mesh,
       const std::vector<int64_t>& dims_mapping,
-      const flat_hash_map<int64_t, phi::ReduceType>& partial_status) {
+      const flat_hash_map<int64_t, phi::ReduceType>& partial_status = {}) {
     return get(ctx,
                ProcessMeshAttribute::get(ctx, mesh),
                dims_mapping,
 
@@ -13,6 +13,7 @@
 // limitations under the License.
 #pragma once
 
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
 #include "paddle/pir/include/core/cast_utils.h"
 #include "paddle/pir/include/core/dll_decl.h"
 #include "paddle/pir/include/core/type.h"
@@ -25,24 +26,46 @@ class IR_API DistTypeInterface
  public:
   struct Concept {
     /// Defined these methods with the interface.
-    explicit Concept(pir::Type (*local_type)(pir::Type))
-        : local_type(local_type) {}
+    explicit Concept(pir::Type (*local_type)(pir::Type),
+                     ProcessMeshAttribute (*process_mesh_attr)(pir::Type),
+                     TensorDistAttribute (*tensor_dist_attr)(pir::Type))
+        : local_type(local_type),
+          process_mesh_attr(process_mesh_attr),
+          tensor_dist_attr(tensor_dist_attr) {}
     pir::Type (*local_type)(pir::Type);
+    ProcessMeshAttribute (*process_mesh_attr)(pir::Type);
+    TensorDistAttribute (*tensor_dist_attr)(pir::Type);
   };
 
   template <class ConcreteType>
   struct Model : public Concept {
     static Type local_type(Type type) {
       return pir::cast<ConcreteType>(type).local_type();
     }
-    Model() : Concept(local_type) {}
+    static ProcessMeshAttribute process_mesh_attr(Type type) {
+      return pir::cast<ConcreteType>(type).process_mesh_attr();
+    }
+
+    static TensorDistAttribute tensor_dist_attr(Type type) {
+      return pir::cast<ConcreteType>(type).tensor_dist_attr();
+    }
+
+    Model() : Concept(local_type, process_mesh_attr, tensor_dist_attr) {}
   };
 
   DistTypeInterface(pir::Type type, Concept *impl)
       : pir::TypeInterfaceBase<DistTypeInterface>(type), impl_(impl) {}
 
   pir::Type local_type() { return impl_->local_type(*this); }
 
+  ProcessMeshAttribute process_mesh_attr() {
+    return impl_->process_mesh_attr(*this);
+  }
+
+  TensorDistAttribute tensor_dist_attr() {
+    return impl_->tensor_dist_attr(*this);
+  }
+
  private:
   Concept *impl_;
 };
 
@@ -21,6 +21,7 @@
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/pir/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/builtin_op.h"
 #include "paddle/pir/include/core/ir_context.h"
 
 namespace paddle {
@@ -155,6 +156,7 @@ void ShardTensorOp::Build(pir::Builder& builder,
                                                 tensor_dist_attr,
                                                 local_shape);
   argument.AddOutput(out_dist_tensor_type);
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void ReShardOp::VerifySig() {
 
@@ -72,6 +72,16 @@ class DistDenseTensorType
         InferLocalDDim(dense_tensor_type.dims(), tensor_dist_attr);
     return get(ctx, dense_tensor_type, tensor_dist_attr, local_ddim);
   }
+
+  // return the replicated dist dense tensor type.
+  static DistDenseTensorType get(pir::IrContext* ctx,
+                                 pir::DenseTensorType dense_tensor_type,
+                                 ProcessMeshAttribute process_mesh_attr) {
+    auto& ddim = dense_tensor_type.dims();
+    auto attr = TensorDistAttribute::get(
+        ctx, process_mesh_attr, std::vector<int64_t>(ddim.size(), -1));
+    return get(ctx, dense_tensor_type, attr, ddim);
+  }
 };
 
 }  // namespace dialect
 
@@ -249,7 +249,8 @@ def GenBuildInputArgsStr(
 
 
 def GenBuildInsertFullForMutableAttribute(
-    op_class_name,
+    args,
+    op_info,
     op_attribute_name_list,
     op_attribute_build_arg_type_list,
     op_mutable_attribute_name_list,
@@ -757,10 +758,8 @@ def GenBuildOutputs(
 
 
 def gen_build_func_str(
-    op_class_name,
-    op_input_name_list,
-    op_input_type_list,
-    op_input_optional_list,
+    args,
+    op_info,
     op_attribute_name_list,
     op_attribute_type_list,
     op_attribute_build_arg_type_list,
@@ -771,18 +770,13 @@ def gen_build_func_str(
     op_non_mutable_attribute_type_list,
     op_non_mutable_attribute_build_arg_type_list,
     op_non_mutable_attribute_default_value_list,
-    op_output_name_list,
-    op_output_type_list,
-    op_output_size_list,
-    op_output_optional_list,
-    op_infer_meta_map,
-    op_inplace_map,
     muta_attr_is_input=False,
     attr_args_is_map=False,
 ):
+    op_input_name_list = op_info.input_name_list
     build_args_for_declare = ""
     build_func = ""
-    build_info_str = OP_INFO_TEMPLATE.format(op_name=op_class_name)
+    build_info_str = OP_INFO_TEMPLATE.format(op_name=op_info.class_name)
 
     build_args_for_declare = GenBuildInputArgsStr(
         op_input_name_list,
@@ -815,7 +809,8 @@ def gen_build_func_str(
     if not muta_attr_is_input:
         inset_full_for_mutable_attributes_str = (
             GenBuildInsertFullForMutableAttribute(
-                op_class_name,
+                args,
+                op_info,
                 op_attribute_name_list,
                 op_attribute_build_arg_type_list,
                 op_mutable_attribute_name_list,
@@ -836,7 +831,7 @@ def gen_build_func_str(
   argument.AddAttributes(argument_attributes);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
   ::pir::PassStopGradientsDefaultly(argument);""".format(
-        op_name=op_class_name
+        op_name=op_info.class_name
     )
 
     GET_ATTRIBUTES_FROM_MAP_TEMPLATE = """
@@ -912,7 +907,7 @@ def gen_build_func_str(
                     data_name = "AsString"
                 get_attributes_str += (
                     GET_ARRAY_ATTRIBUTE_FROM_MAP_TEMPLATE.format(
-                        op_name=op_class_name,
+                        op_name=op_info.class_name,
                         attr_type=attr_type,
                         attribute_name=attr_names[idx],
                         inner_type=inner_type,
@@ -922,38 +917,38 @@ def gen_build_func_str(
             elif "paddle::dialect::IntArrayAttribute" in attr_types[idx]:
                 get_attributes_str += (
                     GET_INTARRAY_ATTRIBUTE_FROM_MAP_TEMPLATE.format(
-                        op_name=op_class_name,
+                        op_name=op_info.class_name,
                         attr_type=attr_type,
                         attribute_name=attr_names[idx],
                     )
                 )
             elif "paddle::dialect::ScalarAttribute" in attr_types[idx]:
                 get_attributes_str += (
                     GET_SCALAR_ATTRIBUTE_FROM_MAP_TEMPLATE.format(
-                        op_name=op_class_name,
+                        op_name=op_info.class_name,
                         attr_type=attr_type,
                         attribute_name=attr_names[idx],
                     )
                 )
             elif "pir::StrAttribute" in attr_types[idx]:
                 get_attributes_str += (
                     GET_STR_ATTRIBUTES_FROM_MAP_TEMPLATE.format(
-                        op_name=op_class_name,
+                        op_name=op_info.class_name,
                         attr_type=attr_type,
                         attribute_name=attr_names[idx],
                         attr_ir_type=attr_types[idx],
                     )
                 )
             else:
                 get_attributes_str += GET_ATTRIBUTES_FROM_MAP_TEMPLATE.format(
-                    op_name=op_class_name,
+                    op_name=op_info.class_name,
                     attr_type=attr_type,
                     attribute_name=attr_names[idx],
                     attr_ir_type=attr_types[idx],
                 )
 
     build_func = OP_BUILD_TEMPLATE.format(
-        op_name=op_class_name,
+        op_name=op_info.class_name,
         build_info=build_info_str,
         build_args=build_args_for_define,
         build_mutable_attributes=inset_full_for_mutable_attributes_str,
 
@@ -1451,10 +1451,8 @@ def AutoCodeGen(
                         build_args_with_muta_attr_not_input_for_declare,
                         build_func_with_muta_attr_not_input,
                     ) = gen_build_func_str(
-                        op_class_name,
-                        op_input_name_list,
-                        op_input_type_list,
-                        op_input_optional_list,
+                        args,
+                        op_info,
                         op_attribute_name_list,
                         op_attribute_type_list,
                         op_attribute_build_arg_type_list,
@@ -1465,23 +1463,15 @@ def AutoCodeGen(
                         op_non_mutable_attribute_type_list,
                         op_non_mutable_attribute_build_arg_type_list,
                         op_non_mutable_attribute_default_value_list,
-                        op_output_name_list,
-                        op_output_type_list,
-                        op_output_size_list,
-                        op_output_optional_list,
-                        op_infer_meta_map,
-                        op_inplace_map,
                         muta_attr_is_input=False,
                     )
                     if len(op_attribute_name_list) > 0:
                         (
                             build_args_with_attr_is_map_for_declare,
                             build_func_with_attr_is_map,
                         ) = gen_build_func_str(
-                            op_class_name,
-                            op_input_name_list,
-                            op_input_type_list,
-                            op_input_optional_list,
+                            args,
+                            op_info,
                             op_attribute_name_list,
                             op_attribute_type_list,
                             op_attribute_build_arg_type_list,
@@ -1492,12 +1482,6 @@ def AutoCodeGen(
                             op_non_mutable_attribute_type_list,
                             op_non_mutable_attribute_build_arg_type_list,
                             op_non_mutable_attribute_default_value_list,
-                            op_output_name_list,
-                            op_output_type_list,
-                            op_output_size_list,
-                            op_output_optional_list,
-                            op_infer_meta_map,
-                            op_inplace_map,
                             muta_attr_is_input=False,
                             attr_args_is_map=True,
                         )
@@ -1508,10 +1492,8 @@ def AutoCodeGen(
                             build_args_with_muta_attr_is_input_for_declare,
                             build_func_with_muta_attr_is_input,
                         ) = gen_build_func_str(
-                            op_class_name,
-                            op_input_name_list,
-                            op_input_type_list,
-                            op_input_optional_list,
+                            args,
+                            op_info,
                             op_attribute_name_list,
                             op_attribute_type_list,
                             op_attribute_build_arg_type_list,
@@ -1522,12 +1504,6 @@ def AutoCodeGen(
                             op_non_mutable_attribute_type_list,
                             op_non_mutable_attribute_build_arg_type_list,
                             op_non_mutable_attribute_default_value_list,
-                            op_output_name_list,
-                            op_output_type_list,
-                            op_output_size_list,
-                            op_output_optional_list,
-                            op_infer_meta_map,
-                            op_inplace_map,
                             muta_attr_is_input=True,
                         )
 
 
@@ -94,11 +94,6 @@ def get_infermeta_inputs_str(
         # add mutable attributes as inputs
         if len(op_mutable_attribute_name_list) > 0:
             for i in range(len(op_mutable_attribute_name_list)):
-                if (
-                    op_mutable_attribute_name_list[i]
-                    not in inuse_infer_meta_args
-                ):
-                    continue
                 infermeta_inputs_str += CREATE_INPUT_VALUE_TEMPLATE.format(
                     input_name=op_mutable_attribute_name_list[i],
                     index=str(i + len(op_input_name_list)),
@@ -297,8 +292,6 @@ def GenBuildOutputsPart2(
     # Prepare mutable attributes
     if mutable_attr_is_input:
         for idx in range(len(op_mutable_attribute_name_list)):
-            if op_mutable_attribute_name_list[idx] not in inuse_infer_meta_args:
-                continue
             attr_dtype = op_mutable_attribute_type_list[idx]
             # int_array
             if attr_dtype[0] == "paddle::dialect::IntArrayAttribute":
@@ -617,13 +610,39 @@ def GenDistBranch(args, op_info):
     TEMPLATE = """
   // Auto Parallel condition
   if(HasDistInput(input_values)) {{
+    ProcessMeshAttribute op_mesh;
+    auto ctx = pir::IrContext::Instance();
+    for(auto value : input_values) {{
+      if (auto dist_interface = value.type().dyn_cast<DistTypeInterface>()) {{
+        op_mesh = dist_interface.process_mesh_attr();
+        break;
+      }}
+    }}"""
+    dist_branch_str = TEMPLATE.format()
+    TEMPLATE = """
+    if(!{name}.FromTensor()) {{
+      auto dist_type = DistDenseTensorType::get(ctx, {name}_.type().dyn_cast<DenseTensorType>(), op_mesh);
+      {name}_.set_type(dist_type);
+      {name}_.defining_op()->set_attribute(
+        kAttrOpDistAttr,
+          OperationDistAttribute::get(
+            ctx,
+            op_mesh,
+            {{dist_type.tensor_dist_attr() }},
+            {{}}
+          )
+      );
+    }}
+    """
+    for mutable_attr_name in op_info.mutable_attribute_name_list:
+        dist_branch_str += TEMPLATE.format(name=mutable_attr_name)
+    TEMPLATE = """
     if(!AllInputAreDist(input_values)) {{
         PADDLE_THROW(common::errors::Unimplemented(
             "Mixed inputs with DenseTensor and DistDenseTensor are not supported yet."));
     }}
-    ProcessMeshAttribute op_mesh = input_values[0].type().dyn_cast<DistDenseTensorType>().process_mesh_attr();
     std::vector<TensorDistAttribute> operand_dist_attrs, result_dist_attrs;"""
-    dist_branch_str = TEMPLATE.format()
+    dist_branch_str += TEMPLATE.format()
     infer_spmd_args_list = []
     # Prepare inputs_meta_tensor & attributes for infer spmd
     for name in op_info.spmd_params:
@@ -680,12 +699,12 @@ def GenDistBranch(args, op_info):
             TEMPLATE = """
     auto dist_attr_{name} = CvtToPirDistAttr(spmd_info.second[{idx}]);
     result_dist_attrs.push_back(dist_attr_{name});
-    argument_outputs.push_back(DistDenseTensorType::get(pir::IrContext::Instance(), {name}_type.dyn_cast<pir::DenseTensorType>(), dist_attr_{name}));
+    argument_outputs.push_back(DistDenseTensorType::get(ctx, {name}_type.dyn_cast<pir::DenseTensorType>(), dist_attr_{name}));
 """
             dist_branch_str += TEMPLATE.format(idx=idx, name=output_name)
     TEMPLATE = """
     attributes[kAttrOpDistAttr] = OperationDistAttribute::get(
-        pir::IrContext::Instance(),
+        ctx,
         op_mesh,
         operand_dist_attrs,
         result_dist_attrs