[Comp] Add p_norm_grad in details.h (#69675)

HydrogenSulfate · cubehan3 · web-flow · commit 86e1ead0dc81 · 2024-12-19T11:03:33.000+08:00
* add p_norm in details.h

* update code

* use shape64 instead of shapee

* add p_norm to ALLOW_DYNAMIC_SHAPE_VJP_OPS and add unittest

* add more unitest

* fix op_name

* fix

* fix take_along_axis_grad when duplicated entries in indices

* update unitest

* reduce matrix scale

* Update paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h

Co-authored-by: cubehan3 &lt;hanqiukun@baidu.com&gt;

* remove VLOG

* Update paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h

Co-authored-by: cubehan3 &lt;hanqiukun@baidu.com&gt;

* Update paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h

Co-authored-by: cubehan3 &lt;hanqiukun@baidu.com&gt;

* fix p_norm_decomp

---------

Co-authored-by: cubehan3 &lt;hanqiukun@baidu.com&gt;
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -132,6 +132,7 @@
     'tanh_grad',
     'transpose_grad',
     'unsqueeze_grad',
+    'p_norm_grad',
 ]
 
 # declare belongs to codegen, but implementation not
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -157,6 +157,7 @@
     'stack_grad',
     'swish_grad',
     'swiglu_grad',
