[0-size Tensor Job2 No.9] Add 0-size Tensor support for paddle.dot (#73450)

co63oc · web-flow · commit 09289069a933 · 2025-06-20T22:01:44.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -704,14 +704,23 @@ bool DotOpInferSymbolicShape(pir::Operation *op,
                         y_rank,
                         x_rank));
   for (size_t i = 0; i < x_rank; ++i) {
+    if (x_shape[i] == 0 || y_shape[i] == 0) {
+      continue;
+    }
     infer_context->AddEqualCstr(x_shape[i], y_shape[i]);
   }
   // Dot OP require both inputs should have the same shape
 
-  x_shape.erase(x_shape.end() - 1);
-  auto output_shape =
-      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)};
-  infer_context->SetShapeOrDataForValue(op->result(0), output_shape);
+  auto out_shape = x_shape;
+  // The output dims need to be modified.
+  if (x_rank == 2 && x_shape[0] != 0 && y_shape[0] == 0) {
+    out_shape[0] = symbol::DimExpr{0};
+  }
+  out_shape.erase(out_shape.end() - 1);
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
   return true;
 }
 
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -1613,6 +1613,9 @@ void DotInferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out) {
           x_dims.to_str()));
   bool shape_match = true;
   for (int i = 0; i < x_rank; ++i) {
+    if (x_dims[i] == 0 || y_dims[i] == 0) {
+      continue;
+    }
     if (x_dims[i] != y_dims[i]) {
       shape_match = false;
       break;
@@ -1627,10 +1630,16 @@ void DotInferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out) {
                         "with input tensor Y: %s",
                         x_dims.to_str(),
                         y_dims.to_str()));
-  std::vector<int64_t> x_dims_vec = common::vectorize(x_dims);
-  std::vector<int64_t> x_dims_vec_cut(x_dims_vec.begin(), x_dims_vec.end() - 1);
-  x_dims = common::make_ddim(x_dims_vec_cut);
-  out->set_dims(x_dims);
+
+  auto out_dims = x_dims;
+  // The output dims need to be modified.
+  if (x_rank == 2 && x_dims[0] != 0 && y_dims[0] == 0) {
+    out_dims[0] = 0;
+  }
+  std::vector<int64_t> out_dims_vec = common::vectorize(out_dims);
+  std::vector<int64_t> out_dims_vec_cut(out_dims_vec.begin(),
+                                        out_dims_vec.end() - 1);
+  out->set_dims(common::make_ddim(out_dims_vec_cut));
   out->set_dtype(x.dtype());
   out->set_layout(x.layout());
 }
diff --git a/paddle/phi/kernels/cpu/dot_kernel.cc b/paddle/phi/kernels/cpu/dot_kernel.cc
@@ -18,6 +18,7 @@
 #include "paddle/phi/core/kernel_registry.h"
 
 #include "paddle/phi/common/complex.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 
@@ -26,6 +27,12 @@ void DotKernel(const Context& dev_ctx,
                const DenseTensor& x,
                const DenseTensor& y,
                DenseTensor* out) {
+  if (x.numel() == 0 || y.numel() == 0) {
+    // x[2, 1], y[2, 0], out[2]
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+    return;
+  }
   if (out->numel() <= 0) {
     return;
   }
diff --git a/paddle/phi/kernels/gpu/dot_kernel.cu b/paddle/phi/kernels/gpu/dot_kernel.cu
@@ -21,15 +21,21 @@
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 
 #include "paddle/phi/common/complex.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
-
 namespace phi {
 
 template <typename T, typename Context>
 void DotKernel(const Context& dev_ctx,
                const DenseTensor& x,
                const DenseTensor& y,
                DenseTensor* out) {
+  if (x.numel() == 0 || y.numel() == 0) {
+    // x[2, 1], y[2, 0], out[2]
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+    return;
+  }
   if (out->numel() <= 0) {
     return;
   }
diff --git a/paddle/phi/kernels/impl/dot_grad_kernel_impl.h b/paddle/phi/kernels/impl/dot_grad_kernel_impl.h
@@ -23,7 +23,6 @@ limitations under the License. */
 #include "paddle/phi/kernels/funcs/complex_functors.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
-
 namespace phi {
 
 template <typename DeviceContext, typename T, typename Enable = void>
@@ -1381,9 +1380,19 @@ void DotGradKernel(const Context& dev_ctx,
                    DenseTensor* dy) {
   if (dx) {
     dev_ctx.template Alloc<T>(dx);
+    if (dx->numel() == 0) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(y.dims())), 0, dy);
+      return;
+    }
   }
   if (dy) {
     dev_ctx.template Alloc<T>(dy);
+    if (dy->numel() == 0) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(x.dims())), 0, dx);
+      return;
+    }
   }
   DotGradFunction<Context, T>()(dev_ctx, &x, &y, &dout, dx, dy);
 }
diff --git a/test/legacy_test/test_dot_op.py b/test/legacy_test/test_dot_op.py
@@ -450,6 +450,41 @@ def test_check_grad_ignore_y(self):
                 )
 
 
+class DotOp_ZeroSize(OpTest):
+    def setUp(self):
+        self.op_type = "dot"
+        self.python_api = paddle.dot
+        self.public_python_api = paddle.dot
+        self.init_shape()
+        self.init_dtype()
+        self.init_input_output()
+
+        self.inputs = {
+            'X': OpTest.np_dtype_to_base_dtype(self.x),
+            'Y': OpTest.np_dtype_to_base_dtype(self.y),
+        }
+        self.outputs = {'Out': self.out}
+        self.attrs = {}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+    def test_check_grad_normal(self):
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True)
+
+    def init_input_output(self):
+        self.x = np.random.uniform(0.1, 1, self.shape).astype(self.dtype)
+        self.y = np.random.uniform(1, 3, self.shape).astype(self.dtype)
+        self.out = np.dot(self.x, self.y).astype(self.dtype)
+
+    def init_dtype(self):
+        self.dtype = np.float64
+
+    def init_shape(self):
+        # return shape []
+        self.shape = [0]
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
