Support PyLayer's output as None (#69674)

xuxinyi389 · web-flow · commit 1dcab8211ba5 · 2024-12-03T15:28:46.000+08:00
* tmp

* fix_range

* fix_test_part_I

* support_none

* polish
diff --git a/paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.cc b/paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.cc
@@ -58,10 +58,11 @@ PyLayerInstruction::PyLayerInstruction(
 
   SetKernelType(AnalyseOpFuncType(op, place));
   VLOG(6) << "finish process analyse kernel type";
-
   for (size_t i = 0; i < pylayer_op.num_results(); ++i) {
-    output_vars_.push_back(value_exec_info->GetScope()->GetVar(
-        value_exec_info->GetValue2VarName().at(pylayer_op.result(i))));
+    if (pylayer_op.result(i) && pylayer_op.result(i).type()) {
+      output_vars_.push_back(value_exec_info->GetScope()->GetVar(
+          value_exec_info->GetValue2VarName().at(pylayer_op.result(i))));
+    }
   }
   VLOG(6) << "finish process output_vars";
 
diff --git a/paddle/fluid/framework/new_executor/instruction/control_flow/yield_instruction.cc b/paddle/fluid/framework/new_executor/instruction/control_flow/yield_instruction.cc
@@ -38,13 +38,18 @@ YieldInstruction::YieldInstruction(size_t id,
       continue;
     }
     auto in = op->operand_source(i);
-    inputs.emplace(in, GetValueIds(in, *value_exe_info));
-    input_vars_.push_back(value_exe_info->GetVarByValue(in));
+    if (in && in.type()) {
+      inputs.emplace(in, GetValueIds(in, *value_exe_info));
+      input_vars_.push_back(value_exe_info->GetVarByValue(in));
+    }
   }
   SetInputs(inputs);
 
   for (size_t i = 0; i < parent_op->num_results(); ++i) {
-    output_vars_.push_back(value_exe_info->GetVarByValue(parent_op->result(i)));
+    if (parent_op->result(i) && parent_op->result(i).type()) {
+      output_vars_.push_back(
+          value_exe_info->GetVarByValue(parent_op->result(i)));
+    }
   }
 
   PADDLE_ENFORCE_EQ(
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.cc
@@ -84,14 +84,17 @@ void PyLayerOp::Build(pir::Builder &builder,             // NOLINT
 
   auto &op = fwd_block->back();
 
-  std::vector<pir::Attribute> outs_stop_gradient;
+  auto outs_stop_gradient_attr = true;
   for (size_t i = 0; i < op.num_operands(); ++i) {
     argument.AddOutput(op.operand(i).type());
     auto bool_attr = op.operand_source(i).attribute<pir::BoolAttribute>(
         pir::kStopGradientAttrName);
-    outs_stop_gradient.push_back(bool_attr ? bool_attr
-                                           : builder.bool_attr(false));
+    if (!bool_attr || (bool_attr && !bool_attr.data())) {
+      outs_stop_gradient_attr = false;
+    }
   }
+  std::vector<pir::Attribute> outs_stop_gradient(
+      op.num_operands(), builder.bool_attr(outs_stop_gradient_attr));
 
   argument.AddAttribute(
       kBackwardFunctionIdAttrName,
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -1524,8 +1524,12 @@ void HandleForPyLayerOp(
   auto old_pylayerop = op_item->dyn_cast<PyLayerOp>();
   std::vector<pir::Type> new_pylayerop_outputs;
   for (size_t i = 0; i < old_pylayerop.num_results(); ++i) {
-    new_pylayerop_outputs.push_back(
-        ConvertOpTypeToKernelType(ctx, old_pylayerop.result(i).type(), place));
+    if (!static_cast<bool>(old_pylayerop.result(i).type())) {
+      new_pylayerop_outputs.push_back(old_pylayerop.result(i).type());
+    } else {
+      new_pylayerop_outputs.push_back(ConvertOpTypeToKernelType(
+          ctx, old_pylayerop.result(i).type(), place));
+    }
   }
 
   // Create PyLayerOp and insert to kernel dialect program
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -755,10 +755,11 @@ static std::vector<std::vector<pir::Value>> GenerateBackwardBlockForPyLayerOp(
   // 1. construct pylayer grad op
   VLOG(6) << "Prepare Outputs for pylayer_grad";
   std::vector<pir::Type> output_types;
+  // NOTE: the last input of pylayer op is create_stack when called
+  // save_for_backward, whose stop_gradient is always True
   for (size_t i = 0; i < inputs_.size(); ++i) {
-    if (!stop_gradients[i][0]) {
-      output_types.push_back(inputs_[i][0].type());
-    }
+    if (inputs_[i][0].type().isa<pir::InletType>()) break;
+    output_types.push_back(inputs_[i][0].type());
   }
 
   VLOG(6) << "Prepare Inputs for pylayer_grad";
@@ -837,12 +838,9 @@ static std::vector<std::vector<pir::Value>> GenerateBackwardBlockForPyLayerOp(
   VLOG(6) << "Update pylayer_grad op finished";
 
   std::vector<std::vector<pir::Value>> res{inputs_.size()};
-  int grad_op_result_index = 0;
   for (size_t i = 0; i < res.size(); ++i) {
     res[i].resize(1);
-    res[i][0] = !stop_gradients[i][0]
-                    ? pylayer_grad->result(grad_op_result_index++)
-                    : pir::Value();
+    res[i][0] = !stop_gradients[i][0] ? pylayer_grad->result(i) : pir::Value();
   }
   return res;
 }
diff --git a/paddle/fluid/pybind/python_callable_registry.cc b/paddle/fluid/pybind/python_callable_registry.cc
@@ -69,11 +69,17 @@ void PirCallPythonFunc(py::object *callable,
 
   for (size_t i = 0; i < out_num; ++i) {
     try {
-      auto py_out_value = py::cast<pir::Value>(ret_tuple[i]);
-      PADDLE_ENFORCE_NOT_NULL(py_out_value.impl(),
-                              common::errors::InvalidArgument(
-                                  "Output value %d should not be nullptr", i));
-      (*outs)[i] = py_out_value;
+      if (ret_tuple[i].is_none()) {
+        VLOG(6) << "Set Output( " << i << " ) value as fake_value";
+        (*outs)[i] = pir::Value(nullptr);
+      } else {
+        auto py_out_value = py::cast<pir::Value>(ret_tuple[i]);
+        PADDLE_ENFORCE_NOT_NULL(
+            py_out_value.impl(),
+            common::errors::InvalidArgument(
+                "Output value %d should not be nullptr", i));
+        (*outs)[i] = py_out_value;
+      }
     } catch (py::cast_error &) {
       PADDLE_THROW(common::errors::InvalidArgument(
           "pybind11::cast to pir::Value error. The %d-th output exception is "
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -873,12 +873,6 @@ def append_yield(
                         # create grad_op
                         before_ops_num = len(bwd_block.ops)
 
-                        # TODO(MarioLulab): `PyLayer.backward` has not supported return `None` yet. Will be supported soon.
-                        if any(zero_flag):
-                            raise ValueError(
-                                "pylayer_op.backward have not supported return `None` yet. Will be supported soon."
-                            )
-
                         with dynamic_shape_prim_vjp_guard(op, inputs):
                             input_grads = paddle.framework.core.call_vjp(
                                 op,
diff --git a/test/dygraph_to_static/test_pylayer.py b/test/dygraph_to_static/test_pylayer.py
@@ -108,7 +108,6 @@ def backward(ctx, dy):
 class cus_tanh_3(PyLayer):
     @staticmethod
     def forward(ctx, x1, x2, func1, func2=paddle.square):
-        ctx.func = func2
         y1 = func1(x1)
         y2 = func1(x2)
         ctx.save_for_backward(y1, y2)
@@ -117,7 +116,7 @@ def forward(ctx, x1, x2, func1, func2=paddle.square):
     @staticmethod
     def backward(ctx, dy1, dy2):
         y1, y2 = ctx.saved_tensor()
-        re1 = dy1 * (1 - ctx.func(y1))
+        re1 = dy1 * (1 - paddle.square(y1))
         re2 = dy2 * (1 - paddle.square(y2))
         return re1, None
 
@@ -510,15 +509,8 @@ def test_func(input1, input2):
 
         self.run_in_pir = False
         self._run_and_compare(input1, input2)
-
-        # TODO(MarioLulab): pylayer_op.backward have not supported return `None` yet. Will be supported soon.
-        with self.assertRaises(Exception) as e:
-            self.run_in_pir = True
-            self._run_and_compare(input1, input2)
-        self.assertTrue(
-            "pylayer_op.backward have not supported return `None` yet. Will be supported soon."
-            in str(e.exception)
-        )
+        self.run_in_pir = True
+        self._run_and_compare(input1, input2)
 
     def test_simple_pylayer_return_none(self):
         @paddle.jit.to_static(full_graph=True)
