bugfix:p_norm test=develop

Author: LCStayingdullCircuit

paddle/phi/kernels/gpu/p_norm_grad_kernel.cu

@@ -50,29 +50,14 @@ struct PNormGradFunctor {
                   DY* dy,
                   const Dim& dim,
                   int size) {
-    auto x_mt = x->template cast<MT>();
-    auto y_mt = y->template cast<MT>();
-    auto dy_mt = dy->template cast<MT>();
-
-    auto norm_pow = y_mt.pow(-this->porder);
-    auto mask_norm_nonzero = (y_mt != static_cast<MT>(0)).template cast<MT>();
-
-    // Set to 0 where porder < 0 and x == 0
-    MT zero = static_cast<MT>(0);
-    auto mask_x_zero = (x_mt == zero).template cast<MT>();
-
-    MT is_porder_negative =
-        this->porder < zero ? static_cast<MT>(1) : static_cast<MT>(0);
-    auto invalid_mask = (mask_x_zero * is_porder_negative);
-    auto safe_pow =
-        x_mt.abs().pow(this->porder) * (static_cast<MT>(1) - invalid_mask);
-
+    auto unstable_term = (*x).abs().pow(this->porder);
+    auto mask = (*x) == x->constant(static_cast<T>(0));
+    auto stable_term =
+        mask.select(x->constant(static_cast<T>(0)), unstable_term);
+    auto self_scaled = (*x).sign() * stable_term;
+    auto norm_term = (*y).pow(-this->porder);
     dx->device(place) =
-        (safe_pow * x_mt.sign() * dy_mt.broadcast(dim) *
-         norm_pow.broadcast(dim) *
-         mask_norm_nonzero.broadcast(dim)  // Mask out positions where norm == 0
-         )
-            .template cast<T>();
+        self_scaled * dy->broadcast(dim) * norm_term.broadcast(dim);
   }
 
   MT porder;

paddle/phi/kernels/gpu/p_norm_kernel.cu

@@ -22,6 +22,9 @@
 #include "paddle/phi/kernels/funcs/reduce_function.h"
 #include "paddle/phi/kernels/gpu/reduce.h"
 
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/phi/kernels/activation_kernel.h"
+
 namespace phi {
 template <typename T>
 struct NonzeroFunctor {
@@ -134,8 +137,24 @@ void PNormKernel(const Context& dev_ctx,
           dev_ctx, *in_x, out_norm, FabsFunctor<T>(), reduce_axis);
     } else if (porder == 2.0) {
       // fast 2-norm
-      phi::funcs::ReduceKernel<T, MT, kps::AddFunctor, SquareFunctor<MT>>(
-          dev_ctx, *in_x, &out_temp, SquareFunctor<MT>(), reduce_axis);
+      using MT = typename phi::dtype::MPTypeTrait<T>::Type;
+      phi::DenseTensor temp_sum_of_squares_hp;
+      temp_sum_of_squares_hp.Resize(out_norm->dims());
+      dev_ctx.template Alloc<MT>(&temp_sum_of_squares_hp);
+      phi::funcs::ReduceKernel<T, MT, kps::AddFunctor, SquareFunctor<T>>(
+          dev_ctx,
+          *in_x,
+          &temp_sum_of_squares_hp,
+          SquareFunctor<T>(),
+          reduce_axis);
+
+      phi::DenseTensor temp_norm_hp;
+      temp_norm_hp.Resize(out_norm->dims());
+      dev_ctx.template Alloc<MT>(&temp_norm_hp);
+      phi::SqrtKernel<MT>(dev_ctx, temp_sum_of_squares_hp, &temp_norm_hp);
+      phi::CastKernel<MT>(dev_ctx, temp_norm_hp, out_norm->dtype(), out_norm);
+      return;
+
     } else if (porder == 3.0) {
       // fast 3-norm
       phi::funcs::ReduceKernel<T, MT, kps::AddFunctor, FabsCubicFunctor<MT>>(