Merge branch 'infer_split_with_num' of https://github.com/liuzhenhai9…

…3/Paddle into infer_split_with_num
PaddlePaddle · Nov 13, 2023 · 0297525 · 0297525
2 parents 45bdb49 + 33bd39a
commit 0297525
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 20 deletions.
diff --git a/paddle/phi/api/yaml/generator/dist_api_gen.py b/paddle/phi/api/yaml/generator/dist_api_gen.py
@@ -846,7 +846,7 @@ def generate_output_creation_code(self) -> str:
                 # SetKernelDistOutput arg
                 dist_output_arg = (
                     "spmd_info.second[0]"
-                    if self.generate_infer_spmd
+                    if self.infer_meta['spmd_rule'] is not None
                     else self.outputs['out_size_expr'][0]
                 )
                 output_creation_code += VECTOR_OUT_CREATION_TEMPLATE.format(
@@ -899,7 +899,7 @@ def generate_output_creation_code(self) -> str:
                     else:
                         dist_output_arg = (
                             f"spmd_info.second[{i}]"
-                            if self.generate_infer_spmd
+                            if self.infer_meta['spmd_rule'] is not None
                             else self.outputs['out_size_expr'][i]
                         )
                         output_creation_code += (

diff --git a/paddle/phi/infermeta/spmd_rules/concat.cc b/paddle/phi/infermeta/spmd_rules/concat.cc
@@ -17,6 +17,8 @@ limitations under the License. */
 #include <limits>
 #include <set>
 
+#include "glog/logging.h"
+
 #include "paddle/phi/infermeta/spmd_rules/elementwise.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
 
@@ -86,7 +88,11 @@ SpmdInfo ConcatInferSpmd(const std::vector<DistMetaTensor>& x, int axis) {
   AlignDimsSharding(
       &input_attrs, tensor_shapes, axis_names, {}, align_axis, true);
 
-  return {{input_attrs}, {input_attrs[non_empty_index]}};
+  auto out_dist_attr =
+      CopyTensorDistAttrForOutput(input_attrs[non_empty_index]);
+  out_dist_attr.set_dims_mapping(input_attrs[non_empty_index].dims_mapping());
+  VLOG(4) << "concat out " << out_dist_attr.to_string();
+  return {{input_attrs}, {out_dist_attr}};
 }
 
 SpmdInfo ConcatInferSpmdReverse(const std::vector<DistMetaTensor>& x,

diff --git a/paddle/phi/infermeta/spmd_rules/utils.cc b/paddle/phi/infermeta/spmd_rules/utils.cc
@@ -274,33 +274,27 @@ void AlignDimsSharding(std::vector<TensorDistAttr>* input_attrs_ptr,
 
   const auto& process_mess = input_attrs[non_empty_index].process_mesh();
   auto has_mismatch = [&](int32_t mesh_dim) {
-    bool mismatch = false;
     for (size_t i = 0; i < n_inputs; i++) {
       if (IsEmpty(tensor_shapes[i])) {
         continue;
       }
       auto& p_a = inputs_placements[non_empty_index][mesh_dim];
       auto& p_b = inputs_placements[i][mesh_dim];
-      if (!p_a->is_shard()) {
-        if (!PlacementEqual(p_a, p_b)) {
-          mismatch = true;
-          break;
+      if (p_a->is_shard() && p_b->is_shard()) {
+        auto a_shard = std::dynamic_pointer_cast<ShardStatus>(p_a);
+        auto b_shard = std::dynamic_pointer_cast<ShardStatus>(p_b);
+        auto a_axis = axis_names[non_empty_index][a_shard->get_axis()];
+        auto b_axis = axis_names[i][b_shard->get_axis()];
+        if (a_axis != b_axis) {
+          return true;
         }
       }
-      if (!p_b->is_shard()) {
-        mismatch = true;
-        break;
-      }
-      auto a_shard = std::dynamic_pointer_cast<ShardStatus>(p_a);
-      auto b_shard = std::dynamic_pointer_cast<ShardStatus>(p_b);
-      auto a_axis = axis_names[non_empty_index][a_shard->get_axis()];
-      auto b_axis = axis_names[i][b_shard->get_axis()];
-      if (a_axis != b_axis) {
-        mismatch = true;
-        break;
+
+      if (!PlacementEqual(p_a, p_b)) {
+        return true;
       }
     }
-    return mismatch;
+    return false;
   };
 
   // a dim can not be sharded twice along diffrent mesh_dim