fix

paddle/fluid/framework/new_executor/instruction/instruction_util.h

@@ -68,5 +68,8 @@ void HandleForInplaceOp(pir::Operation* op,
                         InstructionBase* instr);
 
 void ShareVarBuffer(const Variable* src_var, Variable* dst_var);
+
+std::unordered_set<pir::Value> GetInternalInputs(pir::Block* block);
+std::unordered_set<pir::Value> GetInternalOutputs(pir::Block* block);
 }  // namespace framework
 }  // namespace paddle

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"
@@ -196,6 +197,119 @@ void PyLayerOp::UpdateOutput() {
   VerifyRegion();
 }
 
+void PyLayerOp::UpdateInputOutput() {
+  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"));
+
+  std::unordered_set<pir::Value> global_block_inner_inputs;
+  global_block_inner_inputs =
+      paddle::framework::GetInternalInputs(program_block);
+
+  pir::Block &block = forward_block();
+  std::vector<pir::Value> input_values = inputs();
+  std::vector<pir::Value> output_values = outputs();
+
+  std::unordered_set<pir::Value> inner_inputs;
+  inner_inputs = paddle::framework::GetInternalInputs(&block);
+  std::unordered_set<pir::Value> inner_outputs;
+  inner_outputs = paddle::framework::GetInternalOutputs(&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::Type> new_pylayer_output_types;
+  std::vector<pir::Value> new_pylayer_inputs;
+  std::vector<pir::Value> new_pylayer_yield_inputs;
+
+  for (auto value : input_values) {
+    if (value && (!inner_inputs.count(value))) {
+      need_build_new_pylayer = true;
+      continue;
+    }
+    new_pylayer_inputs.push_back(value);
+  }
+
+  std::vector<int> old_pylayer_outputs_map_to_new_pylayer_outputs_index;
+
+  if (block.back().isa<pir::YieldOp>()) {
+    std::vector<pir::Value> yield_inputs = block.back().operands_source();
+    PADDLE_ENFORCE_EQ(
+        yield_inputs.size(),
+        output_values.size(),
+        common::errors::Unimplemented(
+            "YieldOp's input size(%d) must be equal with "
+            "PyLayer's outpus's output size %d. If Pass modify PyLayer's "
+            "block, the Pass should not modify YieldOp, because YieldOp must "
+            "update with PyLayer outputs together. Otherwise, when updating "
+            "PyLayer outputs, the mapping relationship between the new PyLayer "
+            "and the old PyLayer outputs cannot be known. Therefore, we can't "
+            "use ReplaceAllUsesWith update Value of PyLayer outputs.",
+            yield_inputs.size(),
+            output_values.size()));
+    int index = 0;
+    for (size_t i = 0; i < yield_inputs.size(); i++) {
+      if (yield_inputs[i] && (!inner_outputs.count(yield_inputs[i]))) {
+        PADDLE_ENFORCE_EQ(
+            global_block_inner_inputs.count(output_values[i]),
+            0,
+            common::errors::Unimplemented(
+                "The PyLayer's output not defined in PyLayer's block, "
+                "but used in global block."));
+        need_build_new_pylayer = true;
+        old_pylayer_outputs_map_to_new_pylayer_outputs_index.push_back(-1);
+        continue;
+      }
+      new_pylayer_output_types.push_back(yield_inputs[i].type());
+      new_pylayer_yield_inputs.push_back(yield_inputs[i]);
+      old_pylayer_outputs_map_to_new_pylayer_outputs_index.push_back(index++);
+    }
+  } else {
+    if (!output_values.empty()) {
+      PADDLE_THROW(common::errors::Unimplemented(
+          "The last op of PyLayer block, is not yield_op, but a %s",
+          block.back().name()));
+    }
+  }
+
+  if (need_build_new_pylayer) {
+    ::pir::IrContext *ctx = ::pir::IrContext::Instance();
+    block.pop_back();
+    ::pir::Builder builder = ::pir::Builder(ctx, &block);
+    builder.SetInsertionPointToBlockEnd(&block);
+    builder.Build<pir::YieldOp>(new_pylayer_yield_inputs);
+
+    ::pir::Builder builder2 = ::pir::Builder(ctx, program_block);
+    builder2.set_insertion_point(&(**this));
+    auto new_pylayer = builder2.Build<PyLayerOp>(new_pylayer_inputs,
+                                                 forward_region().TakeBack(),
+                                                 backward_function_id());
+    for (size_t i = 0;
+         i < old_pylayer_outputs_map_to_new_pylayer_outputs_index.size();
+         i++) {
+      if (old_pylayer_outputs_map_to_new_pylayer_outputs_index[i] != -1) {
+        output_values[i].ReplaceAllUsesWith(new_pylayer.result(
+            old_pylayer_outputs_map_to_new_pylayer_outputs_index[i]));
+      }
+    }
+    pir::Block::Iterator iter = **this;
+    program_block->erase(iter);
+  }
+}
+
 }  // namespace dialect
 }  // namespace paddle
 

paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h

@@ -49,6 +49,14 @@ class PyLayerOp : public pir::Op<PyLayerOp> {
     }
     return input_values;
   }
+
+  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::Value input(size_t index) const {
     PADDLE_ENFORCE_LT(
         index,
@@ -75,6 +83,7 @@ class PyLayerOp : public pir::Op<PyLayerOp> {
   void VerifyRegion();
 
   void UpdateOutput();
+  void UpdateInputOutput();
 };
 
 }  // namespace dialect

paddle/fluid/pybind/control_flow_api.cc

@@ -38,6 +38,8 @@
 
 #include "paddle/fluid/pybind/python_callable_registry.h"
 
+#include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
+
 namespace py = pybind11;
 using paddle::dialect::ApiBuilder;
 using paddle::dialect::AssertOp;
@@ -95,6 +97,20 @@ void BindPyLayerOp(py::module* m) {
     return ApiBuilder::Instance().GetBuilder()->Build<PyLayerOp>(
         inputs, std::vector<Type>{}, -1);
   });
+  m->def("updata_pylayer_op", [](pir::Program* program) -> void {
+    std::unordered_set<pir::Value> global_block_inner_inputs;
+    global_block_inner_inputs =
+        paddle::framework::GetInternalInputs(program->block());
+
+    for (auto iter = program->block()->begin();
+         iter != program->block()->end();) {
+      pir::Operation* op_item = &(*iter);
+      ++iter;
+      if (op_item->isa<PyLayerOp>()) {
+        op_item->dyn_cast<PyLayerOp>().UpdateInputOutput();
+      }
+    }
+  });
   py::class_<PyLayerOp> pylayer_op(*m, "PyLayerOp", R"DOC(
     TODO(MarioLulab): Add some docs for pd_op.pylayer
   )DOC");
@@ -103,6 +119,7 @@ void BindPyLayerOp(py::module* m) {
            &PyLayerOp::forward_block,
            return_value_policy::reference)
       .def("update_output", &PyLayerOp::UpdateOutput)
+      .def("update_input_and_output", &PyLayerOp::UpdateInputOutput)
       .def(
           "as_operation", &PyLayerOp::operation, return_value_policy::reference)
       .def("id",