PaddlePaddle
diff --git a/‎paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.cc
Copy file name to clipboard
+4-23Lines changed: 4 additions & 23 deletions b/‎paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.cc
Copy file name to clipboard
+4-23Lines changed: 4 additions & 23 deletions
diff --git a/‎paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.h
Copy file name to clipboard
-3Lines changed: 0 additions & 3 deletions b/‎paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.h
Copy file name to clipboard
-3Lines changed: 0 additions & 3 deletions
diff --git a/‎paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.cc
Copy file name to clipboard
+173Lines changed: 173 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.cc
Copy file name to clipboard
+173Lines changed: 173 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.h
Copy file name to clipboard
+64Lines changed: 64 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.h
Copy file name to clipboard
+64Lines changed: 64 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/instruction/instruction_util.cc
Copy file name to clipboard
+22Lines changed: 22 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/instruction/instruction_util.cc
Copy file name to clipboard
+22Lines changed: 22 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/instruction/instruction_util.h
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/instruction/instruction_util.h
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_interpreter.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/pir_interpreter.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions
@@ -198,27 +198,6 @@ IfInstruction::~IfInstruction() {
   }
 }
 
-void IfInstruction::CopyBranchOutput(const std::vector<std::string>& var_names,
-                                     const PirInterpreter* inter) {
-  for (size_t i = 0; i < var_names.size(); ++i) {
-    auto* inner_var = inter->InnerScope()->GetVar(var_names[i]);
-
-    if (inner_var->IsType<phi::DenseTensor>()) {
-      output_vars_[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
-          inner_var->Get<phi::DenseTensor>());
-
-    } else if (inner_var->IsType<phi::TensorArray>()) {
-      const auto& inner_array = inner_var->Get<phi::TensorArray>();
-      auto* output_array = output_vars_[i]->GetMutable<phi::TensorArray>();
-      // output_array->clear();
-      *output_array = inner_array;
-    } else {
-      PADDLE_THROW(
-          phi::errors::Unimplemented("unsupported type %d", inner_var->Type()));
-    }
-  }
-}
-
 void IfInstruction::Run() {
   bool cond = true;
   if (cond_var_->IsType<phi::DenseTensor>()) {
@@ -257,7 +236,8 @@ void IfInstruction::Run() {
     paddle::platform::DontClearMKLDNNCache(true_branch_inter_->GetPlace());
 #endif
     true_branch_inter_->Run({}, false);
-    CopyBranchOutput(true_branch_outputs_, true_branch_inter_);
+    CopyBranchOutput(
+        true_branch_outputs_, output_vars_, true_branch_inter_->InnerScope());
   } else {
 #ifdef PADDLE_WITH_DNNL
     // Executor on being destroyed clears oneDNN cache and resets
@@ -266,7 +246,8 @@ void IfInstruction::Run() {
     paddle::platform::DontClearMKLDNNCache(false_branch_inter_->GetPlace());
 #endif
     false_branch_inter_->Run({}, false);
-    CopyBranchOutput(false_branch_outputs_, false_branch_inter_);
+    CopyBranchOutput(
+        false_branch_outputs_, output_vars_, false_branch_inter_->InnerScope());
   }
   // copy output
 }
 
@@ -49,9 +49,6 @@ class IfInstruction : public InstructionBase {
   PirInterpreter* FalseBranchInterpreter() const { return false_branch_inter_; }
 
  private:
-  void CopyBranchOutput(const std::vector<std::string>& var_names,
-                        const PirInterpreter* inter);
-
   ::pir::Operation* op_;
 
   std::string cond_name_{"if_instruction"};
 
@@ -0,0 +1,173 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.h"
+
+#include "paddle/fluid/framework/new_executor/interpreter/interpreter_util.h"
+#include "paddle/fluid/framework/new_executor/interpreter/stream_analyzer.h"
+#include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
+#include "paddle/fluid/framework/new_executor/pir_interpreter.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infermeta.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_parser.h"
+#include "paddle/fluid/platform/collective_helper.h"
+#include "paddle/fluid/platform/device_context.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/core/meta_tensor.h"
+#include "paddle/phi/core/type_defs.h"
+
+#include "paddle/pir/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/operation.h"
+#include "paddle/pir/include/core/value.h"
+
+#include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
+#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
+
+#ifdef PADDLE_WITH_DNNL
+#include "paddle/fluid/platform/mkldnn_helper.h"
+#endif
+
+namespace paddle {
+namespace framework {
+
+PyLayerInstruction::PyLayerInstruction(
+    size_t id,
+    const platform::Place& place,
+    pir::Operation* op,
+    ValueExecutionInfo* value_exec_info,
+    interpreter::ExecutionConfig execution_config)
+    : InstructionBase(id, place) {
+  PADDLE_ENFORCE(op->isa<paddle::dialect::PyLayerOp>(),
+                 phi::errors::PreconditionNotMet(
+                     "Cond instruction only support pylayer op"));
+  auto pylayer_op = op->dyn_cast<paddle::dialect::PyLayerOp>();
+  op_ = op;
+
+  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))));
+  }
+  VLOG(6) << "finish process output_vars";
+
+  auto& fwd_block = pylayer_op.forward_block();
+  std::unordered_map<pir::Value, std::vector<int>> inputs;
+  GetInputIds(op, *value_exec_info, &inputs);
+  auto fwd_outside_inputs =
+      GetExternalInputs(&fwd_block, *value_exec_info, &inputs);
+
+  // NOTE(chenxi67): the variable corresponding to container value if a
+  // <VariableRefArray> Type. It will recursively get the ID of internal
+  // variables when use GetValueId() method. However, the copy_var pushed into
+  // the tuple does not have a corresponding ID, and will insert a -1. Here we
+  // remove the value of -1.
+  for (auto& item : inputs) {
+    auto& var_vec = item.second;
+    for (auto it = var_vec.begin(); it != var_vec.end();) {
+      if (*it == -1) {
+        it = var_vec.erase(it);
+      } else {
+        ++it;
+      }
+    }
+  }
+  SetInputs(inputs);
+
+  std::unordered_map<pir::Value, std::vector<int>> outputs;
+  for (size_t i = 0; i < op->num_results(); i++) {
+    pir::Value value = op->result(i);
+    if (value && value.type()) {
+      PADDLE_ENFORCE_EQ(
+          value_exec_info->HasValue(value),
+          true,
+          phi::errors::PreconditionNotMet(
+              "output should in name map, [%d] 'th output of [%s] op",
+              i,
+              "pylayer op"));
+      outputs.emplace(value, GetValueIds(value, *value_exec_info));
+    }
+  }
+
+  for (auto& item : outputs) {
+    auto& var_vec = item.second;
+    for (auto it = var_vec.begin(); it != var_vec.end();) {
+      if (*it == -1) {
+        it = var_vec.erase(it);
+      } else {
+        ++it;
+      }
+    }
+  }
+
+  SetOutputs(outputs);
+  VLOG(6) << "finish process inputs outputs index";
+
+  Scope* fwd_scope = &(value_exec_info->GetScope()->NewScope());
+  auto skip_gc_vars = execution_config.skip_gc_vars;
+  execution_config.skip_gc_vars.clear();
+  execution_config.create_local_scope = true;
+  fwd_inter_ = new PirInterpreter(place,
+                                  {},
+                                  &fwd_block,
+                                  fwd_scope,
+                                  value_exec_info->NewChild(fwd_scope),
+                                  execution_config);
+
+  std::set<std::string> fwd_skip_gc_names_set;
+  for (auto value : GetYiedOpInputs(&fwd_block)) {
+    fwd_outputs_.push_back(fwd_inter_->GetNameByValue(value));
+    fwd_skip_gc_names_.push_back(fwd_inter_->GetNameByValue(value));
+    fwd_skip_gc_names_set.insert(fwd_inter_->GetNameByValue(value));
+  }
+
+  // NOTE(zhangbo): According to the concept of control flow, child scopes
+  // should not control the lifecycle of parent scope variables.
+  for (auto value : fwd_outside_inputs) {
+    fwd_skip_gc_names_.push_back(fwd_inter_->GetNameByValue(value));
+    fwd_skip_gc_names_set.insert(fwd_inter_->GetNameByValue(value));
+  }
+  for (const auto& var_name : skip_gc_vars) {
+    fwd_skip_gc_names_.push_back(var_name);
+    fwd_skip_gc_names_set.insert(var_name);
+  }
+
+  fwd_inter_->SetSkipGcVars(fwd_skip_gc_names_set);
+  VLOG(6) << "finish process forward block interpreter";
+}
+
+PyLayerInstruction::~PyLayerInstruction() {
+  if (fwd_inter_ != nullptr) {
+    delete fwd_inter_;
+  }
+}
+
+void PyLayerInstruction::Run() {
+  VLOG(6) << "start pylayer forward block interpreter";
+
+#ifdef PADDLE_WITH_DNNL
+  // Executor on being destroyed clears oneDNN cache and resets
+  // registered model data layout. This is unwanted for nested
+  // Executors (executors declared inside control ops)
+  paddle::platform::DontClearMKLDNNCache(fwd_inter_->GetPlace());
+#endif
+  fwd_inter_->Run({}, false);
+  CopyBranchOutput(fwd_outputs_, output_vars_, fwd_inter_->InnerScope());
+}
+
+}  // namespace framework
+}  // namespace paddle
@@ -0,0 +1,64 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/fluid/framework/new_executor/instruction/instruction_base.h"
+#include "paddle/fluid/framework/new_executor/interpreter/execution_config.h"
+
+namespace ir {
+class Operation;
+}  // namespace ir
+
+namespace paddle {
+namespace framework {
+class Scope;
+class Value;
+class PirInterpreter;
+class ValueExecutionInfo;
+
+class PyLayerInstruction : public InstructionBase {
+ public:
+  PyLayerInstruction(size_t id,
+                     const platform::Place& place,
+                     ::pir::Operation* op,
+                     ValueExecutionInfo* value_exe_info,
+                     interpreter::ExecutionConfig execution_config);
+
+  ~PyLayerInstruction();
+
+  void Run() override;
+
+  const std::string& Name() const override { return name_; }
+
+  ::pir::Operation* Operation() const override { return op_; }
+
+  PirInterpreter* ForwardInterpreter() const { return fwd_inter_; }
+
+ private:
+  ::pir::Operation* op_;
+
+  std::string name_{"pylayer_instruction"};
+
+  std::vector<Variable*> output_vars_;
+
+  PirInterpreter* fwd_inter_ = nullptr;
+
+  std::vector<std::string> fwd_outputs_;
+
+  std::vector<std::string> fwd_skip_gc_names_;
+};
+
+}  // namespace framework
+}  // namespace paddle
@@ -418,5 +418,27 @@ bool GetCondData(const phi::DenseTensor& cond) {
   return cpu_cond->data<bool>()[0];
 }
 
