[0-size Tensor Job2 No.30] Add 0-size Tensor support for swiglu (#73346)

co63oc · web-flow · commit 6ec5b572188d · 2025-06-17T14:03:09.000+08:00
* Fix

* ci

* Fix
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
@@ -2088,11 +2088,26 @@ bool SwigluOpInferSymbolicShape(pir::Operation *op,
   if (op->operand_source(1)) {
     const auto &y_shape_or_data =
         infer_context->GetShapeOrDataForValue(op->operand_source(1));
+    bool x_size_0 = false;
+    bool y_size_0 = false;
     for (size_t i = 0; i < rank; ++i) {
+      if (x_shape_or_data.shape()[i] == 0) {
+        x_size_0 = true;
+      }
+      if (y_shape_or_data.shape()[i] == 0) {
+        y_size_0 = true;
+      }
+      if (x_shape_or_data.shape()[i] == 0 || y_shape_or_data.shape()[i] == 0) {
+        continue;
+      }
       infer_context->AddEqualCstr(x_shape_or_data.shape()[i],
                                   y_shape_or_data.shape()[i]);
     }
     infer_context->SetShapeOrDataForValue(op->result(0), x_shape_or_data);
+    if (!x_size_0 && y_size_0) {
+      // set y shape
+      infer_context->SetShapeOrDataForValue(op->result(0), y_shape_or_data);
+    }
   } else {
     std::vector<symbol::DimExpr> x_shape = x_shape_or_data.shape();
     // TODO(CINN): Add distribute constraint
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -4452,12 +4452,21 @@ void SwiGLUInferMeta(const MetaTensor& x,
                      const MetaTensor& y,
                      MetaTensor* out) {
   if (y) {
-    PADDLE_ENFORCE_EQ(
-        x.dims(),
-        y.dims(),
-        common::errors::InvalidArgument(
-            "The shape of Input(X) should be equal of the shape of Input(Y)."));
+    auto x_numel = common::product(x.dims());
+    auto y_numel = common::product(y.dims());
+    // skip 0-size
+    if (x_numel != 0 && y_numel != 0) {
+      PADDLE_ENFORCE_EQ(
+          x.dims(),
+          y.dims(),
+          common::errors::InvalidArgument("The shape of Input(X) should be "
+                                          "equal of the shape of Input(Y)."));
+    }
     out->share_meta(x);
+    // If y is 0-size, out is 0-size
+    if (x_numel != 0 && y_numel == 0) {
+      out->set_dims(y.dims());
+    }
   } else {
     auto dims = x.dims();
     PADDLE_ENFORCE_EQ(
diff --git a/paddle/phi/kernels/swiglu_grad_kernel.h b/paddle/phi/kernels/swiglu_grad_kernel.h
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/device_context.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 
@@ -36,12 +37,19 @@ void SwiGLUGradKernel(const Context &dev_ctx,
                       const DenseTensor &dz,
                       DenseTensor *dx,
                       DenseTensor *dy) {
-  if (x.numel() == 0) {
-    if (dx) {
-      dev_ctx.template Alloc<T>(dx);
-    }
+  if (dx && dx->numel() == 0) {
+    dev_ctx.template Alloc<T>(dx);
     if (dy) {
-      dev_ctx.template Alloc<T>(dy);
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+    }
+    return;
+  }
+  if (dy && dy->numel() == 0) {
+    dev_ctx.template Alloc<T>(dy);
+    if (dx) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
     }
     return;
   }
diff --git a/paddle/phi/kernels/swiglu_kernel.h b/paddle/phi/kernels/swiglu_kernel.h
@@ -16,7 +16,7 @@
 
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/device_context.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 
 template <typename T, typename Context>
@@ -28,10 +28,12 @@ void SwiGLUKernel(const Context &dev_ctx,
                   const DenseTensor &x,
                   const paddle::optional<DenseTensor> &y,
                   DenseTensor *z) {
-  if (x.numel() == 0) {
+  // If either x or y has a numel 0, the numel of z is 0.
+  if (z->numel() == 0) {
     dev_ctx.template Alloc<T>(z);
     return;
   }
+
   const auto *x_ptr = x.data<T>();
   auto *z_ptr = dev_ctx.template Alloc<T>(z);
   const auto &dims = x.dims();
diff --git a/paddle/phi/kernels/xpu/swiglu_kernel.cc b/paddle/phi/kernels/xpu/swiglu_kernel.cc
@@ -15,7 +15,7 @@
 #include "paddle/phi/kernels/swiglu_kernel.h"
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 template <typename T, typename Context>
 void SwiGluKernel(const Context& dev_ctx,
@@ -26,6 +26,7 @@ void SwiGluKernel(const Context& dev_ctx,
   using XPUTypefp32 = typename XPUTypeTrait<float>::Type;
   const auto* x_data = x.data<T>();
   auto* z_data = dev_ctx.template Alloc<T>(z);
+  if (z->numel() == 0) return;
   const auto& dims = x.dims();
   int64_t axis = dims.size() - 1;
   auto dims_vec = common::vectorize<int64_t>(dims);
diff --git a/paddle/phi/kernels/xpu/swiglu_kernel_grad.cc b/paddle/phi/kernels/xpu/swiglu_kernel_grad.cc
@@ -22,8 +22,8 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/swiglu_grad_kernel.h"
-
 namespace phi {
 
 template <typename T, typename Context>
@@ -33,6 +33,22 @@ void SwiGluGradKernel(const Context& dev_ctx,
                       const DenseTensor& dz,
                       DenseTensor* dx,
                       DenseTensor* dy) {
+  if (dx && dx->numel() == 0) {
+    dev_ctx.template Alloc<T>(dx);
+    if (dy) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+    }
+    return;
+  }
+  if (dy && dy->numel() == 0) {
+    dev_ctx.template Alloc<T>(dy);
+    if (dx) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+    }
+    return;
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
   const auto* x_data = x.data<T>();
   const auto* dz_data = dz.data<T>();
diff --git a/test/legacy_test/test_swiglu.py b/test/legacy_test/test_swiglu.py
@@ -297,5 +297,47 @@ def test_swiglu(self):
         self.assertEqual(out[1].shape, y.shape)
 
 
+class TestSwigluOp_ZeroSize(OpTest):
+    def config(self):
+        self.x_shape = (0, 128)
+        self.y_shape = (1, 128)
+        self.out_shape = (0, 128)
+
+    def setUp(self):
+        self.config()
+        self.op_type = "swiglu"
+        self.python_api = fused_swiglu_impl
+        self.public_python_api = fused_swiglu_impl
+        x = np.random.uniform(-1, 1, self.x_shape).astype("float64")
+        y = np.random.uniform(-1, 1, self.y_shape).astype("float64")
+        out_grad = np.random.uniform(-1, 1, self.out_shape).astype("float64")
+        res = swiglu(x, y, out_grad)
+        self.inputs = {'x': x, 'y': y}
+        self.outputs = {'out': res[0].numpy()}
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['x', 'y'],
+            'out',
+        )
+
+
+class TestSwigluOp_ZeroSize2(TestSwigluOp_ZeroSize):
+    def config(self):
+        self.x_shape = (1, 128)
+        self.y_shape = (0, 128)
+        self.out_shape = (0, 128)
+
+
+class TestSwigluOp_ZeroSize3(TestSwigluOp_ZeroSize):
+    def config(self):
+        self.x_shape = (0, 128)
+        self.y_shape = (0, 128)
+        self.out_shape = (0, 128)
+
+
 if __name__ == "__main__":
     unittest.main()
