[PIR] pir pylayer update input (#70335)

wanghuancoder · web-flow · commit 4a6d084297e4 · 2024-12-20T09:49:23.000+08:00
* pir pylayer updata input
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.cc
@@ -20,6 +20,7 @@ paddle::dialect::PyLayerOp
 
 #include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
 
+#include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
@@ -39,6 +40,40 @@ paddle::dialect::PyLayerOp
 namespace paddle {
 namespace dialect {
 
+std::unordered_set<pir::Value> GetInternalInputs(pir::Block *block) {
+  std::unordered_set<pir::Value> inner_inputs;
+  for (auto &op : *block) {
+    std::string op_name = op.name();
+    if (op.attributes().count("op_name")) {
+      op_name = op.attributes()
+                    .at("op_name")
+                    .dyn_cast<pir::StrAttribute>()
+                    .AsString();
+    }
+    VLOG(8) << "GetInternalInputs of " << op_name;
+    if (op.num_regions()) {
+      for (size_t i = 0; i < op.num_regions(); ++i) {
+        for (auto &sub_block : op.region(i)) {
+          std::unordered_set<pir::Value> sub_set =
+              GetInternalInputs(&sub_block);
+          inner_inputs.insert(sub_set.begin(), sub_set.end());
+        }
+      }
+    }
+    if (op.isa<pir::TuplePopOp>()) {
+      auto tuple_pop_op = op.dyn_cast<pir::TuplePopOp>();
+      if (tuple_pop_op.has_container()) {
+        inner_inputs.insert(tuple_pop_op.container());
+      }
+    }
+    for (size_t i = 0; i < op.num_operands(); ++i) {
+      inner_inputs.insert(op.operand_source(i));
+      VLOG(10) << op_name << "'s inner_input: " << op.operand_source(i).impl();
+    }
+  }
+  return inner_inputs;
+}
+
 const char *PyLayerOp::attributes_name[1] = {
     kBackwardFunctionIdAttrName};  // NOLINT
 
@@ -196,6 +231,59 @@ void PyLayerOp::UpdateOutput() {
   VerifyRegion();
 }
 
+PyLayerOp PyLayerOp::UpdateInput() {
+  PADDLE_ENFORCE_NOT_NULL(*this,
+                          common::errors::InvalidArgument(
+                              "The pylayer_op in PyLayerOp used to update "
+                              "output can't be nullptr"));
+  auto program_block = parent();
+  PADDLE_ENFORCE_NOT_NULL(
+      program_block,
+      common::errors::InvalidArgument(
+          "The parent block of pylayer_op which used to update "
+          "output can't be nullptr"));
+
+  pir::Block &block = forward_block();
+  std::vector<pir::Value> input_values = inputs();
+
+  std::unordered_set<pir::Value> inner_inputs;
+  inner_inputs = GetInternalInputs(&block);
+
+  for (size_t arg_id = 0; arg_id < block.args_size();) {
+    if (block.arg(arg_id) && (!inner_inputs.count(block.arg(arg_id)))) {
+      block.EraseArg(arg_id);
+      continue;
+    }
+    ++arg_id;
+  }
+
+  bool need_build_new_pylayer = false;
+  std::vector<pir::Value> new_pylayer_inputs;
+
+  for (auto value : input_values) {
+    if (value && (!inner_inputs.count(value))) {
+      need_build_new_pylayer = true;
+      continue;
+    }
+    new_pylayer_inputs.push_back(value);
+  }
+
+  if (need_build_new_pylayer) {
+    ::pir::IrContext *ctx = ::pir::IrContext::Instance();
+
+    ::pir::Builder builder = ::pir::Builder(ctx, program_block);
+    builder.set_insertion_point(&(**this));
+    auto new_pylayer = builder.Build<PyLayerOp>(new_pylayer_inputs,
+                                                forward_region().TakeBack(),
+                                                backward_function_id());
+    (**this).ReplaceAllUsesWith(new_pylayer.outputs());
+    pir::Block::Iterator iter = **this;
+    iter = program_block->erase(iter);
+    return new_pylayer;
+  }
+  return *this;
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h
@@ -57,6 +57,15 @@ class PyLayerOp : public pir::Op<PyLayerOp> {
                                         "num_operands of pylayer op."));
     return operand_source(index);
   }
+
+  std::vector<pir::Value> outputs() const {
+    std::vector<pir::Value> output_values;
+    for (size_t index = 0; index < num_results(); ++index) {
+      output_values.push_back(result(index));
+    }
+    return output_values;
+  }
+
   pir::Block &forward_block();
   pir::Region &forward_region() { return (*this)->region(0); }
 
@@ -75,6 +84,7 @@ class PyLayerOp : public pir::Op<PyLayerOp> {
   void VerifyRegion();
 
   void UpdateOutput();
+  PyLayerOp UpdateInput();
 };
 
 }  // namespace dialect
diff --git a/paddle/fluid/pybind/control_flow_api.cc b/paddle/fluid/pybind/control_flow_api.cc
@@ -103,6 +103,7 @@ void BindPyLayerOp(py::module* m) {
            &PyLayerOp::forward_block,
            return_value_policy::reference)
       .def("update_output", &PyLayerOp::UpdateOutput)
+      .def("update_input", &PyLayerOp::UpdateInput)
       .def(
           "as_operation", &PyLayerOp::operation, return_value_policy::reference)
       .def("id",
