Fix bug caused by split infershape (#40116)

* fix bug caused by split infershape

* revert infer_shape of split

* revert split

paddle/fluid/operators/split_op.cc

@@ -26,6 +26,52 @@ class SplitOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
 
+  void InferShape(framework::InferShapeContext *ctx) const override {
+    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
+                      platform::errors::InvalidArgument(
+                          "Input(X) of SplitOp should not be null."));
+    PADDLE_ENFORCE_GE(ctx->Outputs("Out").size(), 1UL,
+                      platform::errors::InvalidArgument(
+                          "Outputs(Out) of SplitOp should not be empty."));
+    auto in_dims = ctx->GetInputDim("X");
+    auto outs_names = ctx->Outputs("Out");
+    size_t axis = static_cast<size_t>(ctx->Attrs().Get<int>("axis"));
+    size_t num = static_cast<size_t>(ctx->Attrs().Get<int>("num"));
+    std::vector<int> sections = static_cast<std::vector<int>>(
+        ctx->Attrs().Get<std::vector<int>>("sections"));
+    const size_t outs_number = outs_names.size();
+
+    if (sections.size() > 0) {
+      PADDLE_ENFORCE_EQ(
+          sections.size(), outs_number,
+          platform::errors::InvalidArgument("tensor split sections size "
+                                            "should be equal to output size."));
+    }
+
+    if (ctx->HasInput("AxisTensor")) {
+      auto out_dims = phi::make_ddim(std::vector<int>(in_dims.size(), -1));
+      std::vector<framework::DDim> outs_dims(outs_number, out_dims);
+      ctx->SetOutputsDim("Out", outs_dims);
+      for (size_t i = 0; i < outs_number; ++i) {
+        ctx->ShareLoD("X", "Out", 0, i);
+      }
+      return;
+    }
+
+    bool each_section_is_known =
+        (sections.size() > 0 && !ctx->HasInputs("SectionsTensorList"));
+
+    auto outs_dims = UpdateOutsDims(ctx->IsRuntime(), each_section_is_known,
+                                    in_dims, num, sections, axis, outs_number);
+    ctx->SetOutputsDim("Out", outs_dims);
+    if (axis != 0) {
+      // Only pass LoD when not spliting along the first dim.
+      for (size_t i = 0; i < outs_number; ++i) {
+        ctx->ShareLoD("X", "Out", 0, i);
+      }
+    }
+  }
+
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
@@ -125,10 +171,6 @@ Example:
 
 namespace ops = paddle::operators;
 
-DELCARE_INFER_SHAPE_FUNCTOR(split, SplitInferShapeFunctor,
-                            PT_INFER_META(phi::SplitInferMeta));
-
 REGISTER_OPERATOR(split, ops::SplitOp, ops::SplitOpMaker,
                   ops::SplitGradMaker<paddle::framework::OpDesc>,
-                  ops::SplitGradMaker<paddle::imperative::OpBase>,
-                  SplitInferShapeFunctor);
+                  ops::SplitGradMaker<paddle::imperative::OpBase>);

paddle/phi/infermeta/unary.cc

