fix the bug in convert tensor to dist. (#68455)

winter-wang · web-flow · commit 3004bd5754fc · 2024-09-26T11:16:37.000+08:00
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc
@@ -305,7 +305,7 @@ void CvtAllInputsToDist(const std::vector<pir::Value>& inputs,
       if (auto vec_type = type.dyn_cast<pir::VectorType>()) {
         for (size_t idx = 0; idx < vec_type.size(); ++idx) {
           auto inner_type = vec_type[idx];
-          if (!inner_type && !vec_type[idx].isa<DistTypeInterface>()) {
+          if (inner_type && !inner_type.isa<DistTypeInterface>()) {
             CvtValueToDist(value, mesh_attr);
             break;
           }
diff --git a/paddle/fluid/pir/dialect/op_generator/api_gen.py b/paddle/fluid/pir/dialect/op_generator/api_gen.py
@@ -261,7 +261,7 @@
     }}"""
 
 COMBINE_OP_TEMPLATE = """
-    auto {op_name} = ApiBuilder::Instance().GetBuilder()->Build<pir::CombineOp>({in_name});"""
+    auto {op_name} = builtin_combine({in_name}).defining_op<pir::CombineOp>();"""
 
 SPLIT_OP_TEMPLATE = """
     auto {op_name} = ApiBuilder::Instance().GetBuilder()->Build<pir::SplitOp>({in_name});"""
diff --git a/paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
@@ -822,6 +822,7 @@ def GenDistBranch(args, op_info):
             if 'pir::VectorType' in op_info.input_type_list[input_index]:
                 TEMPLATE = """
     std::vector<phi::distributed::DistMetaTensor> vec_dist_meta_{name};
+    {name} = {name}_.type().dyn_cast<pir::VectorType>();
     for(auto& sub_ir_tensor: {name}.data()) {{
       vec_dist_meta_{name}.push_back(CvtToDistMetaTensor(sub_ir_tensor.dyn_cast<DistDenseTensorType>()));
     }}"""
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.cc b/paddle/fluid/pir/dialect/operator/ir/manual_api.cc
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/fluid/pir/dialect/operator/ir/manual_api.h"
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_tools.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
@@ -24,6 +25,11 @@
 namespace paddle::dialect {
 
 pir::Value builtin_combine(const std::vector<pir::Value>& x) {
+  // Auto Parallel condition
+  ProcessMeshAttribute op_mesh;
+  if (HasDistInput(x, &op_mesh)) {
+    CvtAllInputsToDist(x, op_mesh);
+  }
   auto combine_op =
       ApiBuilder::Instance().GetBuilder()->Build<pir::CombineOp>(x);
   return combine_op.out();

Original file line number	Diff line number	Diff line change
`@@ -305,7 +305,7 @@ void CvtAllInputsToDist(const std::vector<pir::Value>& inputs,`
`305`	`305`	`if (auto vec_type = type.dyn_cast<pir::VectorType>()) {`
`306`	`306`	`for (size_t idx = 0; idx < vec_type.size(); ++idx) {`
`307`	`307`	`auto inner_type = vec_type[idx];`
`308`		`- if (!inner_type && !vec_type[idx].isa<DistTypeInterface>()) {`
	`308`	`+ if (inner_type && !inner_type.isa<DistTypeInterface>()) {`
`309`	`309`	`CvtValueToDist(value, mesh_attr);`
`310`	`310`	`break;`
`311`	`311`	`}`