+    'p_norm_grad',
 ]  # custom vjp list of composite op
 
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h b/paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h
@@ -112,37 +112,44 @@ Tensor p_norm_decomp(const Tensor& x,
                      const float epsilon = 1.0e-12f,
                      const bool& keepdim = false,
                      const bool& asvector = false) {
+  // NOTE: if asvector is True, then axis will be ignored
+  // and will reduce all elements in x
+
   auto x_tmp = ConvertToMT<T>(x);
 
   Tensor res;
+  std::vector<int> reduce_axis = {};
+  if (!asvector) {
+    reduce_axis.push_back(axis);
+  }
   if (porder == 0.0) {
     // 0-norm
     auto zero = full_scalar<T>(0, x_tmp.dtype(), x_tmp.place());
     auto none_zero = not_equal<T>(x_tmp, zero);
     res = cast<T>(none_zero, x_tmp.dtype());
-    res = sum<T>(res, {axis}, x_tmp.dtype(), keepdim);
+    res = sum<T>(res, reduce_axis, x_tmp.dtype(), keepdim);
   } else if (porder == 1.0) {
     // 1-norm
     res = abs<T>(x_tmp);
-    res = sum<T>(res, {axis}, x_tmp.dtype(), keepdim);
+    res = sum<T>(res, reduce_axis, x_tmp.dtype(), keepdim);
   } else if (porder == 2.0) {
     // 2-norm
-    res = sqrt<T>(sum<T>(x_tmp * x_tmp, {axis}, x_tmp.dtype(), keepdim));
+    res = sqrt<T>(sum<T>(x_tmp * x_tmp, reduce_axis, x_tmp.dtype(), keepdim));
   } else if (porder == INFINITY) {
     // +INF-norm
     res = abs<T>(x_tmp);
-    res = max<T>(x_tmp, {axis}, keepdim);
+    res = max<T>(x_tmp, reduce_axis, keepdim);
   } else if (porder == -INFINITY) {
     // -INF-norm
     res = abs<T>(x_tmp);
-    res = min<T>(x_tmp, {axis}, keepdim);
+    res = min<T>(x_tmp, reduce_axis, keepdim);
   } else {
     // vanilla p-norm
     auto porder_tensor = full_scalar<T>(porder, x_tmp.dtype(), x_tmp.place());
     auto inv_porder_tensor =
         full_scalar<T>(1 / porder, x_tmp.dtype(), x_tmp.place());
     res = elementwise_pow<T>(x_tmp, porder_tensor);
-    res = sum<T>(res, {axis}, x_tmp.dtype(), keepdim);
+    res = sum<T>(res, reduce_axis, x_tmp.dtype(), keepdim);
     res = elementwise_pow<T>(res, inv_porder_tensor);
   }
 
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -3484,6 +3484,152 @@ void amin_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void p_norm_grad(const Tensor& x,
+                 /*output of forward was reserved for efficient backward*/
+                 const Tensor& out,
+                 const Tensor& out_grad,
+                 float porder,
+                 int axis,
+                 float epsilon,
+                 bool keepdim,
+                 bool asvector,
+                 Tensor* x_grad) {
+  if (x_grad) {
+    if (axis < 0) {
+      axis += x.dims().size();
+    }
+
+    Tensor x_grad_tmp;
+    if (porder == 0.0) {
+      // dx = 0
+      if (has_dynamic_shape(x.shape())) {
+        x_grad_tmp = backend::full_with_tensor<T>(
+            shape64<T>(x), 0, x.dtype(), x.place());
+      } else {
+        x_grad_tmp =
+            full<T>(common::vectorize(x.dims()), 0, x.dtype(), x.place());
+      }
+    } else {
+      /* generic case formula:
+        dx = {
+          dy * y^(1-p) * |x|^(p-1) * sgn(x), if p != +-inf,
+          dy * sgn(x) * (x==y), if p == +-inf.
+        }
+      */
+      Tensor expand_out = out;
+      Tensor expand_out_grad = out_grad;
+      // firstly expand output_grad to same ndim with x for convenience
+      if (!keepdim) {
+        if (has_dynamic_shape(x.shape())) {
+          Tensor expand_shape;
+          if (asvector) {
+            // reduce all dimensions in forward
+            expand_shape = full<T>(std::vector<int64_t>(x.dims().size(), 1),
+                                   1,
+                                   DataType::INT64,
+                                   out_grad.place());
+          } else {
+            // only reduce one dimension in forward
+            expand_shape = shape64<T>(out_grad);
+            std::vector<Tensor> expand_shape_vec;
+            for (int64_t i = 0; i < expand_shape.size(); ++i) {
+              expand_shape_vec.push_back(get_slice<T>(expand_shape, i));
+            }
+            expand_shape_vec.insert(
+                expand_shape_vec.begin() + axis,
+                full<T>({1}, 1, expand_shape.dtype(), expand_shape.place()));
+            expand_shape = concat<T>(expand_shape_vec);
+          }
+          expand_out_grad = backend::reshape<T>(out_grad, expand_shape);
+          expand_out = backend::reshape<T>(out, expand_shape);
+        } else {
+          std::vector<int64_t> expand_shape =
+              common::vectorize(out_grad.dims());
+          if (asvector) {
+            // reduce all dimensions in forward
+            expand_shape = std::vector<int64_t>(x.dims().size(), 1);
+          } else {
+            // only reduce one dimension in forward
+            expand_shape.insert(expand_shape.begin() + axis, 1);
+          }
+          expand_out_grad = reshape<T>(out_grad, expand_shape);
+          expand_out = reshape<T>(out, expand_shape);
+        }
+      }
+
+      if (porder == 1.0) {
+        // dx = dy * sign(x)
+        auto x_sign = sign<T>(x);
+        x_grad_tmp = x_sign * expand_out_grad;
+      } else if (porder == 2.0) {
+        // dx = dy * (x / y)
+        x_grad_tmp = x / expand_out;
+        // fill zero to avoid division by zero
+        Tensor _zero_tensor;
+        if (has_dynamic_shape(x.shape())) {
+          _zero_tensor = backend::full_with_tensor<T>(
+              shape64<T>(x), 0, x.dtype(), x.place());
+        } else {
+          _zero_tensor =
+              full<T>(common::vectorize(x.dims()), 0, x.dtype(), x.place());
+        }
+
+        auto finite_mask = isfinite<T>(x_grad_tmp);
+        x_grad_tmp = where<T>(finite_mask, x_grad_tmp, _zero_tensor);
+        x_grad_tmp = expand_out_grad * (x_grad_tmp);
+
+      } else if (porder == INFINITY || porder == -INFINITY) {
+        // dy * sgn(x) * (x==y), if p == +-inf.
+        auto x_abs = abs<T>(x);
+        auto mask =
+            cast<T>(bitwise_or<T>(equal<T>(x_abs, expand_out), isnan<T>(x_abs)),
+                    expand_out.dtype());
+        auto x_sign = sign<T>(x);
+        x_grad_tmp =
+            x_sign * ((expand_out_grad /
+                       sum<T>(mask, {axis}, expand_out_grad.dtype(), true)) *
+                      mask);
+
+      } else if (porder < 1.0) {
+        // dx = dy * y^(1-p) * |x|^(p-1) * sgn(x)
+        auto x_sign = sign<T>(x);
+        auto x_abs_pow = abs<T>(x);
+        x_abs_pow = x_abs_pow.pow(porder - 1);
+
+        auto x_scaled = x_sign * x_abs_pow;
+        x_grad_tmp = x_scaled * expand_out_grad * expand_out.pow(1 - porder);
+
+      } else if (porder < 2.0) {
+        // dx = dy * y^(1-p) * |x|^(p-1) * sgn(x)
+        auto x_sign = sign<T>(x);
+        auto x_abs_pow = abs<T>(x);
+        x_abs_pow = x_abs_pow.pow(porder - 1);
+
+        // auto scale_v = expand_out_grad / expand_out.pow(porder - 1);
+        // auto _zero_tensor =
+        //     full<T>(common::vectorize(x.dims()), 0.0, x.dtype());
+        // auto out_non_zero_mask = not_equal<T>(expand_out, _zero_tensor);
+        // scale_v = scale_v * cast<T>(out_non_zero_mask, scale_v.dtype());
+        // x_grad_tmp = x_sign * x_abs_pow * scale_v;
+
+        auto scale_v = expand_out_grad * expand_out.pow(1 - porder);
+        x_grad_tmp = x_sign * x_abs_pow * scale_v;
+
+      } else {
+        // dx = dy * y^(1-p) * |x|^(p-1) * sgn(x)
+        auto x_sign = sign<T>(x);
+        auto x_abs_pow = abs<T>(x);
+        x_abs_pow = x_abs_pow.pow(porder - 1);
+
+        auto x_scaled = x_sign * x_abs_pow;
+        x_grad_tmp = x_scaled * expand_out_grad * expand_out.pow(1 - porder);
+      }
+    }
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/primitive/primitive.yaml b/paddle/fluid/primitive/primitive/primitive.yaml
@@ -50,6 +50,7 @@
 - bitwise_xor
 - isinf
 - isnan
+- isfinite
 - remainder
 - scale
 - matmul
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -104,6 +104,7 @@
     "pd_op.trunc",
     "pd_op.unsqueeze",
     "pd_op.where",
+    "pd_op.p_norm",
 ]
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py

Original file line number	Diff line number	Diff line change
`@@ -132,6 +132,7 @@`
`132`	`132`	`'tanh_grad',`
`133`	`133`	`'transpose_grad',`
`134`	`134`	`'unsqueeze_grad',`
	`135`	`+ 'p_norm_grad',`
`135`	`136`	`]`
`136`	`137`
`137`	`138`	`# declare belongs to codegen, but implementation not`
Original file line number	Diff line number	Diff line change
`@@ -104,6 +104,7 @@`
`104`	`104`	`"pd_op.trunc",`
`105`	`105`	`"pd_op.unsqueeze",`
`106`	`106`	`"pd_op.where",`
	`107`	`+ "pd_op.p_norm",`
`107`	`108`	`]`
`108`	`109`
`109`	`110`