[PHI] Fix p_norm kernel for big tensor (#73262)

lshpku · web-flow · commit 0de9e93e5ae9 · 2025-06-12T10:28:48.000+08:00
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3389,7 +3389,7 @@ void PNormInferMeta(const MetaTensor& x,
                         x_rank,
                         x_dim));
 
-  std::vector<int> out_dim_vector;
+  std::vector<int64_t> out_dim_vector;
   if (asvector) {
     if (keepdim) {
       for (int i = 0; i < x_rank; ++i) {
diff --git a/paddle/phi/kernels/gpu/p_norm_kernel.cu b/paddle/phi/kernels/gpu/p_norm_kernel.cu
@@ -77,82 +77,6 @@ struct FabsCubicFunctor {
 };
 #endif
 
-#ifndef PADDLE_WITH_XPU_KP
-
-inline void GetDims(const phi::DDim& dim,
-                    int axis,
-                    int* pre,
-                    int* n,
-                    int* post,
-                    bool asvector) {
-  *pre = 1;
-  *post = 1;
-  *n = static_cast<int>(dim[axis]);
-  if (asvector) {
-    *n = static_cast<int>(product(dim));
-  } else {
-    for (int i = 0; i < axis; ++i) {
-      (*pre) *= static_cast<int>(dim[i]);
-    }
-    for (int i = axis + 1; i < dim.size(); ++i) {
-      (*post) *= static_cast<int>(dim[i]);
-    }
-  }
-}
-
-template <typename T, typename Context>
-void ReducePNormEigen(const Context& dev_ctx,
-                      const DenseTensor& x,
-                      float porder,
-                      int axis,
-                      float epsilon,
-                      bool keepdim,
-                      bool asvector,
-                      DenseTensor* out) {
-  auto xdim = x.dims();
-  if (axis < 0) axis = xdim.size() + axis;
-  int pre = 0, n = 0, post = 0;
-  GetDims(xdim, axis, &pre, &n, &post, asvector);
-
-  for (int i = 0; i < xdim.size(); i++) {
-    PADDLE_ENFORCE_LT(0,
-                      xdim[i],
-                      errors::InvalidArgument(
-                          "The dims of Input(X) should be greater than 0."));
-  }
-
-  auto* place = dev_ctx.eigen_device();
-
-  Eigen::DSizes<int, 3> shape(pre, n, post);
-  Eigen::DSizes<int, 2> norm_shape(pre, post);
-
-  auto x_e = phi::EigenVector<T>::Flatten(x);
-  auto norm_e = phi::EigenVector<T>::Flatten(*out);
-
-  auto xr = x_e.reshape(shape);
-  auto norm = norm_e.reshape(norm_shape);
-
-  // p=0 means number of non-zero elements of (xr)
-  // p=inf means the maximum of |xr|
-  // p=-inf means the minimum of |xr|
-  // otherwise, Lp-norm = pow(sum(pow(|xr|, p)), 1/p)
-  Eigen::DSizes<int, 1> rdim(1);
-  if (porder == 0) {
-    norm.device(*place) =
-        (xr != xr.constant(static_cast<T>(0))).template cast<T>().sum(rdim);
-  } else if (porder == INFINITY) {
-    norm.device(*place) = xr.abs().maximum(rdim);
-  } else if (porder == -INFINITY) {
-    norm.device(*place) = xr.abs().minimum(rdim);
-  } else {
-    norm.device(*place) = xr.abs()
-                              .pow(static_cast<T>(porder))
-                              .sum(rdim)
-                              .pow(static_cast<T>(1.0f / porder));
-  }
-}
-#endif
-
 template <typename T, typename Context>
 void PNormKernel(const Context& dev_ctx,
                  const DenseTensor& x,
@@ -179,70 +103,54 @@ void PNormKernel(const Context& dev_ctx,
     return;
   }
 
-  if (x.numel() > std::numeric_limits<int32_t>::max()) {
-#ifndef PADDLE_WITH_XPU_KP
-    ReducePNormEigen<T, Context>(
-        dev_ctx, *in_x, porder, axis, epsilon, keepdim, asvector, out_norm);
-#else
-    PADDLE_THROW(common::errors::Fatal(
-        "If Input.numel() > INT32_MAX, reduce_sum kernel uses EigenTensor "
-        "sum for reduce_sum function. Such case is only supported on GPU "
-        "now."));
-#endif
+  using MT = typename dtype::MPTypeTrait<T>::Type;
+  if (porder == 0) {
+    phi::funcs::ReduceKernel<T, T, kps::AddFunctor, NonzeroFunctor<T>>(
+        dev_ctx, *in_x, out_norm, NonzeroFunctor<T>(), reduce_axis);
+  } else if (porder == INFINITY) {
+    phi::funcs::ReduceKernel<T, T, kps::MaxFunctor, AbsFunctor<T>>(
+        dev_ctx, *in_x, out_norm, AbsFunctor<T>(), reduce_axis);
+  } else if (porder == -INFINITY) {
+    phi::funcs::ReduceKernel<T, T, kps::MinFunctor, AbsFunctor<T>>(
+        dev_ctx, *in_x, out_norm, AbsFunctor<T>(), reduce_axis);
   } else {
-    using MT = typename dtype::MPTypeTrait<T>::Type;
-    if (porder == 0) {
-      phi::funcs::ReduceKernel<T, T, kps::AddFunctor, NonzeroFunctor<T>>(
-          dev_ctx, *in_x, out_norm, NonzeroFunctor<T>(), reduce_axis);
-    } else if (porder == INFINITY) {
-      phi::funcs::ReduceKernel<T, T, kps::MaxFunctor, AbsFunctor<T>>(
-          dev_ctx, *in_x, out_norm, AbsFunctor<T>(), reduce_axis);
-    } else if (porder == -INFINITY) {
-      phi::funcs::ReduceKernel<T, T, kps::MinFunctor, AbsFunctor<T>>(
-          dev_ctx, *in_x, out_norm, AbsFunctor<T>(), reduce_axis);
-    } else {
 #ifdef _WIN32
+    phi::funcs::ReduceKernel<T, T, kps::AddFunctor, UnsignedPowFunctor<T>>(
+        dev_ctx, *in_x, out_norm, UnsignedPowFunctor<T>(porder), reduce_axis);
+
+    const DenseTensor* tmp_norm = out_norm;
+    std::vector<const DenseTensor*> ins = {tmp_norm};
+    std::vector<DenseTensor*> outs = {out_norm};
+    phi::funcs::ElementwiseKernel<T>(
+        dev_ctx, ins, &outs, UnsignedPowFunctor<T>(1. / porder));
+#else
+    if (porder == 1.0) {
+      // fast 1-norm
+      phi::funcs::ReduceKernel<T, T, kps::AddFunctor, FabsFunctor<T>>(
+          dev_ctx, *in_x, out_norm, FabsFunctor<T>(), reduce_axis);
+    } else if (porder == 2.0) {
+      // fast 2-norm
+      phi::funcs::ReduceKernel<T, T, kps::AddFunctor, SquareFunctor<T>>(
+          dev_ctx, *in_x, out_norm, SquareFunctor<T>(), reduce_axis);
+    } else if (porder == 3.0) {
+      // fast 3-norm
+      phi::funcs::ReduceKernel<T, T, kps::AddFunctor, FabsCubicFunctor<T>>(
+          dev_ctx, *in_x, out_norm, FabsCubicFunctor<T>(), reduce_axis);
+    } else {
+      // vanilla norm
       phi::funcs::ReduceKernel<T, T, kps::AddFunctor, UnsignedPowFunctor<T>>(
           dev_ctx, *in_x, out_norm, UnsignedPowFunctor<T>(porder), reduce_axis);
+    }
 
+    if (porder != 1.0) {
+      // save computation when porder is 1.0
       const DenseTensor* tmp_norm = out_norm;
       std::vector<const DenseTensor*> ins = {tmp_norm};
       std::vector<DenseTensor*> outs = {out_norm};
       phi::funcs::ElementwiseKernel<T>(
           dev_ctx, ins, &outs, UnsignedPowFunctor<T>(1. / porder));
-#else
-      if (porder == 1.0) {
-        // fast 1-norm
-        phi::funcs::ReduceKernel<T, T, kps::AddFunctor, FabsFunctor<T>>(
-            dev_ctx, *in_x, out_norm, FabsFunctor<T>(), reduce_axis);
-      } else if (porder == 2.0) {
-        // fast 2-norm
-        phi::funcs::ReduceKernel<T, T, kps::AddFunctor, SquareFunctor<T>>(
-            dev_ctx, *in_x, out_norm, SquareFunctor<T>(), reduce_axis);
-      } else if (porder == 3.0) {
-        // fast 3-norm
-        phi::funcs::ReduceKernel<T, T, kps::AddFunctor, FabsCubicFunctor<T>>(
-            dev_ctx, *in_x, out_norm, FabsCubicFunctor<T>(), reduce_axis);
-      } else {
-        // vanilla norm
-        phi::funcs::ReduceKernel<T, T, kps::AddFunctor, UnsignedPowFunctor<T>>(
-            dev_ctx,
-            *in_x,
-            out_norm,
-            UnsignedPowFunctor<T>(porder),
-            reduce_axis);
-      }
-
-      if (porder != 1.0) {
-        // save computation when porder is 1.0
-        const DenseTensor* tmp_norm = out_norm;
-        std::vector<const DenseTensor*> ins = {tmp_norm};
-        std::vector<DenseTensor*> outs = {out_norm};
-        phi::funcs::ElementwiseKernel<T>(
-            dev_ctx, ins, &outs, UnsignedPowFunctor<T>(1. / porder));
-      }
-#endif
     }
+#endif
   }
 }
 }  // namespace phi
