[Prim] Add index_add_double_grad in dygraph composite (#70367)

HydrogenSulfate · web-flow · commit 5317b5acd121 · 2024-12-23T11:26:38.000+08:00
* add VLOG for meminfo

* fix index_add_grad nullptr gradient bug

* fix

* add index_add_double_grad in dygraph composite

* update eager_gen

* update UT

* fix for xpu
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -84,6 +84,7 @@
     "gather_nd_double_grad",
     "reshape_double_grad",
     "take_along_axis_double_grad",
+    "index_add_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
@@ -36,4 +36,5 @@
     'index_put_grad',
     'gather_nd_grad',
     'take_along_axis_grad',
+    'index_add_grad',
 ]
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -52,6 +52,7 @@
 - sign
 - sigmoid
 - index_put
+- index_add
 - isnan
 - isfinite
 - take_along_axis
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
@@ -1107,5 +1107,43 @@ void take_along_axis_double_grad(const Tensor& indices,
   }
 }
 
+template <typename T>
+void index_add_double_grad(const Tensor& index,
+                           const Tensor& out_grad,
+                           const paddle::optional<Tensor>& grad_x_grad,
+                           const paddle::optional<Tensor>& grad_add_value_grad,
+                           int axis,
+                           Tensor* grad_out_grad) {
+  if (grad_out_grad) {
+    if (grad_x_grad && grad_add_value_grad) {
+      Tensor grad_out_grad_tmp = grad_x_grad.get();
+      grad_out_grad_tmp = index_add<T>(
+          grad_out_grad_tmp, index, grad_add_value_grad.get(), axis);
+      set_output<T>(grad_out_grad_tmp, grad_out_grad);
+
+    } else if (grad_x_grad) {
+      Tensor grad_out_grad_tmp = grad_x_grad.get();
+      set_output<T>(grad_out_grad_tmp, grad_out_grad);
+
+    } else if (grad_add_value_grad) {
+      Tensor DDadd_value = grad_add_value_grad.get();
+      Tensor grad_out_grad_tmp = full<T>(common::vectorize(out_grad.dims()),
+                                         0,
+                                         DDadd_value.dtype(),
+                                         DDadd_value.place());
+      grad_out_grad_tmp =
+          index_add<T>(grad_out_grad_tmp, index, DDadd_value, axis);
+      set_output<T>(grad_out_grad_tmp, grad_out_grad);
+
+    } else {
+      Tensor grad_out_grad_tmp = full<T>(common::vectorize(out_grad.dims()),
+                                         0,
+                                         out_grad.dtype(),
+                                         out_grad.place());
+      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
@@ -1565,6 +1565,18 @@
     func : imag_grad
     data_type : complex(out_grad)
 
+- backward_op : index_add_double_grad
+  forward : index_add_grad (Tensor index, Tensor add_value, Tensor grad_out, int axis) -> Tensor(grad_x), Tensor(grad_add_value)
+  args : (Tensor index, Tensor grad_out, Tensor grad_x_grad, Tensor grad_add_value_grad, int axis)
+  output : Tensor(grad_out_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param: [grad_out]
+  data_transform :
+    skip_transform : index
+  composite : index_add_double_grad(index, grad_out, grad_x_grad, grad_add_value_grad, axis, grad_out_grad)
+  optional: grad_x_grad, grad_add_value_grad
+
 - backward_op : index_add_grad
   forward : index_add(Tensor x, Tensor index,  Tensor add_value, int axis=0) -> Tensor(out)
   args : (Tensor index, Tensor add_value, Tensor out_grad, int axis)
@@ -1575,6 +1587,7 @@
     func : index_add_grad
     data_type : out_grad
   inplace : (out_grad -> x_grad)
+  backward : index_add_double_grad
 
 - backward_op : index_put_double_grad
   forward : index_put_grad (Tensor x, Tensor[] indices, Tensor value, Tensor grad_out, bool accumulate=false) -> Tensor(grad_x), Tensor(grad_value)
diff --git a/test/prim/prim/vjp/eager/test_comp_eager_index_add_double_grad.py b/test/prim/prim/vjp/eager/test_comp_eager_index_add_double_grad.py
@@ -0,0 +1,58 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+
+
+@param.parameterized_class(
+    ('x', 'index', 'axis', 'value', 'cotangent', 'dtype'),
+    [
+        (
+            np.random.randn(4, 3, 2),  # x
+            np.random.randint(-3, 3, size=(16,)),  # index
+            1,  # axis
+            np.random.randint(0, 3, size=(4, 16, 2)),  # valie
+            np.random.rand(4, 3, 2),  # cotangent
+            np.float32,  # dtype
+        ),
+    ],
+)
+class TestTakeAlongAxisTanhDoubleGrad(unittest.TestCase):
+    def test_index_add_tanh_double_grad(self):
+        x_tensor = paddle.to_tensor(
+            self.x, dtype=self.dtype, stop_gradient=False
+        )
+        value_tensor = paddle.to_tensor(
+            self.value, dtype=self.dtype, stop_gradient=False
+        )
+        index_tensor = paddle.to_tensor(self.index, dtype="int64")
+        dout_tensor = paddle.to_tensor(
+            self.cotangent, dtype=self.dtype, stop_gradient=False
+        )
+        out = paddle.index_add(x_tensor, index_tensor, self.axis, value_tensor)
+
+        out = paddle.tanh(out)
+
+        dx = paddle.grad(out, x_tensor, dout_tensor, create_graph=True)[0]
+
+        ddx = paddle.grad(dx, dout_tensor, create_graph=True)[0]
+
+
+if __name__ == '__main__':
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -84,6 +84,7 @@`
`84`	`84`	`"gather_nd_double_grad",`
`85`	`85`	`"reshape_double_grad",`
`86`	`86`	`"take_along_axis_double_grad",`
	`87`	`+ "index_add_double_grad",`
`87`	`88`	`]`
`88`	`89`
`89`	`90`	`# white ops list whose kernel can automatically do type promotion.`
Original file line number	Diff line number	Diff line change
`@@ -36,4 +36,5 @@`
`36`	`36`	`'index_put_grad',`
`37`	`37`	`'gather_nd_grad',`
`38`	`38`	`'take_along_axis_grad',`
	`39`	`+ 'index_add_grad',`
`39`	`40`	`]`