[Prim] Add where double grad (#68125)

HydrogenSulfate · web-flow · commit ad331b5cf9e8 · 2024-09-14T15:23:20.000+08:00
* add where_double_grad and p_norm in eager prim

* add unitest for where comp hight order diff

* add unitest for norm 1-order comp diff

* update yaml and annotations

* add if

* fix typo

* remove p_norm_grad

* update owner of key composite

* restore test_where_op.py

* Update eager_gen.py
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -78,6 +78,7 @@
     "abs_triple_grad",
     "exp_double_grad",
     "log_double_grad",
+    "where_double_grad",
 ]
 
 # white ops list whose kernel can automatically do type promotion.
diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py
@@ -31,4 +31,5 @@
     'silu_grad',
     'exp_grad',
     'abs_double_grad',
+    'where_grad',
 ]
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -26,6 +26,7 @@
 - gather_nd
 - log
 - max
+- min
 - maximum
 - minimum
 - prod
diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h
@@ -158,6 +158,32 @@ void maximum_double_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void where_double_grad(const Tensor& condition,
+                       const paddle::optional<Tensor>& grad_x_grad,
+                       const paddle::optional<Tensor>& grad_y_grad,
+                       Tensor* grad_out_grad) {
+  if (grad_out_grad) {
+    if (grad_x_grad && grad_y_grad) {
+      // ddz = ddx * cond + (1-cond) * ddy
+      auto condition_mask = cast<T>(condition, grad_x_grad.get().dtype());
+      auto ddout = condition_mask * grad_x_grad.get() +
+                   (1 - condition_mask) * grad_y_grad.get();
+      set_output<T>(ddout, grad_out_grad);
+    } else if (grad_x_grad) {
+      // ddz = ddx * cond
+      auto condition_mask = cast<T>(condition, grad_x_grad.get().dtype());
+      auto ddout = condition_mask * grad_x_grad.get();
+      set_output<T>(ddout, grad_out_grad);
+    } else if (grad_y_grad) {
+      // ddz = (1-cond) * ddy
+      auto condition_mask = cast<T>(condition, grad_y_grad.get().dtype());
+      auto ddout = (1 - condition_mask) * grad_y_grad.get();
+      set_output<T>(ddout, grad_out_grad);
+    }
+  }
+}
+
 template <typename T>
 void tanh_triple_grad(const Tensor& out,
                       const Tensor& grad_out_forward,
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -3580,6 +3580,7 @@
   kernel :
     func : where_grad
   no_need_buffer : x, y
+  backward : where_double_grad
 
 - backward_op : yolo_loss_grad
   forward : yolo_loss (Tensor x, Tensor gt_box, Tensor gt_label, Tensor gt_score, int[] anchors={}, int[] anchor_mask={}, int class_num =1 , float ignore_thresh=0.7, int downsample_ratio=32, bool use_label_smooth=true, float scale_x_y=1.0) -> Tensor(loss), Tensor(objectness_mask), Tensor(gt_match_mask)
@@ -3709,3 +3710,13 @@
   kernel:
     func: unpool_grad
     data_type: x
+
+- backward_op: where_double_grad
+  forward: where_grad (Tensor condition, Tensor x, Tensor y, Tensor grad_out) -> Tensor(grad_x), Tensor(grad_y)
+  args: (Tensor condition, Tensor grad_x_grad, Tensor grad_y_grad)
+  output: Tensor(grad_out_grad)
+  infer_meta :
+    func : GeneralUnaryGradInferMeta
+    param : [condition]
+  composite: where_double_grad(condition, grad_x_grad, grad_y_grad, grad_out_grad)
+  optional: grad_x_grad, grad_y_grad
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -3994,7 +3994,7 @@
     {warpctcgrad : WarpCTCGrad, loss : Loss}
 
 - op : where
-  backward : where_grad
+  backward : where_grad, where_double_grad
   inputs :
     {condition : Condition, x : X, y : Y}
   outputs :
diff --git a/test/prim/prim/vjp/test_comp_high_grad.py b/test/prim/prim/vjp/test_comp_high_grad.py
@@ -437,6 +437,129 @@ def test_high_grad(self):
             self.func_triple(p)
 
 
+@param.parameterized_class(
+    ('shape1'),
+    [
+        ([2],),
+        ([2, 3],),
+        ([2, 3, 4],),
+        ([2, 3, 3, 4],),
+    ],
+)
+class TestWhereHighGradCheck(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.shape1 = cls.shape1
+
+    def _where_wrapper(self, x):
+        return paddle.where(x[0], x[1], x[2])
+
+    @prog_scope()
+    def _func_double(self, place):
+        shape1 = self.shape1
+        eps = 0.0005
+        dtype = np.float64
+
+        cond = paddle.static.data('cond', shape1, dtype="bool")
+        cond.stop_gradient = True
+        cond.persistable = True
+
+        x = paddle.static.data('x', shape1, dtype=dtype)
+        x.stop_gradient = False
+        x.persistable = True
+
+        y = paddle.static.data('y', shape1, dtype=dtype)
+        y.stop_gradient = False
+        y.persistable = True
+
+        out = paddle.where(cond, x, y)
+
+        # generate random data
+        cond_arr = np.random.choice([False, True], size=shape1)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+
+        # where double grad only has CompositeOpMaker,don't need set prim_flag
+        from paddle.base import core
+
+        core._set_prim_backward_enabled(True)
+        gradient_checker.double_grad_check(
+            [cond, x, y],
+            y=out,
+            x_init=[cond_arr, x_arr, y_arr],
+            place=place,
+            eps=eps,
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self._where_wrapper,
+            [cond, x, y],
+            y=out,
+            x_init=[cond_arr, x_arr, y_arr],
+            place=place,
+        )
+        core._set_prim_backward_enabled(False)
+
+    @prog_scope()
+    def _func_triple(self, place):
+        shape1 = self.shape1
+        eps = 0.0005
+        dtype = np.float64
+
+        cond = paddle.static.data('cond', shape1, dtype="bool")
+        cond.stop_gradient = True
+        cond.persistable = True
+
+        x = paddle.static.data('x', shape1, dtype=dtype)
+        x.stop_gradient = False
+        x.persistable = True
+
+        y = paddle.static.data('y', shape1, dtype=dtype)
+        y.stop_gradient = False
+        y.persistable = True
+
+        out = paddle.where(cond, x, y)
+
+        # generate random data
+        cond_arr = np.random.choice([False, True], size=shape1)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+
+        # where double grad only has CompositeOpMaker,don't need set prim_flag
+        from paddle.base import core
+
+        core._set_prim_backward_enabled(True)
+        gradient_checker.triple_grad_check(
+            [cond, x, y],
+            y=out,
+            x_init=[cond_arr, x_arr, y_arr],
+            place=place,
+            eps=eps,
+        )
+        gradient_checker.triple_grad_check_for_dygraph(
+            self.silu_wrapper,
+            [cond, x, y],
+            y=out,
+            x_init=[cond_arr, x_arr, y_arr],
+            place=place,
+        )
+        core._set_prim_backward_enabled(False)
+
+    def test_high_grad(self):
+        paddle.enable_static()
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not core.is_compiled_with_cuda()
+        ):
+            places.append(base.CPUPlace())
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self._func_double(p)
+            self._func_triple(p)
+
+
 @param.parameterized_class(
     ('shape1'),
     [
diff --git a/tools/check_file_diff_approvals.sh b/tools/check_file_diff_approvals.sh
@@ -407,8 +407,8 @@ for CHANGE_FILE in ${ALL_CHANGE_YAML_FILES}; do
     fi
 done
 if [ "${BAN_COMP_MESSAGE}" != "" ] && [ "${GIT_PR_ID}" != "" ]; then
-    echo_line="If you need to change the key composite, you must have one RD (Charles-hit(wanghao), cyber-pioneer(chenzhuo), cxxly(chenxiaoxu)) review and approve. \nThe code that do not meet the specification are as follows:\n${BAN_COMP_MESSAGE}\n"
-    check_approval 1 Charles-hit cyber-pioneer cxxly
+    echo_line="If you need to change the key composite, you must have one RD (xiaoguoguo626807(wangruting), cubehan3(hanqiukun)) review and approve. \nThe code that do not meet the specification are as follows:\n${BAN_COMP_MESSAGE}\n"
+    check_approval 1 xiaoguoguo626807 cubehan3
 fi
 
 NEW_OP_ADDED=`git diff --name-only --diff-filter=A upstream/$BRANCH |grep -oE ".+_op..*" || true`

Original file line number	Diff line number	Diff line change
`@@ -78,6 +78,7 @@`
`78`	`78`	`"abs_triple_grad",`
`79`	`79`	`"exp_double_grad",`
`80`	`80`	`"log_double_grad",`
	`81`	`+ "where_double_grad",`
`81`	`82`	`]`
`82`	`83`
`83`	`84`	`# white ops list whose kernel can automatically do type promotion.`
Original file line number	Diff line number	Diff line change
`@@ -31,4 +31,5 @@`
`31`	`31`	`'silu_grad',`
`32`	`32`	`'exp_grad',`
`33`	`33`	`'abs_double_grad',`
	`34`	`+ 'where_grad',`
`34`	`35`	`]`
-Original file line number
+Diff line change
 - gather_nd
 - log
 - max
 +- min
 - maximum
 - minimum
 - prod