[0-size Tensor ] Add 0-size Tensor support for paddle.linalg.vecdot API. (#72694)

DanielSun11 · web-flow · commit 3e7645db98ef · 2025-05-28T16:38:58.000+08:00
* refix the 0-size tensor  for paddle.linalg.vecdot

* fotmat code with google style

* fix:format with google style

* remove "print" in unittest

* add unittest for onednn kernel

* format the unittest case
diff --git a/paddle/phi/kernels/cpu/reduce_sum_grad_kernel.cc b/paddle/phi/kernels/cpu/reduce_sum_grad_kernel.cc
@@ -31,6 +31,10 @@ void ReduceSumGradKernel(const Context& dev_ctx,
                          bool reduce_all,
                          DenseTensor* x_grad) {
   reduce_all = recompute_reduce_all(x, dims, reduce_all);
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   ReduceGradKernel<Context, T, funcs::SumGradFunctor, true>(dev_ctx,
                                                             x,
                                                             paddle::none,
diff --git a/paddle/phi/kernels/cpu/reduce_sum_kernel.cc b/paddle/phi/kernels/cpu/reduce_sum_kernel.cc
@@ -12,10 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <set>
-
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
 
+#include <set>
+
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cpu/reduce.h"
@@ -36,57 +36,12 @@ void SumRawKernel(const Context& dev_ctx,
     out_dtype = out->dtype();
   }
   if (x.numel() == 0) {
-    auto x_dims = x.dims();
-    std::vector<int> out_dims;
-    if (reduce_all) {
-      if (keep_dim) {
-        out_dims.resize(x_dims.size(), 1);
-      } else {
-        out_dims = std::vector<int>();
-      }
-    } else {
-      std::set<int> reduce_dims;
-      auto dims_vec = dims.GetData();
-      for (auto dim : dims_vec) {
-        PADDLE_ENFORCE_GE(dim,
-                          -x_dims.size(),
-                          common::errors::InvalidArgument(
-                              "The dimension index is out of range, "
-                              "expected index >= %d, but received %d.",
-                              -x_dims.size(),
-                              dim));
-        PADDLE_ENFORCE_LT(dim,
-                          x_dims.size(),
-                          common::errors::InvalidArgument(
-                              "The dimension index is out of range, "
-                              "expected index < %d, but received %d.",
-                              x_dims.size(),
-                              dim));
-        if (dim < 0) {
-          dim += x_dims.size();
-        }
-        reduce_dims.insert(dim);
-      }
-      if (keep_dim) {
-        out_dims.resize(x_dims.size());
-        for (int i = 0; i < x_dims.size(); ++i) {
-          if (reduce_dims.count(i)) {
-            out_dims[i] = 1;
-          } else {
-            out_dims[i] = x_dims[i];
-          }
-        }
-      } else {
-        for (int i = 0; i < x_dims.size(); ++i) {
-          if (!reduce_dims.count(i)) {
-            out_dims.push_back(x_dims[i]);
-          }
-        }
-      }
-    }
-    out->Resize(phi::make_ddim(out_dims));
     dev_ctx.template Alloc<T>(out);
-    FullKernel<T, Context>(dev_ctx, out_dims, 0, out_dtype, out);
+    FullKernel<T, Context>(dev_ctx,
+                           phi::IntArray(common::vectorize(out->dims())),
+                           0,
+                           out_dtype,
+                           out);
     return;
   }
   phi::Reduce<CPUContext, T, phi::funcs::SumFunctor>(
diff --git a/paddle/phi/kernels/kps/reduce_kernel.cu b/paddle/phi/kernels/kps/reduce_kernel.cu
@@ -14,6 +14,7 @@
 
 #include <limits>
 #include <set>
+
 #include "paddle/phi/common/complex.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/phi/core/kernel_registry.h"
@@ -210,57 +211,12 @@ void SumRawKernel(const Context& dev_ctx,
     out_dtype = out->dtype();
   }
   if (x.numel() == 0) {
-    auto x_dims = x.dims();
-    std::vector<int> out_dims;
-    if (reduce_all) {
-      if (keep_dim) {
-        out_dims.resize(x_dims.size(), 1);
-      } else {
-        out_dims = std::vector<int>();
-      }
-    } else {
-      std::set<int> reduce_dims;
-      auto dims_vec = dims.GetData();
-      for (auto dim : dims_vec) {
-        PADDLE_ENFORCE_GE(dim,
-                          -x_dims.size(),
-                          common::errors::InvalidArgument(
-                              "The dimension index is out of range, "
-                              "expected index >= %d, but received %d.",
-                              -x_dims.size(),
-                              dim));
-        PADDLE_ENFORCE_LT(dim,
-                          x_dims.size(),
-                          common::errors::InvalidArgument(
-                              "The dimension index is out of range, "
-                              "expected index < %d, but received %d.",
-                              x_dims.size(),
-                              dim));
-        if (dim < 0) {
-          dim += x_dims.size();
-        }
-        reduce_dims.insert(dim);
-      }
-      if (keep_dim) {
-        out_dims.resize(x_dims.size());
-        for (int i = 0; i < x_dims.size(); ++i) {
-          if (reduce_dims.count(i)) {
-            out_dims[i] = 1;
-          } else {
-            out_dims[i] = x_dims[i];
-          }
-        }
-      } else {
-        for (int i = 0; i < x_dims.size(); ++i) {
-          if (!reduce_dims.count(i)) {
-            out_dims.push_back(x_dims[i]);
-          }
-        }
-      }
-    }
-    out->Resize(phi::make_ddim(out_dims));
     dev_ctx.template Alloc<T>(out);
-    FullKernel<T, Context>(dev_ctx, out_dims, 0, out_dtype, out);
+    FullKernel<T, Context>(dev_ctx,
+                           phi::IntArray(common::vectorize(out->dims())),
+                           0,
+                           out_dtype,
+                           out);
     return;
   }
 
diff --git a/paddle/phi/kernels/onednn/reduce_sum_grad_kernel.cc b/paddle/phi/kernels/onednn/reduce_sum_grad_kernel.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/phi/kernels/reduce_sum_grad_kernel.h"
+
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/onednn/reduce_kernel_impl.h"
 #include "paddle/phi/kernels/reduce_kernel_impl.h"
@@ -26,6 +27,10 @@ void SumGradKernel(const Context& dev_ctx,
                    bool keep_dim,
                    bool reduce_all,
                    DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   reduce_all = recompute_reduce_all(x, dims, reduce_all);
   ReduceGradKernel<T, Context>(dev_ctx,
                                x,
diff --git a/paddle/phi/kernels/onednn/reduce_sum_kernel.cc b/paddle/phi/kernels/onednn/reduce_sum_kernel.cc
@@ -13,7 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
+
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/onednn/reduce_kernel_impl.h"
 
 namespace phi {
@@ -26,6 +28,15 @@ void SumRawKernel(const Context& dev_ctx,
                   DataType out_dtype UNUSED,
                   DenseTensor* out) {
   reduce_all = recompute_reduce_all(x, dims, reduce_all);
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    FullKernel<T, Context>(dev_ctx,
+                           phi::IntArray(common::vectorize(out->dims())),
+                           0,
+                           out_dtype,
+                           out);
+    return;
+  }
   ReduceKernel<T, Context>(dev_ctx,
                            x,
                            dims,
diff --git a/paddle/phi/kernels/xpu/reduce_sum_grad_kernel.cc b/paddle/phi/kernels/xpu/reduce_sum_grad_kernel.cc
@@ -31,6 +31,10 @@ void ReduceSumGradKernel(const Context& dev_ctx,
                          bool reduce_all,
                          DenseTensor* x_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   reduce_all = recompute_reduce_all(x, dims_arr, reduce_all);
   auto dims = dims_arr.GetData();
   dev_ctx.Alloc(x_grad, x.dtype());
diff --git a/paddle/phi/kernels/xpu/reduce_sum_kernel.cc b/paddle/phi/kernels/xpu/reduce_sum_kernel.cc
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/xpu/xpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/xpu/reduce.h"
 
 namespace phi {
@@ -31,6 +32,15 @@ void SumRawKernel(const Context& dev_ctx,
   if (out_dtype == DataType::UNDEFINED && out->dtype() != x.dtype()) {
     out_dtype = out->dtype();
   }
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    FullKernel<T, Context>(dev_ctx,
+                           phi::IntArray(common::vectorize(out->dims())),
+                           0,
+                           out_dtype,
+                           out);
+    return;
+  }
   XPUReduce<Context, T, phi::SumFunctor>(
       dev_ctx, x, dims.GetData(), keep_dim, reduce_all, out_dtype, out);
 }
diff --git a/test/legacy_test/test_reduce_op.py b/test/legacy_test/test_reduce_op.py
@@ -207,6 +207,52 @@ def test_check_grad(self):
             )
 
 
+class TestSumOp3D0size(TestSumOp3Dim):
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_pir_onednn=True)
+
+    def calc_gradient(self):
+        x = self.inputs["X"]
+        grad = np.ones(x.shape, dtype=x.dtype)
+        return (grad,)
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['X'],
+            'Out',
+            user_defined_grads=self.calc_gradient(),
+            check_prim=True,
+            check_prim_pir=True,
+            check_pir=True,
+            check_pir_onednn=True,
+        )
+
+
+class TestSumOp3D0size1(TestSumOp3D0size):
+    def init_input(self):
+        self.x = np.random.uniform(0, 0.1, (5, 0, 10)).astype(self.dtype)
+
+    def init_attrs(self):
+        self.attrs = {'dim': (0, 1, 2)}
+
+
+class TestSumOp3D0size2(TestSumOp3D0size):
+    def init_input(self):
+        self.x = np.random.uniform(0, 0.1, (0, 6, 10)).astype(self.dtype)
+
+    def init_attrs(self):
+        self.attrs = {'dim': (0, 1, 2)}
+
+
+class TestSumOp3D0size3(TestSumOp3D0size):
+    def init_input(self):
+        self.x = np.random.uniform(0, 0.1, (4, 6, 0)).astype(self.dtype)
+
+    def init_attrs(self):
+        self.attrs = {'dim': (0, 1, 2)}
+
+
 create_test_fp16_class(TestSumOp)
 create_test_fp16_class(TestSumOp_ZeroDim)
 create_test_fp16_class(TestSumOp5D)
