PaddlePaddle
diff --git a/‎paddle/fluid/framework/CMakeLists.txt
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions b/‎paddle/fluid/framework/CMakeLists.txt
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions
diff --git a/‎paddle/fluid/framework/executor_gc_helper.cc
Copy file name to clipboard
+43-21Lines changed: 43 additions & 21 deletions b/‎paddle/fluid/framework/executor_gc_helper.cc
Copy file name to clipboard
+43-21Lines changed: 43 additions & 21 deletions
diff --git a/‎paddle/fluid/framework/ir/graph_helper.cc
Copy file name to clipboard
+12-4Lines changed: 12 additions & 4 deletions b/‎paddle/fluid/framework/ir/graph_helper.cc
Copy file name to clipboard
+12-4Lines changed: 12 additions & 4 deletions
diff --git a/‎paddle/fluid/framework/ir/memory_optimize_pass/CMakeLists.txt
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/fluid/framework/ir/memory_optimize_pass/CMakeLists.txt
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/memory_optimize_pass/eager_deletion_pass.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/fluid/framework/ir/memory_optimize_pass/eager_deletion_pass.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/memory_optimize_pass/pylayer_op_eager_deletion_pass.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions b/‎paddle/fluid/framework/ir/memory_optimize_pass/pylayer_op_eager_deletion_pass.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/fluid/operators/controlflow/CMakeLists.txt
Copy file name to clipboard
+11-2Lines changed: 11 additions & 2 deletions b/‎paddle/fluid/operators/controlflow/CMakeLists.txt
Copy file name to clipboard
+11-2Lines changed: 11 additions & 2 deletions
diff --git a/‎paddle/fluid/operators/controlflow/pylayer_op.cc
Copy file name to clipboard
+25-4Lines changed: 25 additions & 4 deletions b/‎paddle/fluid/operators/controlflow/pylayer_op.cc
Copy file name to clipboard
+25-4Lines changed: 25 additions & 4 deletions
@@ -696,6 +696,7 @@ cc_library(
   DEPS while_op_helper
        recurrent_op_helper
        conditional_block_op_helper
+       pylayer_op_helper
        scope
        proto_desc
        operator
@@ -1014,8 +1015,9 @@ else()
          monitor)
 endif()
 
