[0-size Tensor Job2 No.2、3、10] Add 0-size Tensor support for equal_all [fluid_ops] (#73550)

* Fix * Fix * Fix * Fix

Author: co63oc

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc

@@ -282,7 +282,9 @@ bool BincountOpInferSymbolicShape(
  "The 'shape' of Input(Weights) must be 1-D tensor. "
  "But the dimension of Input(Weights) is [%d]",
  weights_dims.size()));
- infer_context->AddEqualCstr(weights_dims[0], x_dims[0]);
+ if (x_dims[0] != 0) {
+ infer_context->AddEqualCstr(weights_dims[0], x_dims[0]);
+ }
  }
 
  symbol::DimExpr out_unknown = infer_context->GetNextSymName();

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc

@@ -82,22 +82,32 @@ bool AddNOpInferSymbolicShape(pir::Operation *op,
  "should be larger than 0. But received X's dimensions %d.",
  inputs_shape.size()));
  symbol::TensorShapeOrDataDimExprs candidate_shape = inputs_shape.front();
+ std::vector<symbol::DimExpr> candidate_shape_vec = candidate_shape.shape();
  for (size_t i = 1; i < inputs_shape.size(); ++i) {
  // 0D tensor
  if (inputs_shape[i].shape().size() == 0) {
  continue;
  }
  if (candidate_shape.shape().size() == 0) {
  candidate_shape = inputs_shape[i];
+ candidate_shape_vec = candidate_shape.shape();
  continue;
  }
- for (size_t j = 0; j < candidate_shape.shape().size(); ++j) {
- infer_context->AddEqualCstr(candidate_shape.shape()[j],
- inputs_shape[i].shape()[j]);
+ for (size_t j = 0; j < candidate_shape_vec.size(); ++j) {
+ if (candidate_shape_vec[j] != 0) {
+ if (inputs_shape[i].shape()[j] != 0) {
+ infer_context->AddEqualCstr(candidate_shape_vec[j],
+ inputs_shape[i].shape()[j]);
+ } else {
+ candidate_shape_vec[j] = symbol::DimExpr{0};
+ }
+ }
  }
  }
  infer_context->SetShapeOrDataForValue(
- op->result(0), symbol::ShapeOrDataDimExprs{candidate_shape});
+ op->result(0),
+ symbol::ShapeOrDataDimExprs{
+ symbol::TensorShapeOrDataDimExprs(candidate_shape_vec)});
 
  return true;
 }

paddle/phi/infermeta/binary.cc

@@ -295,16 +295,18 @@ void BincountInferMeta(const MetaTensor& x,
  "But the dimension of Input(Weights) is [%d]",
  weights_dim.size()));
 