@@ -508,17 +508,6 @@ void SplitInferMeta(const MetaTensor& x,
                     const Scalar& axis,
                     std::vector<MetaTensor*> out,
                     MetaConfig config) {
-  if (!config.is_runtime) {
-    if (axis.FromTensor() || num_or_sections.FromTensor()) {
-      auto out_dims = phi::make_ddim(std::vector<int>(x.dims().size(), -1));
-      for (auto* item : out) {
-        item->set_dims(out_dims);
-        item->share_lod(x);
-      }
-      return;
-    }
-  }
-
   int axis_value = axis.to<int>();
   int rank = x.dims().size();
   PADDLE_ENFORCE_EQ(
@@ -533,16 +522,15 @@ void SplitInferMeta(const MetaTensor& x,
     axis_value = axis_value + rank;
   }
 
-  std::vector<phi::DDim> out_dims(out.size(), x.dims());
-
   auto input_axis_dim = x.dims().at(axis_value);
   auto num_or_sections_data = num_or_sections.GetData();
+  // step1: get formated sections
+  std::vector<int64_t> sections;
   // num_or_sections is a number
   if (num_or_sections_data.size() == 1) {
-    if (config.is_runtime || input_axis_dim > 0) {
     int num = num_or_sections_data.at(0);
-      PADDLE_ENFORCE_EQ(
-          input_axis_dim % num,
+
+    PADDLE_ENFORCE_EQ(input_axis_dim % num,
                       0,
                       phi::errors::InvalidArgument(
                           "The input's size along the split dimension "
@@ -553,14 +541,8 @@ void SplitInferMeta(const MetaTensor& x,
                           x.dims(),
                           axis_value));
 
-      size_t out_axis_dim = input_axis_dim / num;
-      for (auto& out_dim : out_dims) {
-        out_dim[axis_value] = out_axis_dim;
-      }
-    } else {
-      for (auto& out_dim : out_dims) {
-        out_dim[axis_value] = -1;
-      }
+    for (int i = 0; i < num; ++i) {
+      sections.push_back(input_axis_dim / num);
     }
   } else {
     // num_or_sections is a sections
@@ -568,9 +550,10 @@ void SplitInferMeta(const MetaTensor& x,
     int unknow_dim_idx = -1;
     int num_of_unknow = 0;
     int sum_of_section = 0;
-    std::vector<int64_t> sections = num_or_sections_data;
 
     for (size_t i = 0; i < num_or_sections_data.size(); ++i) {
+      sections.push_back(num_or_sections_data[i]);
+
       if (num_or_sections_data[i] == unknow_dim_val) {
         num_of_unknow++;
         unknow_dim_idx = i;
@@ -622,22 +605,31 @@ void SplitInferMeta(const MetaTensor& x,
               x.dims(),
               axis_value));
     }
-    for (size_t i = 0; i < out_dims.size(); ++i) {
+  }
+
+  // setp2: fill out dims
+  std::vector<phi::DDim> out_dims(sections.size(), x.dims());
+  if (config.is_runtime || input_axis_dim > 0) {
+    for (size_t i = 0; i < sections.size(); ++i) {
       out_dims[i][axis_value] = sections[i];
     }
+  } else {
+    for (size_t i = 0; i < sections.size(); ++i) {
+      out_dims[i][axis_value] = -1;
+    }
   }
 
-  for (size_t i = 0; i < out.size(); ++i) {
+  for (size_t i = 0; i < sections.size(); ++i) {
     if (axis_value != 0) {
       // Only pass LoD when not spliting along the first dim.
-      out.at(i)->set_dtype(x.dtype());
-      out.at(i)->set_dims(out_dims[i]);
-      out.at(i)->set_layout(x.layout());
+      out[i]->set_dtype(x.dtype());
+      out[i]->set_dims(out_dims[i]);
+      out[i]->set_layout(x.layout());
     } else {
-      out.at(i)->set_dtype(x.dtype());
-      out.at(i)->set_dims(out_dims[i]);
-      out.at(i)->set_layout(x.layout());
-      out.at(i)->share_lod(x);
+      out[i]->set_dtype(x.dtype());
+      out[i]->set_dims(out_dims[i]);
+      out[i]->set_layout(x.layout());
+      out[i]->share_lod(x);
     }
   }
 }

paddle/phi/kernels/cpu/split_kernel.cc

@@ -28,6 +28,23 @@ void SplitKernel(const Context& dev_ctx,
                  const ScalarArray& num_or_sections,
                  const Scalar& axis_scalar,
                  std::vector<DenseTensor*> outs) {
+  // need to infershape output
+  if (num_or_sections.FromTensor() || axis_scalar.FromTensor()) {
+    std::vector<MetaTensor> out_metas;
+    out_metas.reserve(outs.size());
+    std::vector<MetaTensor*> out_metas_ptr;
+    for (size_t i = 0; i < outs.size(); ++i) {
+      out_metas.push_back(outs[i]);
+      out_metas_ptr.push_back(&out_metas.back());
+    }
+
+    phi::SplitInferMeta(x, num_or_sections, axis_scalar, out_metas_ptr, true);
+
+    for (size_t i = 0; i < out_metas.size(); ++i) {
+      outs[i]->Resize(out_metas[i].dims());
+    }
+  }
+
   std::vector<const DenseTensor*> shape_refer;
   for (size_t j = 0; j < outs.size(); ++j) {
     dev_ctx.template Alloc<T>(outs[j]);

paddle/phi/kernels/gpu/split_kernel.cu

@@ -27,6 +27,23 @@ void SplitKernel(const Context& dev_ctx,
                  const ScalarArray& num_or_sections,
                  const Scalar& axis_scalar,
                  std::vector<DenseTensor*> outs) {
+  // need to infershape output
+  if (num_or_sections.FromTensor() || axis_scalar.FromTensor()) {
+    std::vector<MetaTensor> out_metas;
+    out_metas.reserve(outs.size());
+    std::vector<MetaTensor*> out_metas_ptr;
+    for (size_t i = 0; i < outs.size(); ++i) {
+      out_metas.push_back(outs[i]);
+      out_metas_ptr.push_back(&out_metas.back());
+    }
+
+    phi::SplitInferMeta(x, num_or_sections, axis_scalar, out_metas_ptr, true);
+
+    for (size_t i = 0; i < out_metas.size(); ++i) {
+      outs[i]->Resize(out_metas[i].dims());
+    }
+  }
+
   std::vector<const DenseTensor*> shape_refer;
   for (size_t j = 0; j < outs.size(); ++j) {
     dev_ctx.template Alloc<T>(outs[j]);