[0-size Tensor No.269、282、339] Add 0-size Tensor support for subtract API. (#72914)

DanielSun11 · web-flow · commit 93c7be143727 · 2025-05-28T11:07:14.000+08:00
* fix 0 size Tensor for expand kernel in onednn

* format

* expand_grad support 0 size Tensor

* fix bug and add  unittest

* modify the unittest for onednn kernel

* modify expand unittest to force enable onednn op

* fix subtract

* fix  unittest
diff --git a/paddle/phi/kernels/cpu/elementwise_subtract_grad_kernel.cc b/paddle/phi/kernels/cpu/elementwise_subtract_grad_kernel.cc
@@ -17,9 +17,9 @@
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cpu/elementwise_grad.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/elementwise_functor.h"
 #include "paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h"
-
 namespace phi {
 
 template <typename T, typename Context>
@@ -32,6 +32,23 @@ void SubtractGradKernel(const Context& dev_ctx,
                         DenseTensor* dy) {
   // skip out
   auto* out = &dout;
+  if (dout.numel() == 0) {
+    if (dx) {
+      dev_ctx.template Alloc<T>(dx);
+      if (dx->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      dev_ctx.template Alloc<T>(dy);
+      if (dy->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
   ElementwiseSubGrad<T>(dev_ctx, x, y, *out, dout, dx, dy, axis);
 }
 
diff --git a/paddle/phi/kernels/cpu/elementwise_subtract_kernel.cc b/paddle/phi/kernels/cpu/elementwise_subtract_kernel.cc
@@ -27,8 +27,7 @@ void SubtractKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     DenseTensor* out) {
-  if (x.numel() == 0 || y.numel() == 0) {
-    out->Resize(out->dims());
+  if (out && out->numel() == 0) {
     dev_ctx.template Alloc<T>(out);
     return;
   }
diff --git a/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu b/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu
@@ -40,6 +40,23 @@ void SubtractGradKernel(const Context& dev_ctx,
                         DenseTensor* dy) {
   // skip out
   auto* out = &dout;
+  if (dout.numel() == 0) {
+    if (dx) {
+      dev_ctx.template Alloc<T>(dx);
+      if (dx->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      dev_ctx.template Alloc<T>(dy);
+      if (dy->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
   if (dx != nullptr && dy != nullptr && (dx->dims() == dy->dims())) {
     elementwise_sub_grad<T>(dev_ctx, x, y, *out, dout, dx, dy);
   } else {
diff --git a/paddle/phi/kernels/kps/elementwise_kernel.cu b/paddle/phi/kernels/kps/elementwise_kernel.cu
@@ -32,8 +32,7 @@ void SubtractKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     DenseTensor* out) {
-  if (x.numel() == 0 || y.numel() == 0) {
-    out->Resize(out->dims());
+  if (out->numel() == 0) {
     dev_ctx.template Alloc<T>(out);
     return;
   }
diff --git a/paddle/phi/kernels/onednn/elementwise_grad_kernel.cc b/paddle/phi/kernels/onednn/elementwise_grad_kernel.cc
@@ -19,6 +19,7 @@
 
 #include "paddle/phi/backends/onednn/onednn_reuse.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 namespace funcs {
@@ -121,6 +122,23 @@ void ElementwiseGradKernel(const OneDNNContext& dev_ctx,
   const auto& onednn_engine = dev_ctx.GetEngine();
   // oneDNN's binary is optimized for broadcasting y into x, so in other case
   // we have to swap tensors to achieve optimal performance
+  if (dout.numel() == 0) {
+    if (dx) {
+      dev_ctx.template Alloc<T>(dx);
+      if (dx->numel() != 0) {
+        phi::Full<T, OneDNNContext>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      dev_ctx.template Alloc<T>(dy);
+      if (dy->numel() != 0) {
+        phi::Full<T, OneDNNContext>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
   bool swap_x_y = false;
   auto* non_const_x = &x;
   auto* non_const_y = &y;
diff --git a/paddle/phi/kernels/onednn/elementwise_kernel.cc b/paddle/phi/kernels/onednn/elementwise_kernel.cc
@@ -54,6 +54,10 @@ void ElementwiseKernel(const OneDNNContext& dev_ctx,
                        const DenseTensor& y,
                        int axis,
                        DenseTensor* out) {
+  if (out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   const auto& onednn_engine = dev_ctx.GetEngine();
 
   auto* non_const_x = &x;
diff --git a/paddle/phi/kernels/xpu/elementwise_subtract_grad_kernel.cc b/paddle/phi/kernels/xpu/elementwise_subtract_grad_kernel.cc
@@ -16,6 +16,7 @@ limitations under the License. */
 #include "paddle/phi/backends/xpu/xpu_context.h"
 #include "paddle/phi/backends/xpu/xpu_header.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/xpu/elementwise.h"
 
 namespace phi {
@@ -29,6 +30,24 @@ void SubtractGradKernel(const Context& dev_ctx,
                         DenseTensor* dy) {
   using XPUType = typename XPUTypeTrait<T>::Type;
 
+  if (dout.numel() == 0) {
+    if (dx) {
+      dev_ctx.template Alloc<T>(dx);
+      if (dx->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      dev_ctx.template Alloc<T>(dy);
+      if (dy->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
+
   auto f = [](xpu::Context* ctx,
               const XPUType* x,
               const XPUType* y,
diff --git a/paddle/phi/kernels/xpu/elementwise_subtract_kernel.cc b/paddle/phi/kernels/xpu/elementwise_subtract_kernel.cc
@@ -24,6 +24,10 @@ void SubtractKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     DenseTensor* out) {
+  if (out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto f = [](xpu::Context* ctx,
               const XPUType* x,
diff --git a/test/legacy_test/test_elementwise_sub_op.py b/test/legacy_test/test_elementwise_sub_op.py
@@ -32,13 +32,16 @@ def setUp(self):
         self.public_python_api = paddle.subtract
         self.prim_op_type = "prim"
         self.init_dtype()
+        self.init_inputs()
+        self.outputs = {'Out': self.inputs['X'] - self.inputs['Y']}
+        self.if_check_prim()
+        self.if_enable_cinn()
+
+    def init_inputs(self):
         self.inputs = {
             'X': np.random.uniform(0.1, 1, [2, 3, 4, 5]).astype(self.dtype),
             'Y': np.random.uniform(0.1, 1, [2, 3, 4, 5]).astype(self.dtype),
         }
-        self.outputs = {'Out': self.inputs['X'] - self.inputs['Y']}
-        self.if_check_prim()
-        self.if_enable_cinn()
 
     def init_dtype(self):
         self.dtype = np.float64
@@ -1151,6 +1154,38 @@ def test_dygraph_sub(self):
         paddle.enable_static()
 
 
+class TestElementwiseOpZeroSize(TestElementwiseOp):
+    def init_inputs(self):
+        self.attrs = {'enable_check_eager_comp': False}
+        self.inputs = {
+            'X': np.random.uniform(0.1, 1, [2, 0, 4, 5]).astype(self.dtype),
+            'Y': np.random.uniform(0.1, 1, [2, 0, 4, 5]).astype(self.dtype),
+        }
+
+    def if_check_prim(self):
+        self.check_prim = False
+        self.check_prim_pir = False
+
+    def test_check_grad_normal(self):
+        pass
+
+
+class TestElementwiseOpZeroSize2(TestElementwiseOpZeroSize):
+    def init_inputs(self):
+        self.inputs = {
+            'X': np.random.uniform(0.1, 1, [2, 1, 4, 5]).astype(self.dtype),
+            'Y': np.random.uniform(0.1, 1, [2, 0, 4, 5]).astype(self.dtype),
+        }
+
+
+class TestElementwiseOpZeroSize3(TestElementwiseOpZeroSize):
+    def init_inputs(self):
+        self.inputs = {
+            'X': np.random.uniform(0.1, 1, [2, 1, 0, 5]).astype(self.dtype),
+            'Y': np.random.uniform(0.1, 1, [2, 1, 1, 5]).astype(self.dtype),
+        }
+
+
 class TestTensorSubAPIWarnings(unittest.TestCase):
     def test_warnings(self):
         with paddle.pir_utils.OldIrGuard():
diff --git a/test/legacy_test/test_subtract_op.py b/test/legacy_test/test_subtract_op.py
@@ -44,8 +44,12 @@ def test_static_api(self):
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
         ):
-            data_x = paddle.static.data("x", shape=[10, 15], dtype="float32")
-            data_y = paddle.static.data("y", shape=[10, 15], dtype="float32")
+            data_x = paddle.static.data(
+                "x", shape=self.input_x.shape, dtype="float32"
+            )
+            data_y = paddle.static.data(
+                "y", shape=self.input_y.shape, dtype="float32"
+            )
             result_max = paddle.subtract(data_x, data_y)
             exe = paddle.static.Executor(self.place)
             (res,) = exe.run(
@@ -57,8 +61,12 @@ def test_static_api(self):
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
         ):
-            data_x = paddle.static.data("x", shape=[10, 15], dtype="float32")
-            data_z = paddle.static.data("z", shape=[15], dtype="float32")
+            data_x = paddle.static.data(
+                "x", shape=self.input_x.shape, dtype="float32"
+            )
+            data_z = paddle.static.data(
+                "z", shape=self.input_z.shape, dtype="float32"
+            )
             result_max = paddle.subtract(data_x, data_z)
             exe = paddle.static.Executor(self.place)
             (res,) = exe.run(
@@ -70,8 +78,12 @@ def test_static_api(self):
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
         ):
-            data_a = paddle.static.data("a", shape=[3], dtype="int64")
-            data_c = paddle.static.data("c", shape=[3], dtype="int64")
+            data_a = paddle.static.data(
+                "a", shape=self.input_a.shape, dtype="int64"
+            )
+            data_c = paddle.static.data(
+                "c", shape=self.input_b.shape, dtype="int64"
+            )
             result_max = paddle.subtract(data_a, data_c)
             exe = paddle.static.Executor(self.place)
             (res,) = exe.run(
@@ -83,8 +95,12 @@ def test_static_api(self):
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
         ):
-            data_b = paddle.static.data("b", shape=[3], dtype="int64")
-            data_c = paddle.static.data("c", shape=[3], dtype="int64")
+            data_b = paddle.static.data(
+                "b", shape=self.input_b.shape, dtype="int64"
+            )
+            data_c = paddle.static.data(
+                "c", shape=self.input_c.shape, dtype="int64"
+            )
             result_max = paddle.subtract(data_b, data_c)
             exe = paddle.static.Executor(self.place)
             (res,) = exe.run(
@@ -121,5 +137,26 @@ def test_dynamic_api(self):
         np.testing.assert_allclose(res, self.np_expected4, rtol=1e-05)
 
 
+class ApiSubtractTestZeroSize(ApiSubtractTest):
+    def setUp(self):
+        if core.is_compiled_with_cuda():
+            self.place = core.CUDAPlace(0)
+        else:
+            self.place = core.CPUPlace()
+
+        self.input_x = np.random.rand(0, 15).astype("float32")
+        self.input_y = np.random.rand(1, 15).astype("float32")
+        self.input_z = np.random.rand(15).astype("float32")
+        self.input_a = np.random.rand(3).astype('int64')
+        self.input_b = np.random.rand(3).astype('int64')
+        self.input_c = np.random.rand(3).astype('int64')
+
+        self.np_expected1 = np.subtract(self.input_x, self.input_y)
+        self.np_expected2 = np.subtract(self.input_x, self.input_z)
+        self.np_expected3 = np.subtract(self.input_a, self.input_c)
+        self.np_expected4 = np.subtract(self.input_b, self.input_c)
+
+
 if __name__ == "__main__":
+    paddle.enable_static()
     unittest.main()
diff --git a/test/mkldnn/test_elementwise_sub_onednn_op.py b/test/mkldnn/test_elementwise_sub_onednn_op.py
@@ -343,6 +343,32 @@ def test_check_grad_ignore_y(self):
 '''
 
 
+class TestOneDNNElementwiseSubOpZeroSize(TestOneDNNElementwiseSubOp):
+    def init_input_output(self):
+        self.x = np.random.uniform(0.1, 1, [0, 17]).astype(self.dtype)
+        self.y = np.random.uniform(0.1, 1, [1, 17]).astype(self.dtype)
+        self.out = np.subtract(self.x, self.y)
+
+    def test_check_grad_ignore_x(self):
+        self.check_grad_with_place(
+            core.CPUPlace(),
+            ["Y"],
+            "Out",
+            check_pir_onednn=True,
+        )
+
+    def test_check_grad_ignore_y(self):
+        self.check_grad_with_place(
+            core.CPUPlace(),
+            ["X"],
+            "Out",
+            check_pir_onednn=True,
+        )
+
+    def if_check_prim(self):
+        self.check_prim = True
+
+
 if __name__ == '__main__':
     enable_static()
     unittest.main()