+void CopyBranchOutput(const std::vector<std::string>& var_names,
+                      const std::vector<Variable*>& output_vars,
+                      Scope* inner_scope) {
+  for (size_t i = 0; i < var_names.size(); ++i) {
+    auto* inner_var = inner_scope->GetVar(var_names[i]);
+
+    if (inner_var->IsType<phi::DenseTensor>()) {
+      output_vars[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
+          inner_var->Get<phi::DenseTensor>());
+
+    } else if (inner_var->IsType<phi::TensorArray>()) {
+      const auto& inner_array = inner_var->Get<phi::TensorArray>();
+      auto* output_array = output_vars[i]->GetMutable<phi::TensorArray>();
+      // output_array->clear();
+      *output_array = inner_array;
+    } else {
+      PADDLE_THROW(
+          phi::errors::Unimplemented("unsupported type %d", inner_var->Type()));
+    }
+  }
+}
+
 }  // namespace framework
 }  // namespace paddle
@@ -66,5 +66,10 @@ void InsertInplacedExternalInputsToOuts(
     std::unordered_map<pir::Value, std::vector<int>>* outputs);
 
 bool GetCondData(const phi::DenseTensor& cond);
+
+void CopyBranchOutput(const std::vector<std::string>& var_names,
+                      const std::vector<Variable*>& output_vars,
+                      Scope* inner_scope);
+
 }  // namespace framework
 }  // namespace paddle