- PADDLE_ENFORCE_EQ(
- weights_dim[0],
- input_dim[0],
- common::errors::InvalidArgument(
- "The 'shape' of Input(Weights) must be equal to the 'shape' of "
- "Input(X)."
- "But received: the 'shape' of Input(Weights) is [%s],"
- "the 'shape' of Input(X) is [%s]",
- weights_dim,
- input_dim));
+ if (input_dim[0] != 0) {
+ PADDLE_ENFORCE_EQ(
+ weights_dim[0],
+ input_dim[0],
+ common::errors::InvalidArgument(
+ "The 'shape' of Input(Weights) must be equal to the 'shape' of "
+ "Input(X)."
+ "But received: the 'shape' of Input(Weights) is [%s],"
+ "the 'shape' of Input(X) is [%s]",
+ weights_dim,
+ input_dim));
+ }
  }
  out->set_dims(common::make_ddim({-1}));
  if (weights) {

paddle/phi/infermeta/multiary.cc

@@ -415,19 +415,30 @@ void AddNInferMeta(const std::vector<const MetaTensor*>& x,
  continue;
  }
  is_all_0d_tensor = false;
- if (common::product(in_dim) == 0) {
+ // use the first dimension
+ if (i == 0) {
  in_dim = x_dim;
  } else {
  if (config.is_runtime) {
- PADDLE_ENFORCE_EQ(in_dim,
- x_dim,
- common::errors::InvalidArgument(
- "The input tensor X of AddNOp must"
- " have same shape. But received X[0]'s shape = "
- "[%s], X[%d]'s shape = [%s].",
- in_dim,
- i,
- x_dim));
+ for (int j = 0; j < x_dim.size(); ++j) {
+ if (in_dim[j] != 0) {
+ if (x_dim[j] == 0) {
+ // update the 0 dim
+ in_dim[j] = 0;
+ } else {
+ PADDLE_ENFORCE_EQ(
+ in_dim[j],
+ x_dim[j],
+ common::errors::InvalidArgument(
+ "The input tensor X of AddNOp must"
+ " have same shape. But received X[0]'s shape = "
+ "[%s], X[%d]'s shape = [%s].",
+ in_dim,
+ i,
+ x_dim));
+ }
+ }
+ }
  } else {
  PADDLE_ENFORCE_EQ(
  in_dim.size(),

paddle/phi/kernels/cpu/add_n_kernel.cc

@@ -25,6 +25,9 @@ void AddNKernel(const Context& dev_ctx,
  DenseTensor* out) {
  size_t in_num = x.size();
  dev_ctx.template Alloc<T>(out);
+ if (out && out->numel() == 0) {
+ return;
+ }
 
  bool in_place = false;
  if (!x.empty() && x[0]->initialized() && DenseTensor::classof(x[0])) {

paddle/phi/kernels/cpu/bincount_kernel.cc

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
@@ -33,9 +34,11 @@ void BincountInner(const Context& dev_ctx,
  auto input_numel = input->numel();
 
  if (input_data == nullptr) {
- phi::DDim out_dim{0};
+ phi::DDim out_dim{minlength};
  output->Resize(out_dim);
- dev_ctx.template Alloc<InputT>(output);
+ // Since minlength may >0 , so fill with 0.
+ phi::Full<InputT, Context>(
+ dev_ctx, phi::IntArray(common::vectorize(output->dims())), 0, output);
  return;
  }
 

paddle/phi/kernels/cpu/compare_kernel.cc

@@ -70,6 +70,8 @@ inline void CompareAllKernelImpl(const Context& dev_ctx,
 
  if (x.dims() != y.dims()) {
  out_data[0] = false;
+ } else if (x.numel() == 0) { // shape equal and numel is 0, return true
+ out_data[0] = true;
  } else {
  DenseTensor tmp;
  tmp.Resize(x.dims());

paddle/phi/kernels/gpu/add_n_kernel.cu

@@ -79,6 +79,10 @@ template <typename T, typename Context>
 void AddNKernel(const Context &dev_ctx,
  const std::vector<const TensorBase *> &x,
  DenseTensor *out) {
+ if (out && out->numel() == 0) {
+ dev_ctx.template Alloc<T>(out);
+ return;
+ }
  const size_t in_num = x.size();
  for (int i = 0; i < in_num; ++i) {
  PADDLE_ENFORCE_EQ(

paddle/phi/kernels/gpu/bincount_kernel.cu

@@ -17,9 +17,9 @@
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/backends/gpu/gpu_primitives.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
-
 namespace phi {
 
 using phi::PADDLE_CUDA_NUM_THREADS;
@@ -99,9 +99,10 @@ void BincountCUDAInner(const Context& dev_ctx,
  int64_t input_numel = static_cast<int64_t>(input->numel());
 
  if (input_data == nullptr) {
- phi::DDim out_dim{0};
+ phi::DDim out_dim{minlength};
  output->Resize(out_dim);
- dev_ctx.template Alloc<T>(output);
+ phi::Full<InputT, Context>(
+ dev_ctx, phi::IntArray(common::vectorize(output->dims())), 0, output);
  return;
  }
 

paddle/phi/kernels/kps/compare_kernel.cu

@@ -80,6 +80,12 @@ inline void CompareAllKernelImpl(const Context& ctx,
  thrust::fill(out_dev_ptr, out_dev_ptr + 1, false);
  return;
  }
+ // shape equal and numel is 0, return true
+ if (x.numel() == 0) {
+ thrust::device_ptr<bool> out_dev_ptr(out_data);
+ thrust::fill(out_dev_ptr, out_dev_ptr + 1, true);
+ return;
+ }
 
  DenseTensor tmp;
  tmp.Resize(x.dims());