PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
Copy file name to clipboard
+29-20Lines changed: 29 additions & 20 deletions b/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
Copy file name to clipboard
+29-20Lines changed: 29 additions & 20 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/control_flow_api.cc
Copy file name to clipboard
+31-2Lines changed: 31 additions & 2 deletions b/‎paddle/fluid/pybind/control_flow_api.cc
Copy file name to clipboard
+31-2Lines changed: 31 additions & 2 deletions
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+13-1Lines changed: 13 additions & 1 deletion b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+13-1Lines changed: 13 additions & 1 deletion
@@ -48,6 +48,7 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
   argument.output_types.swap(output_types);
   argument.AddRegion().emplace_back();
   argument.AddRegion().emplace_back();
+  cond.set_attribute(kStopGradientAttrName, builder.bool_attr(true));
 }
 
 void IfOp::Build(pir::Builder &builder,             // NOLINT
@@ -289,17 +290,30 @@ void WhileOp::Build(pir::Builder &builder,             // NOLINT
   argument.AddInput(cond);
   argument.AddInputs(inputs);
   auto &body = argument.AddRegion().emplace_back();
+  std::vector<pir::Attribute> outs_stop_gradient;
   for (auto val : inputs) {
     argument.AddOutput(val.type());
-    body.AddArgument(val.type());
+    auto arg = body.AddArgument(val.type());
+
+    auto bool_attr = val.attribute<pir::BoolAttribute>(kStopGradientAttrName);
+    arg.set_attribute(kStopGradientAttrName,
+                      bool_attr ? bool_attr : builder.bool_attr(false));
+    outs_stop_gradient.push_back(bool_attr ? bool_attr
+                                           : builder.bool_attr(false));
   }
+
+  argument.AddAttribute(
+      kStopGradientAttrName,
+      pir::ArrayAttribute::get(builder.ir_context(), outs_stop_gradient));
+
   cond.set_attribute(kStopGradientAttrName, builder.bool_attr(true));
 }
 pir::Block &WhileOp::body() {
   pir::Region &body_region = (*this)->region(0);
   if (body_region.empty()) body_region.emplace_back();
   return body_region.front();
 }
+
 pir::Value WhileOp::cond() { return (*this)->operand_source(0); }
 
 void WhileOp::Print(pir::IrPrinter &printer) {
@@ -367,6 +381,14 @@ std::vector<std::vector<pir::OpResult>> WhileOp::Vjp(
                         "the outputs size is %d.",
                         inputs.size(),
                         outputs.size()));
+  PADDLE_ENFORCE_EQ(inputs.size(),
+                    out_grads.size() + 1,
+                    phi::errors::InvalidArgument(
+                        "while op's inputs' size should equal to "
+                        "output_grads' size + 1, Now the inputs's size is %d ."
+                        "the output_grads size is %d.",
+                        inputs.size(),
+                        out_grads.size()));
   PADDLE_ENFORCE_EQ(stop_gradients[0][0],
                     true,
                     phi::errors::InvalidArgument(
@@ -377,27 +399,12 @@ std::vector<std::vector<pir::OpResult>> WhileOp::Vjp(
 
   std::vector<pir::Type> output_types;
   std::vector<pir::Value> loop_vars;
-  size_t index = 0;
 
-  for (; index < outputs.size(); ++index) {
+  for (size_t index = 0; index < out_grads.size(); ++index) {
     if (!stop_gradients[index + 1][0]) {
       loop_vars.push_back(out_grads[index][0]);
     }
   }
-  for (++index; index < inputs.size(); ++index) {
-    if (!stop_gradients[index][0]) {
-      auto fwd_type = inputs[index][0].type().dyn_cast<DenseTensorType>();
-      PADDLE_ENFORCE_NE(
-          fwd_type,
-          pir::Type(),
-          phi::errors::InvalidArgument(
-              "The forward value type must be dense tensor type."));
-      auto shape = vectorize(fwd_type.dims());
-      auto dtype = TransToPhiDataType(fwd_type.dtype());
-      auto full_op = builder.Build<FullOp>(shape, 0.0, dtype, phi::CPUPlace());
-      loop_vars.push_back(full_op.out());
-    }
-  }
   auto while_grad = builder.Build<WhileOp>(cond_val, loop_vars);
 
   std::vector<std::vector<pir::OpResult>> res(inputs.size());
@@ -426,9 +433,7 @@ std::vector<std::vector<pir::OpResult>> TuplePushOpVjpInterfaceModel::Vjp(
   res[0].resize(1);
   for (size_t i = 1u; i < inputs.size(); ++i) {
     res[i].resize(1);
-    if (!stop_gradients[i][0]) {
-      res[i][0] = pop_op.result(i - 1);
-    }
+    res[i][0] = pop_op.result(i - 1);
   }
   return res;
 }
@@ -439,6 +444,10 @@ void HasElementsOp::Build(pir::Builder &builder,             // NOLINT
   argument.AddInput(container);
   argument.AddOutput(
       DenseTensorType::get(builder.ir_context(), builder.bool_type(), {1}));
+  std::vector<pir::Attribute> outs_stop_gradient{builder.bool_attr(true)};
+  argument.AddAttribute(
+      kStopGradientAttrName,
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
 }
 void HasElementsOp::VerifySig() {
   VLOG(4) << "Verifying inputs, outputs ,attributes for: HasElementsOp.";
 
@@ -16,6 +16,7 @@
 #include <vector>
 
 #include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
+#include "paddle/pir/core/block.h"
 #include "paddle/pir/core/op_base.h"
 
 namespace paddle {
@@ -78,6 +79,7 @@ class WhileOp : public pir::Op<WhileOp, VjpInterface> {
                     const std::vector<pir::Value> &inputs);
   pir::Block &body();
   pir::Value cond();
+  const pir::Block::ArgListType &block_args() { return body().args(); }
   void Print(pir::IrPrinter &printer);  // NOLINT
   void VerifySig() {}
   void VerifyRegion() {}
 
@@ -35,6 +35,7 @@
 
 namespace py = pybind11;
 using paddle::dialect::ApiBuilder;
+using paddle::dialect::HasElementsOp;
 using paddle::dialect::IfOp;
 using paddle::dialect::WhileOp;
 using pir::Block;
@@ -88,7 +89,10 @@ void BindWhileOp(py::module* m) {
     WhileOp in python api.
   )DOC");
   while_op.def("body", &WhileOp::body, return_value_policy::reference)
-      .def("as_operation", &WhileOp::operation, return_value_policy::reference);
+      .def("as_operation", &WhileOp::operation, return_value_policy::reference)
+      .def("block_arguments",
+           &WhileOp::block_args,
+           return_value_policy::reference);
 }
 
 void GetUsedExternalValueImpl(
@@ -126,6 +130,30 @@ std::vector<Value> GetUsedExternalValue(const Operation& op) {
   return used_values;
 }
 
+Value BuildHasElementsOp(Operation& fwd_op) {  // NOLINT
+  PADDLE_ENFORCE(fwd_op.isa<WhileOp>(),
+                 phi::errors::PreconditionNotMet(
+                     "param op of BuildHasElementsOp must be while op."));
+  auto fwdop = fwd_op.dyn_cast<WhileOp>();
+  TuplePushOp push_op;
+  for (auto iter = fwdop.body().rbegin(); iter != fwdop.body().rend(); ++iter) {
+    if (iter->isa<TuplePushOp>()) {
+      push_op = iter->dyn_cast<TuplePushOp>();
+      PADDLE_ENFORCE_EQ(push_op.container().use_empty(),
+                        false,
+                        phi::errors::InvalidArgument(
+                            "The last container in foward while op must used "
+                            "after construct while_grad op"));
+      break;
+    }
+  }
+  auto new_cond = ApiBuilder::Instance()
+                      .GetBuilder()
+                      ->Build<HasElementsOp>(push_op.container())
+                      .out();
+  return new_cond;
+}
+
 void BuildPipeForBlock(Block* block) {
   PADDLE_ENFORCE_NOT_NULL(
       block,
@@ -193,16 +221,17 @@ void PyIfOp::UpdateOutput() {
 void BindControlFlowApi(py::module* m) {
   m->def("get_used_external_value", GetUsedExternalValue);
   m->def("build_pipe_for_block", BuildPipeForBlock);
+  m->def("cf_has_elements", BuildHasElementsOp);
   m->def("cf_yield", [](py::list inputs) {
     std::vector<Value> input_values;
     for (auto input : inputs) {
       input_values.push_back(input.cast<Value>());
     }
     ApiBuilder::Instance().GetBuilder()->Build<YieldOp>(input_values);
   });
-
   BindIfOp(m);
   BindWhileOp(m);
 }
+
 }  // namespace pybind
 }  // namespace paddle
@@ -320,12 +320,16 @@ void BindBlock(py::module *m) {
           "front",
           [](Block &self) { return &self.front(); },
           return_value_policy::reference)
+      .def_property_readonly(
+          "parent_op",
+          [](Block &self) { return self.GetParentOp(); },
+          return_value_policy::reference)
       .def_property_readonly(
           "program",
           [](Block &self) { return self.GetParentOp()->GetParentProgram(); },
           return_value_policy::reference)
       .def_property_readonly(
-          "get_parent",
+          "parent_block",
           [](Block &self) { return self.GetParentOp()->GetParent(); },
           return_value_policy::reference)
       .def_property_readonly("ops",
@@ -717,6 +721,14 @@ void BindValue(py::module *m) {
                 kAttrIsPersisable,
                 BoolAttribute::get(pir::IrContext::Instance(), persistable));
           })
+      .def("all_used_ops",
+           [](Value &self) -> py::list {
+             py::list op_list;
+             for (auto it = self.use_begin(); it != self.use_end(); ++it) {
+               op_list.append(it.owner());
+             }
+             return op_list;
+           })
       .def(
           "get_defining_op",
           [](Value self) -> pir::Operation * {