-target_link_libraries(executor while_op_helper executor_gc_helper
-                      recurrent_op_helper conditional_block_op_helper)
+target_link_libraries(
+  executor while_op_helper executor_gc_helper recurrent_op_helper
+  conditional_block_op_helper pylayer_op_helper)
 
 cc_library(
   parallel_executor
 
@@ -24,6 +24,7 @@
 #include "paddle/fluid/framework/operator.h"
 #include "paddle/fluid/framework/var_desc.h"
 #include "paddle/fluid/operators/controlflow/conditional_block_op_helper.h"
+#include "paddle/fluid/operators/controlflow/pylayer_op_helper.h"
 #include "paddle/fluid/operators/controlflow/recurrent_op_helper.h"
 #include "paddle/fluid/operators/controlflow/while_op_helper.h"
 #include "paddle/fluid/platform/enforce.h"
@@ -226,6 +227,8 @@ GetEagerDeletionCleanVarsForPartial(const ProgramDesc &origin_program,
     auto global_block_ops = CreateOpsFromBlock(program.Block(0));
     operators::PrepareSafeEagerDeletionOnConditionalOpAndConditionalGradOp(
         program, 0, global_block_ops);
+    operators::PrepareSafeEagerDeletionOnPyLayerOpAndPyLayerGradOp(
+        program, 0, global_block_ops);
     operators::PrepareSafeEagerDeletionOnWhileOpAndWhileGradOp(
         program, 0, global_block_ops);
     operators::PrepareSafeEagerDeletionOnRecurrentOpAndRecurrentGradOp(
@@ -240,35 +243,54 @@ GetEagerDeletionCleanVarsForPartial(const ProgramDesc &origin_program,
 
   const char *kSubBlock = "sub_block";
   const char *kSkipEagerDeletionVars = "skip_eager_deletion_vars";
+  // NOTE: pylayer op contains may contain two blocks: forward block and
+  // backward block
+  const char *kBlocks = "blocks";
 
   for (size_t i = 0; i < block_num; ++i) {
     const auto &block = program.Block(i);
     size_t op_num = block.OpSize();
     for (size_t j = 0; j < op_num; ++j) {
       auto *op = block.Op(static_cast<int>(j));
-      if (!op->HasAttr(kSubBlock) || !op->HasAttr(kSkipEagerDeletionVars)) {
+      if ((!op->HasAttr(kSubBlock) && !op->HasAttr(kBlocks)) ||
+          !op->HasAttr(kSkipEagerDeletionVars)) {
         continue;
       }
-      auto sub_block_id = op->GetAttrIfExists<BlockDesc *>(kSubBlock)->ID();
-      PADDLE_ENFORCE_GE(sub_block_id,
-                        0,
-                        platform::errors::PermissionDenied(
-                            "sub_block id must be non-negative number"));
-      PADDLE_ENFORCE_LT(sub_block_id,
-                        block_num,
-                        platform::errors::PermissionDenied(
-                            "sub_block id exceeds max block num"));
-      PADDLE_ENFORCE_EQ(
-          found_skip_vars[sub_block_id],
-          false,
-          platform::errors::PermissionDenied(
-              "there are 2 ops which refer to the same sub_block %d",
-              sub_block_id));
-
-      found_skip_vars[sub_block_id] = true;
-      auto sub_block_skip_vars =
-          op->GetAttrIfExists<std::vector<std::string>>(kSkipEagerDeletionVars);
-      skip_vars_on_each_block[sub_block_id] = std::move(sub_block_skip_vars);
+
+      std::vector<int32_t> sub_block_ids;
+      if (op->HasAttr(kSubBlock)) {
+        sub_block_ids.push_back(
+            op->GetAttrIfExists<BlockDesc *>(kSubBlock)->ID());
+      } else if (op->HasAttr(kBlocks)) {
+        const auto &blocks =
+            op->GetAttrIfExists<std::vector<BlockDesc *>>(kBlocks);
+        for (const auto &block : blocks) {
+          sub_block_ids.push_back(block->ID());
+        }
+      }
+
+      for (auto sub_block_id : sub_block_ids) {
+        PADDLE_ENFORCE_GE(sub_block_id,
+                          0,
+                          platform::errors::PermissionDenied(
+                              "sub_block id must be non-negative number"));
+        PADDLE_ENFORCE_LT(sub_block_id,
+                          block_num,
+                          platform::errors::PermissionDenied(
+                              "sub_block id exceeds max block num"));
+        PADDLE_ENFORCE_EQ(
+            found_skip_vars[sub_block_id],
+            false,
+            platform::errors::PermissionDenied(
+                "there are 2 ops which refer to the same sub_block %d",
+                sub_block_id));
+
+        found_skip_vars[sub_block_id] = true;
+        auto sub_block_skip_vars =
+            op->GetAttrIfExists<std::vector<std::string>>(
+                kSkipEagerDeletionVars);
+        skip_vars_on_each_block[sub_block_id] = std::move(sub_block_skip_vars);
+      }
     }
   }
 
 
@@ -704,10 +704,18 @@ static void GetGraphOpDesc(const std::vector<Node *> &nodes,
         ops->emplace_back(depend_desc);
         VLOG(4) << "add depend op";
       }
-      if (n->Name() == "while" || n->Name() == "while_grad" ||
-          n->Name() == "conditional_block" ||
-          n->Name() == "conditional_block_grad" || n->Name() == "recurrent" ||
-          n->Name() == "recurrent_grad") {
+
+      const std::unordered_set<std::string> control_flow_ops = {
+          "while",
+          "while_grad",
+          "conditional_block",
+          "conditional_block_grad",
+          "recurrent",
+          "recurrent_grad",
+          "pylayer",
+          "pylayer_grad"};
+
+      if (control_flow_ops.count(n->Name())) {
         VLOG(1) << "Update control op attr: skip_eager_deletion_vars";
         UpdateControlOpSkipEagerDeletionVars(*n, graph, graph_idx, n->Name());
       }
 
@@ -6,6 +6,10 @@ cc_library(
   conditional_block_op_eager_deletion_pass
   SRCS conditional_block_op_eager_deletion_pass.cc
   DEPS conditional_block_op_helper graph_helper pass computation_op_handle)
+cc_library(
+  pylayer_op_eager_deletion_pass
+  SRCS pylayer_op_eager_deletion_pass.cc
+  DEPS pylayer_op_helper graph_helper pass computation_op_handle)
 cc_library(
   while_op_eager_deletion_pass
   SRCS while_op_eager_deletion_pass.cc
@@ -31,6 +35,7 @@ set(EAGER_DELETETION_PASS_DEPS
     graph_helper
     pass
     conditional_block_op_eager_deletion_pass
+    pylayer_op_eager_deletion_pass
     while_op_eager_deletion_pass
     recurrent_op_eager_deletion_pass
     reference_count_pass_helper)
 
@@ -294,6 +294,10 @@ void EagerDeletionPass::ApplyImpl(ir::Graph *graph) const {
           "conditional_block_op_eager_deletion_pass");
   conditional_block_op_eager_deletion_pass->Apply(graph);
 
+  auto pylayer_op_eager_deletion_pass =
+      ir::PassRegistry::Instance().Get("pylayer_op_eager_deletion_pass");
+  pylayer_op_eager_deletion_pass->Apply(graph);
+
   auto while_op_eager_deletion_pass =
       ir::PassRegistry::Instance().Get("while_op_eager_deletion_pass");
   while_op_eager_deletion_pass->Apply(graph);
@@ -321,6 +325,7 @@ REGISTER_PASS(eager_deletion_pass, paddle::framework::ir::EagerDeletionPass)
     .RequirePassAttr(paddle::framework::ir::kGarbageCollector);
 
 USE_PASS(conditional_block_op_eager_deletion_pass);
+USE_PASS(pylayer_op_eager_deletion_pass);
 USE_PASS(while_op_eager_deletion_pass);
 USE_PASS(recurrent_op_eager_deletion_pass);
 #ifdef PADDLE_WITH_CINN
 
@@ -0,0 +1,102 @@
+// Copyright (c) 2019 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/details/computation_op_handle.h"
+#include "paddle/fluid/framework/ir/graph_helper.h"
+#include "paddle/fluid/framework/ir/pass.h"
+#include "paddle/fluid/operators/controlflow/op_variant.h"
+#include "paddle/fluid/operators/controlflow/pylayer_op_helper.h"
+namespace paddle {
+namespace framework {
+namespace ir {
+using OpVariant = operators::OpVariant;
+class PyLayerOpEagerDeletionPass : public Pass {
+ protected:
+  void ApplyImpl(Graph *graph) const override {
+    auto all_ops = ir::FilterByNodeWrapper<details::OpHandleBase>(*graph);
+
+    // Find all pylayer_op and pylayer_grad_op
+    std::unordered_map<
+        size_t,
+        std::pair<std::vector<OpVariant>, std::vector<OpVariant>>>
+        target_ops;
+    for (auto *op : all_ops) {
+      auto compute_op = dynamic_cast<details::ComputationOpHandle *>(op);
+      if (compute_op == nullptr) continue;
+
+      if (compute_op->Name() == "pylayer") {
+        target_ops[compute_op->GetScopeIdx()].first.emplace_back(
+            compute_op->GetOp());
+      } else if (compute_op->Name() == "pylayer_grad") {
+        target_ops[compute_op->GetScopeIdx()].second.emplace_back(
+            compute_op->GetOp());
+      }
+    }
+
+    // NOTE(Aurelius84): In case of @to_static, after we finish executing
+    // forward graph, some necessaray variable in step_scope of pylayer_op
+    // should be kept for backward graph.
+    if (graph->IsConstructedByPartialProgram()) {
+      PADDLE_ENFORCE_LE(target_ops.size(),
+                        1,
+                        platform::errors::InvalidArgument(
+                            "Unsupported multi devices if graph is constructed "
+                            "with partial program."));
+      size_t scope_idx = 0;
+      auto &pylayer_ops = target_ops[scope_idx].first;
+      auto &pylayer_grad_ops = target_ops[scope_idx].second;
+
+      auto all_ops = graph->OriginProgram().Block(0).AllOps();
+      if (pylayer_ops.empty()) {
+        operators::AppendOpVariantByOpName(
+            all_ops, std::string("pylayer"), &pylayer_ops);
+      } else if (pylayer_grad_ops.empty()) {
+        operators::AppendOpVariantByOpName(
+            all_ops, std::string("pylayer_grad"), &pylayer_grad_ops);
+      } else {
+        PADDLE_THROW("One of pylayer_ops or pylayer_grad_ops should be empty.");
+      }
+    }
+
+    for (auto &ops_pair : target_ops) {
+      auto &pylayer_ops = ops_pair.second.first;
+      auto &pylayer_grad_ops = ops_pair.second.second;
+      operators::PrepareSafeEagerDeletionOnPyLayerOpAndPyLayerGradOp(
+          graph->OriginProgram(), pylayer_ops, pylayer_grad_ops);
+    }
+
+    for (auto op_hander : all_ops) {
+      auto *compute_op =
+          dynamic_cast<details::ComputationOpHandle *>(op_hander);
+      if (compute_op == nullptr) continue;
+      if (compute_op->Name() == "pylayer" ||
+          compute_op->Name() == "pylayer_grad") {
+        ir::Node *op_node = op_hander->Node();
+        auto *op_base = compute_op->GetOp();
+        if (op_base->Attrs().count("skip_eager_deletion_vars")) {
+          op_node->Op()->SetAttr(
+              "skip_eager_deletion_vars",
+              op_base->Attrs().at("skip_eager_deletion_vars"));
+        }
+      }
+    }
+  }
+};
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(pylayer_op_eager_deletion_pass,
+              paddle::framework::ir::PyLayerOpEagerDeletionPass);
@@ -25,6 +25,7 @@
 #include "paddle/fluid/framework/new_executor/interpreter/static_build.h"
 #include "paddle/fluid/memory/stats.h"
 #include "paddle/fluid/operators/controlflow/conditional_block_op_helper.h"
+#include "paddle/fluid/operators/controlflow/pylayer_op_helper.h"
 #include "paddle/fluid/operators/controlflow/recurrent_op_helper.h"
 #include "paddle/fluid/operators/controlflow/while_op_helper.h"
 #include "paddle/fluid/operators/ops_extra_info.h"
@@ -571,6 +572,8 @@ void BuildOpFuncList(const platform::Place& place,
     const ProgramDesc& main_program = *block.Program();
     operators::PrepareSafeEagerDeletionOnConditionalOpAndConditionalGradOp(
         main_program, block.ID(), ops_unique);
+    operators::PrepareSafeEagerDeletionOnPyLayerOpAndPyLayerGradOp(
+        main_program, block.ID(), ops_unique);
     operators::PrepareSafeEagerDeletionOnWhileOpAndWhileGradOp(
         main_program, block.ID(), ops_unique);
     operators::PrepareSafeEagerDeletionOnRecurrentOpAndRecurrentGradOp(
@@ -611,6 +614,8 @@ void BuildOpFuncList(const platform::Place& place,
     const std::set<std::string> ops_with_var_not_in_scope = {
         "conditional_block",
         "conditional_block_grad",
+        "pylayer",
+        "pylayer_grad"
         "recurrent_grad",
         "rnn_memory_helper",
         "rnn_memory_helper_grad",
 
@@ -3,13 +3,17 @@ if(WITH_UNITY_BUILD)
   # Load Unity Build rules for operators in paddle/fluid/operators/controlflow.
   include(unity_build_rule.cmake)
 endif()
-register_operators(EXCLUDES conditional_block_op DEPS naive_executor
+register_operators(EXCLUDES conditional_block_op pylayer_op DEPS naive_executor
                    standalone_executor)
 
 cc_library(
   conditional_block_op
   SRCS conditional_block_op.cc
   DEPS executor standalone_executor)
+cc_library(
+  pylayer_op
+  SRCS pylayer_op.cc
+  DEPS standalone_executor)
 cc_library(
   op_variant
   SRCS op_variant.cc
@@ -18,6 +22,10 @@ cc_library(
   conditional_block_op_helper
   SRCS conditional_block_op_helper.cc
   DEPS op_variant operator conditional_block_op)
+cc_library(
+  pylayer_op_helper
+  SRCS pylayer_op_helper.cc
+  DEPS op_variant operator pylayer_op)
 cc_library(
   recurrent_op_helper
   SRCS recurrent_op_helper.cc
@@ -28,7 +36,8 @@ cc_library(
   DEPS op_variant operator)
 
 if(WITH_UNITY_BUILD)
-  target_link_libraries(paddle_operators_controlflow_unity conditional_block_op)
+  target_link_libraries(paddle_operators_controlflow_unity conditional_block_op
+                        pylayer_op)
 else()
   target_link_libraries(conditional_block_infer_op conditional_block_op)
 endif()
 
@@ -51,14 +51,35 @@ void PyLayerOp::CreateInterpreter(
         dev_place, block, cur_scope, execution_config));
     VLOG(10) << "[interpreterCore] created:" << core_;
   } else {
-    // NOTE: Borrowed from
-    // `paddle/fluid/operators/controlflow/control_flow_op_helper.h`
-    // TODO(MarioLulab): Add PyLayer Helper ?
     BuildScopeForControlFlowOp(*core_, block, cur_scope);
     core_->reset_scope(cur_scope);
   }
 }
 
+class PyLayerForwardOpProtoMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput(PyLayerOp::kInputs, "The input variables of the sub-block.")
+        .AsDuplicable();
+    AddOutput(PyLayerOp::kOutputs, "The output variables of the sub-block.")
+        .AsDuplicable();
+    AddOutput(
+        PyLayerOp::kScope,
+        "(std::vector<Scope*>) The scope of static pylayer block, used for "
+        "passing intermediate variables between forward and backward.");
+    AddAttr<std::vector<framework::BlockDesc *>>(
+        "blocks",
+        "The blocks of PyLayer operator where blocks[0] indicates the forward "
+        "block and blocks[1] indicates the backward block.");
+    AddComment(R"DOC(PyLayer operator
+
+The PyLayer Operator is designed to support `@to_static` for `PyLayer in Dynamic Graph`.
+
+
+)DOC");
+  }
+};
+
 class PyLayerForwardOp : public PyLayerOp {
  public:
   PyLayerForwardOp(const std::string &type,
@@ -109,7 +130,7 @@ class PyLayerForwardOp : public PyLayerOp {
 class PyLayerForwardInferShape : public framework::InferShapeBase {
  public:
   void operator()(framework::InferShapeContext *context) const override {
-    // TODO(MarioLulab): do nothing.
+    // NOTE(MarioLulab): do nothing.
   }
 };