@@ -234,6 +234,7 @@ const std::unordered_set<std::string> SpecialOps = {
     paddle::dialect::DataOp::name(),
     pir::ShadowOutputOp::name(),
     paddle::dialect::IfOp::name(),
+    paddle::dialect::PyLayerOp::name(),
     paddle::dialect::WhileOp::name(),
     pir::StackCreateOp::name(),
 };
@@ -648,6 +649,13 @@ void HandleForSpecialOp(pir::Operation* op,
       auto if_op_out_value = if_op->result(i);
       BuildValue(if_op_out_value, var_name_prefix, value_exe_info);
     }
+  } else if (op->isa<paddle::dialect::PyLayerOp>()) {
+    auto pylayer_op = op->dyn_cast<paddle::dialect::PyLayerOp>();
+
+    for (size_t i = 0; i < pylayer_op->num_results(); ++i) {
+      auto pylayer_op_out_value = pylayer_op->result(i);
+      BuildValue(pylayer_op_out_value, var_name_prefix, value_exe_info);
+    }
   } else if (op->isa<paddle::dialect::WhileOp>()) {
     auto while_op = op->dyn_cast<paddle::dialect::WhileOp>();
 
 
@@ -45,10 +45,12 @@ namespace framework {
 
 class IfInstruction;
 class WhileInstruction;
+class PyLayerInstruction;
 class ValueExecutionInfo {
  public:
   friend class IfInstruction;
   friend class WhileInstruction;
+  friend class PyLayerInstruction;
 
   explicit ValueExecutionInfo(Scope* scope) : scope_(scope) {}
 
 
@@ -52,6 +52,7 @@
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/assert_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/has_elements_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.h"
+#include "paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/select_input_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/select_output_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/tuple_pop_instruction.h"
@@ -712,6 +713,14 @@ void PirInterpreter::BuildInstruction() {
             {&op.dyn_cast<paddle::dialect::IfOp>().false_block(),
              dynamic_cast<IfInstruction*>(vec_instruction_base_.back().get())
                  ->FalseBranchInterpreter()});
+      } else if (op.isa<paddle::dialect::PyLayerOp>()) {
+        vec_instruction_base_.emplace_back(std::make_unique<PyLayerInstruction>(
+            op_idx++, place_, &op, value_exe_info_.get(), execution_config_));
+        sub_blocks_.insert(
+            {&op.dyn_cast<paddle::dialect::PyLayerOp>().forward_block(),
+             dynamic_cast<PyLayerInstruction*>(
+                 vec_instruction_base_.back().get())
+                 ->ForwardInterpreter()});
       } else if (op.isa<paddle::dialect::WhileOp>()) {
         vec_instruction_base_.emplace_back(std::make_unique<WhileInstruction>(
             op_idx++, place_, &op, value_exe_info_.get(), execution_config_));