[Phi] Change the output format of C++ backward api (Part2)

paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py

@@ -24,10 +24,11 @@
 ops_to_fill_zero_for_empty_grads = set([
     "split_grad", "rnn_grad", "matmul_double_grad", "matmul_triple_grad",
     "sigmoid_double_grad", "sigmoid_triple_grad", "add_double_grad",
-    "add_triple_grad", "multiply_double_grad", "multiply_triple_grad",
-    "conv2d_grad_grad", "batch_norm_double_grad", "tanh_double_grad",
-    "tanh_triple_grad", "subtract_double_grad", "divide_double_grad",
-    "log_double_grad", "elu_double_grad", "leaky_relu_double_grad"
+    "add_triple_grad", "multiply_grad", "multiply_double_grad",
+    "multiply_triple_grad", "conv2d_grad_grad", "batch_norm_double_grad",
+    "tanh_double_grad", "tanh_triple_grad", "subtract_double_grad",
+    "divide_double_grad", "log_double_grad", "elu_double_grad",
+    "leaky_relu_double_grad"
 ])
 
 # For API dispatch used at python-level

paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py

@@ -1321,7 +1321,11 @@ def GenerateNodeDefinition(self, grad_node_creation_str,
     }}
     api_output[i].reserve(returns[i].size());
     for (size_t j = 0; j < returns[i].size(); ++j) {{
-      api_output[i].push_back(&returns[i][j]);
+      if (out_metas[i][j].IsStopGradient()) {{
+        api_output[i].push_back(nullptr);
+      }} else {{
+        api_output[i].push_back(&returns[i][j]);
+      }}
     }}
   }}
 """
@@ -1392,7 +1396,7 @@ def GenerateNodeDefinition(self, grad_node_creation_str,
                 if IsPlainTensorType(rtype):
                     output_autograd_meta = f"""
   auto& {transformed_tensor_name} = returns[{pos}][0];
-  egr::AutogradMeta* {output_autograd_meta_name} = egr::EagerUtils::autograd_meta(&{transformed_tensor_name});"""
+  egr::AutogradMeta* {output_autograd_meta_name} = returns[{pos}][0].initialized() ? egr::EagerUtils::autograd_meta(&{transformed_tensor_name}) : nullptr;"""
 
                 else:
                     assert IsVectorTensorType(rtype)

paddle/fluid/eager/grad_node_info.cc

