[Precision Depth Alignment]support p_norm_grad (#78477)

Author: AlAuAu

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -2806,8 +2806,8 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
       bool asvector = PADDLE_GET_CONST(bool, desc.GetAttr("asvector"));
       int axis = PADDLE_GET_CONST(int, desc.GetAttr("axis"));
-      float porder = PADDLE_GET_CONST(float, desc.GetAttr("porder"));
-      if (asvector || porder != 2.0f || axis != -1) {
+      double porder = PADDLE_GET_CONST(double, desc.GetAttr("porder"));
+      if (asvector || porder != 2.0 || axis != -1) {
         VLOG(3) << op_type
                 << " op only support asvector=False, porder=2, axis = -1.";
         return false;

paddle/fluid/pir/serialize_deserialize/patch/4.yaml

@@ -121,6 +121,12 @@ op_patches:
         object : using_weighted_combine
         type : pir::BoolAttribute
         data : "false"
+  - op_name : pd_op.p_norm
+    actions :
+      - action : modify_attr
+        object : porder
+        type : pir::DoubleAttribute
+        data : 2
   - op_name : pd_op.layer_norm
     actions:
       - action : modify_attr

paddle/phi/infermeta/unary.cc

@@ -3781,7 +3781,7 @@ void PixelUnshuffleInferMeta(const MetaTensor& x,
 }
 
 void PNormInferMeta(const MetaTensor& x,
-                    float porder,
+                    double porder,
                     int axis,
                     float epsilon,
                     bool keepdim,

paddle/phi/infermeta/unary.h

@@ -591,7 +591,7 @@ PADDLE_API void PixelUnshuffleInferMeta(const MetaTensor& x,
                                         MetaTensor* out);
 
 PADDLE_API void PNormInferMeta(const MetaTensor& x,
-                               float porder,
+                               double porder,
                                int axis,
                                float epsilon,
                                bool keepdim,

paddle/phi/kernels/cpu/p_norm_grad_kernel.cc

@@ -44,7 +44,7 @@ void PNormGradKernel(const Context& dev_ctx,
                      const DenseTensor& x,
                      const DenseTensor& out,
                      const DenseTensor& out_grad,
-                     float porder,
+                     double porder,
                      int axis,
                      float epsilon,
                      bool keepdim UNUSED,
@@ -88,8 +88,8 @@ void PNormGradKernel(const Context& dev_ctx,
                         xr.sign() * norm_dy.broadcast(bcast);
   } else {
     dx.device(*place) =
-        (xr.abs()).pow(porder - 1.0f) /
-        ((norm.broadcast(bcast)).pow(porder - 1.0f) + xr.constant(eps));
+        (xr.abs()).pow(porder - 1.0) /
+        ((norm.broadcast(bcast)).pow(porder - 1.0) + xr.constant(eps));
     dx.device(*place) = dx * norm_dy.broadcast(bcast) * xr.sign();
   }
 }

paddle/phi/kernels/cpu/p_norm_kernel.cc

@@ -44,7 +44,7 @@ inline void GetDims(
 template <typename T, typename Context>
 void PNormKernel(const Context& dev_ctx,
                  const DenseTensor& x,
-                 float porder,
+                 double porder,
                  int axis,
                  float epsilon UNUSED,
                  bool keepdim UNUSED,
@@ -90,7 +90,7 @@ void PNormKernel(const Context& dev_ctx,
   } else if (porder == -INFINITY) {
     norm.device(*place) = xr.abs().minimum(rdim);
   } else {
-    norm.device(*place) = xr.abs().pow(porder).sum(rdim).pow(1.0f / porder);
+    norm.device(*place) = xr.abs().pow(porder).sum(rdim).pow(1.0 / porder);
   }
 }
 }  // namespace phi