[PIR] fix bug in dy2st split program. (#62030)

winter-wang · web-flow · commit 2a472d9ad1d6 · 2024-02-26T16:41:29.000+08:00
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
@@ -556,11 +556,13 @@ void HandleForSpecialOp(pir::Operation* op,
     auto value = op->operand_source(0);
 
     Scope* scope = const_cast<Scope*>(value_exe_info->GetScope());
-    if (value.defining_op()->HasAttribute(kAttrIsPersistable) &&
-        value.attribute<pir::BoolAttribute>(kAttrIsPersistable).data()) {
-      VLOG(6) << "Handle for builtin.shadow_output persistable value:"
-              << var_name;
-      scope = const_cast<Scope*>(value_exe_info->GetScope()->root());
+    if (auto bool_atttr =
+            value.attribute<pir::BoolAttribute>(kAttrIsPersistable)) {
+      if (bool_atttr.data()) {
+        VLOG(6) << "Handle for builtin.shadow_ouptut persistable value:"
+                << var_name;
+        scope = const_cast<Scope*>(value_exe_info->GetScope()->root());
+      }
     }
 
     // change operand name to param_name
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -74,6 +74,7 @@
 #include "paddle/pir/include/core/program.h"
 #include "paddle/pir/include/core/type.h"
 #include "paddle/pir/include/core/value.h"
+#include "paddle/pir/include/core/visitors.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_dialect.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_attribute.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
@@ -749,6 +750,12 @@ void BindValue(py::module *m) {
             } else if (auto data_op =
                            self.defining_op<paddle::dialect::DataOp>()) {
               return data_op.attribute<pir::StrAttribute>("name").AsString();
+            } else if (auto block_arg = self.dyn_cast<BlockArgument>()) {
+              if (block_arg.is_kwarg()) {
+                return block_arg.keyword();
+              } else {
+                return "arg_" + std::to_string(block_arg.index());
+              }
             } else {
               PADDLE_THROW(phi::errors::InvalidArgument(
                   "Currently, we can only get name of Value that "
@@ -952,9 +959,11 @@ AnalysisMiddleVariable(const Program &program,
                                             forward_inputs.end());
   range_block_do(
       program.block(), backward_range, [&backward_inputs](Operation *op) {
-        for (auto &t : op->operands()) {
-          backward_inputs.insert(t.source());
-        }
+        pir::Walk(op, [&](Operation *inner_op) {
+          for (auto &t : inner_op->operands()) {
+            backward_inputs.insert(t.source());
+          }
+        });
       });
 
   range_block_do(
@@ -1215,28 +1224,34 @@ SplitedResult SplitForwardBackward(
 
   // counter = 0;
   if (has_backward) {
-    VLOG(4) << "start create backward inputs, inserting pd.data ops.";
-    VLOG(4) << "Create pd.data for backward program: fo, start with input_"
-            << counter;
+    VLOG(4) << "start create backward inputs, creating keyword argument.";
+    VLOG(4)
+        << "Create keyword argument for backward program: fo, start with input_"
+        << counter;
     std::for_each(
         forward_outputs.begin(), forward_outputs.end(), create_kwarg_fn);
-    VLOG(4) << "Create pd.data for backward program: fx, start with input_"
-            << counter;
+    VLOG(4)
+        << "Create keyword argument for backward program: fx, start with input_"
+        << counter;
     std::for_each(
         forward_inputs.begin(), forward_inputs.end(), create_kwarg_fn);
-    VLOG(4) << "Create pd.data for backward program: fp, start with input_"
-            << counter;
+    VLOG(4)
+        << "Create keyword argument for backward program: fp, start with input_"
+        << counter;
     std::for_each(
         forward_params.begin(), forward_params.end(), create_kwarg_fn);
-    VLOG(4) << "Create pd.data for backward program: fm, start with input_"
-            << counter;
+    VLOG(4)
+        << "Create keyword argument for backward program: fm, start with input_"
+        << counter;
     std::for_each(middle_values.begin(), middle_values.end(), create_kwarg_fn);
-    VLOG(4) << "Create pd.data for backward program: fo_g, start with input_"
+    VLOG(4) << "Create keyword argument for backward program: fo_g, start with "
+               "input_"
             << counter;
     std::for_each(forward_outputs_grads.begin(),
                   forward_outputs_grads.end(),
                   create_kwarg_fn);
-    VLOG(4) << "Create pd.data for backward program end. input_" << counter;
+    VLOG(4) << "Create keyword argument for backward program end. input_"
+            << counter;
   }
 
   // counter = 0;
diff --git a/paddle/pir/include/dialect/control_flow/ir/cf_op.h b/paddle/pir/include/dialect/control_flow/ir/cf_op.h
@@ -57,6 +57,7 @@ class IR_API TuplePushOp : public Op<TuplePushOp, SideEffectTrait> {
                     Value inlet,
                     std::initializer_list<Value> element_list);
   void VerifySig();
+  void VerifyRegion();
 
   Value container() { return container_interface().container(); }
   Value inlet() { return operand_source(0); }
@@ -84,6 +85,7 @@ class IR_API TuplePopOp : public Op<TuplePopOp, SideEffectTrait> {
                     OperationArgument &argument,  // NOLINT
                     Value outlet);
   void VerifySig();
+  void VerifyRegion();
 
   Value container() { return container_interface().container(); }
   Value inlet() { return container_interface().inlet(); }
diff --git a/paddle/pir/src/core/ir_printer.cc b/paddle/pir/src/core/ir_printer.cc
@@ -216,6 +216,16 @@ void IrPrinter::PrintRegion(const Region& region) {
 void IrPrinter::PrintBlock(const Block& block) {
   os << indentation() << "{\n";
   AddIndentation();
+  if (!block.kwargs_empty()) {
+    os << indentation() << "^kw:";
+    auto cur = block.kwargs_begin(), end = block.kwargs_end();
+    PrintValue(cur->second);
+    while (++cur != end) {
+      os << ", ";
+      PrintValue(cur->second);
+    }
+    os << "\n";
+  }
   for (auto& item : block) {
     PrintOperation(&item);
     os << "\n";
@@ -241,10 +251,11 @@ void IrPrinter::PrintValue(Value v) {
     aliases_[key] = new_name;
     os << new_name;
   } else {
-    std::string new_name = "%arg" + std::to_string(cur_block_argument_number_);
-    cur_block_argument_number_++;
-    aliases_[key] = new_name;
-    os << new_name;
+    auto arg = v.dyn_cast<BlockArgument>();
+    os << (aliases_[key] =
+               arg.is_kwarg()
+                   ? "%kwarg_" + arg.keyword()
+                   : "%arg_" + std::to_string(cur_block_argument_number_++));
   }
 }
 
diff --git a/paddle/pir/src/core/value_impl.cc b/paddle/pir/src/core/value_impl.cc
@@ -27,9 +27,9 @@ void ValueImpl::set_first_use(OpOperandImpl *first_use) {
   uint32_t offset = kind();
   first_use_offseted_by_kind_ = reinterpret_cast<OpOperandImpl *>(
       reinterpret_cast<uintptr_t>(first_use) + offset);
-  VLOG(10) << "The index of this value is " << offset
-           << ". Offset and set first use: " << first_use << " -> "
-           << first_use_offseted_by_kind_ << ".";
+  VLOG(10) << "The index of this value is: " << offset
+           << ". The address of this value is: " << this
+           << ". This value first use is: " << first_use << ".";
 }
 
 std::string ValueImpl::PrintUdChain() {
@@ -56,8 +56,7 @@ ValueImpl::ValueImpl(Type type, uint32_t kind) : id_(GenerateId()) {
   first_use_offseted_by_kind_ = reinterpret_cast<OpOperandImpl *>(
       reinterpret_cast<uintptr_t>(nullptr) + kind);
   VLOG(10) << "Construct a ValueImpl whose's kind is " << kind
-           << ". The offset first_use address is: "
-           << first_use_offseted_by_kind_;
+           << ". The value_impl address is: " << this;
 }
 
 }  // namespace detail
diff --git a/paddle/pir/src/dialect/control_flow/ir/cf_op.cc b/paddle/pir/src/dialect/control_flow/ir/cf_op.cc
@@ -54,8 +54,6 @@ void TuplePushOp::VerifySig() {
   IR_ENFORCE(num_operands() >= 1u, "The size of inputs must no less than 1.");
   IR_ENFORCE(operand_source(0).type().isa<InletType>(),
              "The first input of cf.tuple_push must be inlet_type.");
-  IR_ENFORCE(operand_source(0).HasOneUse(),
-             "The inlet value of cf.tuple_push can only be used once.");
 
   // No attributes should be verify.
 
@@ -64,6 +62,13 @@ void TuplePushOp::VerifySig() {
   VLOG(4) << "End Verifying for TuplePushOp.";
 }
 
+void TuplePushOp::VerifyRegion() {
+  // Note(winter-wang):Constraints on the number of uses can only can be placed
+  // in VerifyRegion, Otherwise cloning would fail.
+  IR_ENFORCE(operand_source(0).HasOneUse(),
+             "The inlet value of cf.tuple_push can only be used once.");
+}
+
 size_t TuplePushOp::tuple_size() {
   auto operands_size = num_operands();
   IR_ENFORCE(operands_size >= 1u,
@@ -96,12 +101,15 @@ void TuplePopOp::VerifySig() {
   IR_ENFORCE(num_operands() == 1u, "The size of inputs must equal to 1.");
   IR_ENFORCE(operand_source(0).type().isa<OutletType>(),
              "The first input of cf.tuple_pop must be outlet_type.");
-  IR_ENFORCE(operand_source(0).HasOneUse(),
-             "The outlet value of cf.tuple_pop can only be used once.");
 
   // No attributes should be verify.
 
   // Verify outputs:
+}
+
+void TuplePopOp::VerifyRegion() {
+  IR_ENFORCE(operand_source(0).HasOneUse(),
+             "The outlet value of cf.tuple_pop can only be used once.");
 
   // Verify stack validity:
   auto pop_op = container_interface().tuple_pop_op();
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -746,6 +746,9 @@ def _append_backward_desc(self, train_runnable_program: RunableProgram):
         params = train_runnable_program.param_values
         combined_inputs = list(itertools.chain(inputs, params))
         forward_end_idx = len(program.global_block().ops)
+        forward_end_op = None
+        if forward_end_idx > 0:
+            forward_end_op = program.global_block().ops[-1]
         grad_info_map = [None] * len(combined_inputs)
         with backend_guard(self._backend):
             check_type(
@@ -798,6 +801,11 @@ def _append_backward_desc(self, train_runnable_program: RunableProgram):
                             )
                         ),
                     )
+                    if forward_end_op is not None:
+                        for idx, op in enumerate(program.global_block().ops):
+                            if op == forward_end_op:
+                                forward_end_idx = idx + 1
+                                break
 
             if self._hooker:
                 (
