add tanh_triple_grad composite logic (#56072)

lxd-cumt · web-flow · commit 7c4a355642e8 · 2023-08-10T16:22:18.000+08:00
* decompose tanh_triple_grad and add it into prim_white_list test=develop

* fix TanhTripleGradKernel bugs test=develop

* decompose tanh_triple_grad test=develop
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -72,7 +72,7 @@
     "subtract_double_grad",
     "add_triple_grad",
     "silu_double_grad",
-    "tanh_double_grad",
+    "tanh_triple_grad",
 ]
 
 # dict of special api that forward api's output will affect bacward api's output
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
@@ -53,6 +53,76 @@ void tanh_double_grad(const Tensor& out,
   }
 }
 
+template <typename T>
+void tanh_triple_grad(const Tensor& out,
+                      const Tensor& grad_out_forward,
+                      const Tensor& grad_x_grad_forward,
+                      const paddle::optional<Tensor>& grad_out_new_grad,
+                      const paddle::optional<Tensor>& grad_out_grad_grad,
+                      Tensor* out_grad,
+                      Tensor* grad_out_forward_grad,
+                      Tensor* grad_x_grad_forward_grad) {
+  if (out_grad) {
+    if (grad_out_grad_grad) {
+      if (grad_out_new_grad) {
+        auto out_grad_tmp =
+            (-2 * out * grad_x_grad_forward * grad_out_grad_grad.get()) -
+            (2 * grad_out_forward * grad_x_grad_forward *
+             grad_out_new_grad.get());
+        set_output<T>(out_grad_tmp, out_grad);
+      } else {
+        auto out_grad_tmp =
+            -2 * out * grad_x_grad_forward * grad_out_grad_grad.get();
+        set_output<T>(out_grad_tmp, out_grad);
+      }
+    } else {
+      if (grad_out_new_grad) {
+        auto out_grad_tmp = -(2 * grad_out_forward * grad_x_grad_forward *
+                              grad_out_new_grad.get());
+        set_output<T>(out_grad_tmp, out_grad);
+      } else {
+        auto out_grad_tmp = 0 * out;
+        set_output<T>(out_grad_tmp, out_grad);
+      }
+    }
+  }
+
+  if (grad_out_forward_grad) {
+    if (grad_out_new_grad) {
+      auto grad_out_forward_grad_tmp =
+          -2 * out * grad_x_grad_forward * grad_out_new_grad.get();
+      set_output<T>(grad_out_forward_grad_tmp, grad_out_forward_grad);
+    } else {
+      auto grad_out_forward_grad_tmp = 0 * out;
+      set_output<T>(grad_out_forward_grad_tmp, grad_out_forward_grad);
+    }
+  }
+
+  if (grad_x_grad_forward_grad) {
+    if (grad_out_grad_grad) {
+      if (grad_out_new_grad) {
+        auto grad_x_grad_forward_grad_tmp =
+            (1 - (out * out)) * grad_out_grad_grad.get() -
+            2 * out * grad_out_forward * grad_out_new_grad.get();
+        set_output<T>(grad_x_grad_forward_grad_tmp, grad_x_grad_forward_grad);
+      } else {
+        auto grad_x_grad_forward_grad_tmp =
+            (1 - (out * out)) * grad_out_grad_grad.get();
+        set_output<T>(grad_x_grad_forward_grad_tmp, grad_x_grad_forward_grad);
+      }
+    } else {
+      if (grad_out_new_grad) {
+        auto grad_x_grad_forward_grad_tmp =
+            -(2 * out * grad_out_forward * grad_out_new_grad.get());
+        set_output<T>(grad_x_grad_forward_grad_tmp, grad_x_grad_forward_grad);
+      } else {
+        auto grad_x_grad_forward_grad_tmp = 0 * grad_x_grad_forward;
+        set_output<T>(grad_x_grad_forward_grad_tmp, grad_x_grad_forward_grad);
+      }
+    }
+  }
+}
+
 template <typename T>
 void matmul_double_grad(const Tensor& x,
                         const Tensor& y,
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -2345,6 +2345,7 @@
     param : [out, out, grad_x_grad_forward]
   kernel :
     func : tanh_triple_grad
+  composite : tanh_triple_grad(out, grad_out_forward, grad_x_grad_forward, grad_out_new_grad, grad_out_grad_grad, out_grad, grad_out_forward_grad, grad_x_grad_forward_grad)
   inplace : (grad_x_grad_forward -> grad_out_forward_grad)
   optional : grad_out_new_grad, grad_out_grad_grad
 
diff --git a/paddle/phi/kernels/impl/activation_grad_impl.h b/paddle/phi/kernels/impl/activation_grad_impl.h
@@ -189,11 +189,11 @@ void TanhTripleGradKernel(const Context& dev_ctx,
     dev_ctx.template Alloc<T>(d_dout);
   }
   if (d_out_new) {
-    d_dout->Resize(out.dims());
+    d_out_new->Resize(out.dims());
     dev_ctx.template Alloc<T>(d_out_new);
   }
   if (d_ddx) {
-    d_dout->Resize(ddx.dims());
+    d_ddx->Resize(ddx.dims());
     dev_ctx.template Alloc<T>(d_ddx);
   }
   funcs::TanhTripleGradFunctor<T> functor;

Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@`
`72`	`72`	`"subtract_double_grad",`
`73`	`73`	`"add_triple_grad",`
`74`	`74`	`"silu_double_grad",`
`75`		`- "tanh_double_grad",`
	`75`	`+ "tanh_triple_grad",`
`76`	`76`	`]`
`77`	`77`
`78`	`78`	`# dict of special api that forward api's output will affect bacward api's output`
Original file line number	Diff line number	Diff line change
`@@ -189,11 +189,11 @@ void TanhTripleGradKernel(const Context& dev_ctx,`
`189`	`189`	`dev_ctx.template Alloc<T>(d_dout);`
`190`	`190`	`}`
`191`	`191`	`if (d_out_new) {`
`192`		`- d_dout->Resize(out.dims());`
	`192`	`+ d_out_new->Resize(out.dims());`
`193`	`193`	`dev_ctx.template Alloc<T>(d_out_new);`
`194`	`194`	`}`
`195`	`195`	`if (d_ddx) {`
`196`		`- d_dout->Resize(ddx.dims());`
	`196`	`+ d_ddx->Resize(ddx.dims());`
`197`	`197`	`dev_ctx.template Alloc<T>(d_ddx);`
`198`	`198`	`}`
`199`	`199`	`funcs::TanhTripleGradFunctor<T> functor;`