@@ -218,6 +218,8 @@ void GradNodeBase::SetGradOutMeta(const paddle::experimental::Tensor& fwd_in,
   // Set Stop_gradient
   if (fwd_in_meta) {
     meta.SetStopGradient(fwd_in_meta->StopGradient());
+  } else {
+    meta.SetStopGradient(true);
   }
   // Set Adj Edges
   if (fwd_in_meta && !fwd_in_meta->StopGradient()) {

paddle/fluid/eager/utils.cc

@@ -165,12 +165,14 @@ void EagerUtils::SetHistory(std::vector<AutogradMeta*>* autograd_metas,
 
 void EagerUtils::SetHistory(AutogradMeta* autograd_meta,
                             const std::shared_ptr<GradNodeBase>& grad_node) {
-  if (autograd_meta->GradNode()) {
-    VLOG(7) << "Should not set grad node twice, original node is:"
-            << autograd_meta->GradNode()->name()
-            << "current is: " << grad_node->name();
+  if (autograd_meta) {
+    if (autograd_meta->GradNode()) {
+      VLOG(7) << "Should not set grad node twice, original node is:"
+              << autograd_meta->GradNode()->name()
+              << "current is: " << grad_node->name();
+    }
+    autograd_meta->SetGradNode(grad_node);
   }
-  autograd_meta->SetGradNode(grad_node);
 }
 
 void EagerUtils::SetOutRankWithSlot(std::vector<AutogradMeta*>* targets,
@@ -181,7 +183,7 @@ void EagerUtils::SetOutRankWithSlot(std::vector<AutogradMeta*>* targets,
   }
 }
 void EagerUtils::SetOutRankWithSlot(AutogradMeta* target, size_t slot_id) {
-  target->SetSingleOutRankWithSlot(slot_id, 0);
+  if (target) target->SetSingleOutRankWithSlot(slot_id, 0);
 }
 
 std::shared_ptr<egr::EagerVariable> EagerUtils::TrySyncToVar(
@@ -450,17 +452,13 @@ void EagerUtils::FillZeroForEmptyGradInputs(
       paddle::experimental::Tensor& grad = (*in_grads)[i][j];
       if (!grad.initialized()) {
         const GradSlotMeta& grad_in_meta = grad_in_metas[i][j];
-        PADDLE_ENFORCE(
-            grad_in_meta.HasTensorMeta(),
-            paddle::platform::errors::Fatal(
-                "Unable to fill empty grad inputs due to empty GradSlotMeta"));
-
-        const auto& tensor_meta = grad_in_meta.GetTensorMeta();
-        phi::Place place = grad_in_meta.GetPlace();
-
-        auto tensor_with_zero = paddle::experimental::full(
-            phi::vectorize(tensor_meta.dims), 0.0, tensor_meta.dtype, place);
-        grad.set_impl(tensor_with_zero.impl());
+        if (grad_in_meta.HasTensorMeta()) {
+          const auto& tensor_meta = grad_in_meta.GetTensorMeta();
+          auto tensor_with_zero = paddle::experimental::full(
+              phi::vectorize(tensor_meta.dims), 0.0, tensor_meta.dtype,
+              grad_in_meta.GetPlace());
+          grad.set_impl(tensor_with_zero.impl());
+        }
       }
     }
   }

paddle/phi/api/lib/kernel_dispatch.h

@@ -109,7 +109,12 @@ struct KernelKeyParser : ArgsIterator<KernelKeyParser> {
     }
   }
 
-  void operator()(const Tensor& x) { AssignKernelKeySet(*x.impl()); }
+  void operator()(const Tensor& x) {
+    const auto* tensor = x.impl().get();
+    if (tensor) {
+      AssignKernelKeySet(*tensor);
+    }
+  }
 
   void operator()(const std::vector<Tensor>& x) {
     const phi::TensorBase& tensor = *x.at(0).impl();

paddle/phi/kernels/activation_grad_kernel.h

@@ -17,6 +17,7 @@ limitations under the License. */
 #include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/infermeta/unary.h"
+#include "paddle/utils/optional.h"
 
 namespace phi {
 
@@ -136,7 +137,7 @@ void SigmoidTripleGradKernel(const Context& dev_ctx,
                              const DenseTensor& dout,
                              const DenseTensor& ddx,
                              const DenseTensor& d_dout_new,
-                             const DenseTensor& d_ddout,
+                             paddle::optional<const DenseTensor&> d_ddout,
                              DenseTensor* d_out_new,
                              DenseTensor* d_dout,
                              DenseTensor* d_ddx);

paddle/phi/kernels/funcs/activation_functor.h

@@ -1428,16 +1428,19 @@ struct SigmoidTripleGradFunctor : public BaseActivationFunctor<T> {
         GET_DATA_SAFELY(Out, "Input", "Out", "SigmoidTripleGrad"));
     auto dout = EigenVector<T>::Flatten(
         GET_DATA_SAFELY(dOut, "Input", "DOut", "SigmoidTripleGrad"));
-    auto d_ddOut = EigenVector<T>::Flatten(
-        GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "SigmoidTripleGrad"));
     auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
         d_dOut_New, "Input", "D_DOut_New", "SigmoidTripleGrad"));
 
     if (d_Out_New) {
       auto d_OutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
           d_Out_New, "Output", "D_OutNew", "SigmoidTripleGrad"));
-      d_OutNew.device(*d) = (ddx - static_cast<T>(2) * out * ddx) * d_ddOut -
-                            static_cast<T>(2) * dout * ddx * d_dOutNew;
+      d_OutNew.device(*d) = -static_cast<T>(2) * dout * ddx * d_dOutNew;
+      if (d_DDOut) {
+        auto d_ddOut = EigenVector<T>::Flatten(
+            GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "SigmoidTripleGrad"));
+        d_OutNew.device(*d) =
+            (ddx - static_cast<T>(2) * out * ddx) * d_ddOut + d_OutNew;
+      }
     }
     if (d_d_Out) {
       auto d_dOut = EigenVector<T>::Flatten(
@@ -1449,8 +1452,12 @@ struct SigmoidTripleGradFunctor : public BaseActivationFunctor<T> {
       auto d_ddx = EigenVector<T>::Flatten(
           GET_DATA_SAFELY(d_DDx, "Output", "D_DDx", "SigmoidTripleGrad"));
       d_ddx.device(*d) =
-          (static_cast<T>(1) - out) * out * d_ddOut +
           (static_cast<T>(1) - static_cast<T>(2) * out) * dout * d_dOutNew;
+      if (d_DDOut) {
+        auto d_ddOut = EigenVector<T>::Flatten(
+            GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "SigmoidTripleGrad"));
+        d_ddx.device(*d) = d_ddx + (static_cast<T>(1) - out) * out * d_ddOut;
+      }
     }
   }
   static constexpr ActBwdOpFwdDeps FwdDeps() {

paddle/phi/kernels/impl/activation_grad_impl.h

@@ -265,7 +265,7 @@ void SigmoidTripleGradKernel(const Context& dev_ctx,
                              const DenseTensor& dout,
                              const DenseTensor& ddx,
                              const DenseTensor& d_dout_new,
-                             const DenseTensor& d_ddout,
+                             paddle::optional<const DenseTensor&> d_ddout,
                              DenseTensor* d_out_new,
                              DenseTensor* d_dout,
                              DenseTensor* d_ddx) {
@@ -274,19 +274,19 @@ void SigmoidTripleGradKernel(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::SigmoidTripleGradFunctor<T> functor;
   functor(dev_ctx,
           &out,
           &ddx,
           &dout,
-          &d_ddout,
+          d_ddout.get_ptr(),
           &d_dout_new,
           d_dout,
           d_out_new,

python/paddle/utils/code_gen/backward.yaml

@@ -1659,6 +1659,7 @@
     param : [out, fwd_grad_out, grad_grad_x]
   kernel :
     func : sigmoid_triple_grad
+  optional : grad_grad_out_grad
 
 - backward_api : silu_grad
   forward : silu (Tensor x) -> Tensor(out)