[Prim] Add gather nd double grad (#69143)

HydrogenSulfate · web-flow · commit 0d542f917a21 · 2024-11-07T14:17:45.000+08:00
* optimize gather_nd double grad

* add unitest

* update no_need_buffer for index_put_double_grad

* reduce scale of TestGatherNdHighGradCheck check Tensor to avoid OOM

* restore code line

* disable gather_nd_grad for testing gather_nd_double_grad

* use UnchangedInferMeta instead of GatherNdDoubleGradInferMeta

* fix
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -81,6 +81,7 @@
     "where_double_grad",
     "bmm_double_grad",
     "index_put_double_grad",
+    "gather_nd_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
@@ -34,4 +34,5 @@
     'where_grad',
     'bmm_grad',
     'index_put_grad',
+    'gather_nd_grad',
 ]
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
@@ -1049,5 +1049,17 @@ void index_put_double_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void gather_nd_double_grad(const Tensor& grad_out,
+                           const Tensor& index,
+                           const Tensor& grad_x_grad,
+                           Tensor* grad_out_grad) {
+  if (grad_out_grad) {
+    // ddout = gather_nd(ddx, index)
+    auto grad_out_grad_tmp = gather_nd<T>(grad_x_grad, index);
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 }  // namespace prim
 }  // namespace paddle
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -1292,6 +1292,15 @@
   composite : gather_grad(x, index, out_grad, axis, x_grad)
   no_need_buffer : x
 
+- backward_op : gather_nd_double_grad
+  forward : gather_nd_grad (Tensor x, Tensor index, Tensor grad_out) -> Tensor(grad_x)
+  args : (Tensor grad_out, Tensor index, Tensor grad_x_grad)
+  output : Tensor(grad_out_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param: [grad_out]
+  composite : gather_nd_double_grad(grad_out, index, grad_x_grad, grad_out_grad)
+
 - backward_op : gather_nd_grad
   forward : gather_nd (Tensor x, Tensor index) -> Tensor(out)
   args : (Tensor x, Tensor index, Tensor out_grad)
@@ -1302,6 +1311,7 @@
     func : gather_nd_grad
   composite : gather_nd_grad(x, index, out_grad, x_grad)
   no_need_buffer : x
+  backward : gather_nd_double_grad
 
 - backward_op : gaussian_inplace_grad
   forward : gaussian_inplace(Tensor x, float mean=0, float std=1.0, int seed=0) -> Tensor(out)
@@ -1554,15 +1564,15 @@
   inplace : (out_grad -> x_grad)
 
 - backward_op : index_put_double_grad
-  forward : index_put_grad (Tensor x, Tensor[] indices, Tensor value, Tensor out_grad, bool accumulate=false) -> Tensor(grad_x), Tensor(grad_value)
+  forward : index_put_grad (Tensor x, Tensor[] indices, Tensor value, Tensor grad_out, bool accumulate=false) -> Tensor(grad_x), Tensor(grad_value)
   args : (Tensor x, Tensor[] indices, Tensor value, Tensor grad_x_grad, Tensor grad_value_grad, bool accumulate=false)
-  output : Tensor(out_grad_grad)
+  output : Tensor(grad_out_grad)
   infer_meta :
     func : UnchangedInferMeta
     param: [x]
   data_transform :
     skip_transform : indices
-  composite : index_put_double_grad(x, indices, value, grad_x_grad, grad_value_grad, accumulate, out_grad_grad)
+  composite : index_put_double_grad(x, indices, value, grad_x_grad, grad_value_grad, accumulate, grad_out_grad)
   optional: grad_x_grad, grad_value_grad
 
 - backward_op : index_put_grad
diff --git a/test/prim/prim/vjp/test_comp_high_grad.py b/test/prim/prim/vjp/test_comp_high_grad.py
@@ -1212,5 +1212,100 @@ def test_high_grad(self):
                         self.func_triple(p, x_stop, y_stop)
 
 
+@param.parameterized_class(
+    ('x_shape', 'index_shape'),
+    [
+        ([16], [10]),
+        ([16, 16], [20, 2]),
+        ([12, 13, 14], [88, 1]),
+        ([12, 13, 14], [88, 2]),
+        ([12, 13, 14], [88, 3]),
+        ([23, 32, 32], [23 * 32 * 32, 1]),
+        ([23, 32, 32], [23 * 32 * 32, 2]),
+        ([23, 32, 32], [23 * 32 * 32, 3]),
+    ],
+)
+class TestGatherNdHighGradCheck(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.x_shape = cls.x_shape
+        cls.index_shape = cls.index_shape
+
+    def _grad(self, y, x, order):
+        u = y
+        dx = paddle.ones_like(x)
+        for _ in range(order):
+            dx = paddle.grad(u, x, create_graph=True)[0]
+            u = dx
+        return dx
+
+    def func_double(self, place, x_stop, y_stop):
+        x = paddle.randn(self.x_shape).astype("float32").to(device=place)
+        x.stop_gradient = x_stop
+        n_index = self.index_shape[0]
+        index_dim_size = self.index_shape[1] if len(self.index_shape) > 1 else 1
+        index = [
+            paddle.to_tensor(
+                np.random.randint(
+                    -self.x_shape[i], self.x_shape[i], size=[n_index]
+                ),
+                place=place,
+            )
+            for i in range(index_dim_size)
+        ]
+        index = paddle.stack(index, axis=-1)  # [N, D]
+
+        z = paddle.gather_nd(x, index)
+        z = paddle.tanh(z)
+
+        if not x.stop_gradient:
+            dzdx = self._grad(z, x, 2)
+
+    def func_triple(self, place, x_stop, y_stop):
+        x = paddle.randn(self.x_shape).astype("float32").to(device=place)
+        x.stop_gradient = x_stop
+        n_index = self.index_shape[0]
+        index_dim_size = self.index_shape[1] if len(self.index_shape) > 1 else 1
+        index = [
+            paddle.to_tensor(
+                np.random.randint(
+                    -self.x_shape[i], self.x_shape[i], size=[n_index]
+                ),
+                place=place,
+            )
+            for i in range(index_dim_size)
+        ]
+        index = paddle.stack(index, axis=-1)  # [N, D]
+
+        # wraping with tanh to enable high order gradient
+        z = paddle.gather_nd(x, index)
+        z = paddle.tanh(z)
+
+        if not x.stop_gradient:
+            dzdx = self._grad(z, x, 3)
+
+    def test_high_grad(self):
+        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))
+        core._set_prim_backward_enabled(True)
+        core._set_prim_backward_blacklist("gather_nd_grad")
+
+        for p in places:
+            for x_stop in [False, True]:
+                for y_stop in [False, True]:
+                    with dygraph_guard():
+                        self.func_double(p, x_stop, y_stop)
+                        self.func_triple(p, x_stop, y_stop)
+
+        core._set_prim_backward_enabled(False)
+
+
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -81,6 +81,7 @@`
`81`	`81`	`"where_double_grad",`
`82`	`82`	`"bmm_double_grad",`
`83`	`83`	`"index_put_double_grad",`
	`84`	`+ "gather_nd_double_grad",`
`84`	`85`	`]`
`85`	`86`
`86`	`87`	`# white ops list whose kernel can automatically do type promotion.`
Original file line number	Diff line number	Diff line change
`@@ -34,4 +34,5 @@`
`34`	`34`	`'where_grad',`
`35`	`35`	`'bmm_grad',`
`36`	`36`	`'index_put_grad',`
	`37`	`+ 'gather_nd_grad',`
`37`	`38`	`]`