PaddlePaddle
diff --git a/‎paddle/fluid/framework/prune.cc
Copy file name to clipboard
+155-26Lines changed: 155 additions & 26 deletions b/‎paddle/fluid/framework/prune.cc
Copy file name to clipboard
+155-26Lines changed: 155 additions & 26 deletions
diff --git a/‎python/paddle/jit/dy2static/py_layer.py
Copy file name to clipboard
+13-2Lines changed: 13 additions & 2 deletions b/‎python/paddle/jit/dy2static/py_layer.py
Copy file name to clipboard
+13-2Lines changed: 13 additions & 2 deletions
diff --git a/‎python/paddle/static/io.py
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions b/‎python/paddle/static/io.py
Copy file name to clipboard
+14Lines changed: 14 additions & 0 deletions
@@ -30,6 +30,8 @@ const char kRecurrent[] = "recurrent";  // NOLINT
 const char kStates[] = "states";        // NOLINT
 const char kExStates[] = "ex_states";   // NOLINT
 
+const char kPyLayer[] = "pylayer";  // NOLINT
+
 bool HasDependentInputVar(
     const proto::OpDesc& op_desc,
     const std::unordered_set<std::string>& dependent_vars) {
@@ -86,6 +88,23 @@ int GetSubBlockIndex(const proto::OpDesc& op_desc) {
   return -1;
 }
 
+void GetSubBlocksIndices(const proto::OpDesc& op_desc,
+                         std::vector<int>* indices) {
+  for (auto& attr : op_desc.attrs()) {
+    if (attr.type() == proto::AttrType::BLOCKS) {
+      PADDLE_ENFORCE_GT(
+          attr.blocks_idx_size(),
+          0,
+          platform::errors::NotFound(
+              "Attribute blocks is not found in operator %s", op_desc.type()));
+      indices->resize(attr.blocks_idx_size());
+      for (int i = 0; i < attr.blocks_idx_size(); i++) {
+        (*indices)[i] = attr.blocks_idx(i);
+      }
+    }
+  }
+}
+
 void SetSubBlockIndex(proto::OpDesc* op_desc, int sub_idx) {
   for (auto& attr : *op_desc->mutable_attrs()) {
     if (attr.type() == proto::AttrType::BLOCK) {
@@ -99,10 +118,43 @@ void SetSubBlockIndex(proto::OpDesc* op_desc, int sub_idx) {
   }
 }
 
+void SetSubBlocksIndices(proto::OpDesc* op_desc,
+                         const std::vector<int>& sub_indices) {
+  for (auto& attr : *op_desc->mutable_attrs()) {
+    if (attr.type() == proto::AttrType::BLOCKS) {
+      PADDLE_ENFORCE_GT(
+          attr.blocks_idx_size(),
+          0,
+          platform::errors::NotFound(
+              "Attribute blocks is not found in operator %s", op_desc->type()));
+      attr.clear_blocks_idx();
+      for (auto idx : sub_indices) {
+        attr.add_blocks_idx(idx);
+      }
+    }
+  }
+}
+
 bool HasSubBlock(const proto::OpDesc& op_desc) {
   return GetSubBlockIndex(op_desc) > 0;
 }
 
+bool HasSubBlocks(const proto::OpDesc& op_desc) {
+  // ``blocks_idx_size() == 0`` indicates no sub blocks.
+  for (auto& attr : op_desc.attrs()) {
+    if (attr.type() == proto::AttrType::BLOCKS) {
+      PADDLE_ENFORCE_GT(
+          attr.blocks_idx_size(),
+          0,
+          platform::errors::NotFound(
+              "Attribute blocks is not found in operator %s", op_desc.type()));
+      return true;
+    }
+  }
+
+  return false;
+}
+
 int GetOpRole(const proto::OpDesc& op_desc) {
   for (auto& attr : op_desc.attrs()) {
     if (attr.name() == OpProtoAndCheckerMaker::OpRoleAttrName()) {
@@ -150,14 +202,15 @@ int FindMapByValue(const std::map<int, int>& m, int val) {
 }
 
 // In other two cases, the op that has feed vars as output vars is dependent:
-// 1. op has subblock, like while/for/ifelse/recurrent
+// 1. op has subblock, like while/for/ifelse/recurrent/pylayer
 // 2. op is in subblock
 bool IsSubBlockDependent(const proto::OpDesc& op_desc,
                          const std::set<std::string>& feed_vars,
                          int parent_block_id) {
   for (auto& var : op_desc.outputs()) {
     for (auto& argu : var.arguments()) {
-      if ((HasSubBlock(op_desc) || parent_block_id != -1) &&
+      if ((HasSubBlock(op_desc) || HasSubBlocks(op_desc) ||
+           parent_block_id != -1) &&
           feed_vars.count(argu) != 0) {
         return true;
       }
@@ -289,7 +342,7 @@ void prune_impl(const proto::ProgramDesc& input,
     if (should_run[i]) {
       auto* op = op_field->Add();
       *op = input.blocks(block_id).ops(static_cast<int>(i));
-      if (HasSubBlock(*op)) {
+      if (HasSubBlock(*op) || HasSubBlocks(*op)) {
         VLOG(2) << "Pruning op which has sub block: " << op->type();
         // create sub_block_dependent_vars here to help prune the sub block
         std::unordered_set<std::string> sub_block_dependent_vars;
@@ -321,15 +374,41 @@ void prune_impl(const proto::ProgramDesc& input,
             }
           }
         }
-        // GetSubBlockIndex(*op) is the idx of the sub_block in the input desc
-        // output_block_id is the idx of the current block in the output desc
-        prune_impl(input,
-                   output,
-                   GetSubBlockIndex(*op),
-                   output_block_id,
-                   &sub_block_dependent_vars,
-                   feed_var_names,
-                   pruned_origin_block_id_map);
+        if (HasSubBlock(*op)) {
+          // GetSubBlockIndex(*op) is the idx of the sub_block in the input desc
+          // output_block_id is the idx of the current block in the output desc
+          prune_impl(input,
+                     output,
+                     GetSubBlockIndex(*op),
+                     output_block_id,
+                     &sub_block_dependent_vars,
+                     feed_var_names,
+                     pruned_origin_block_id_map);
+        } else if (HasSubBlocks(*op)) {
+          // GetSubBlocksIndices(*op) are the indices of the sub_blocks in the
+          // input desc output_block_id is the idx of the current block in the
+          // output desc
+          std::vector<int> sub_indices;
+          GetSubBlocksIndices(*op, &sub_indices);
+          for (auto& sub_index : sub_indices) {
+            // create a copy of dependent_vars to avoid being overwrited by the
+            // other sub_block
+            std::unordered_set<std::string> dependent_vars_copy =
+                sub_block_dependent_vars;
+            prune_impl(input,
+                       output,
+                       sub_index,
+                       output_block_id,
+                       &dependent_vars_copy,
+                       feed_var_names,
+                       pruned_origin_block_id_map);
+          }
+        } else {
+          PADDLE_ENFORCE(false,
+                         platform::errors::PreconditionNotMet(
+                             "Attr Block or Blocks must exist when recursively "
+                             "calling prune_impl"));
+        }
       }
     }
   }
@@ -402,12 +481,29 @@ std::map<int, int> Prune(const proto::ProgramDesc& input,
         int origin_sub_idx = GetSubBlockIndex(op_desc);
         auto sub_idx =
             FindMapByValue(pruned_origin_block_id_map, origin_sub_idx);
-        PADDLE_ENFORCE_NE(sub_idx,
-                          -1,
-                          platform::errors::NotFound(
-                              "The origin sub block id should be found in "
-                              "pruned_progin_block_id_map"));
+        PADDLE_ENFORCE_NE(
+            sub_idx,
+            -1,
+            platform::errors::NotFound(
+                "The origin sub block id should be found in "
+                "pruned_progin_block_id_map when the op has sub_block"));
         SetSubBlockIndex(&op_desc, sub_idx);
+      } else if (HasSubBlocks(op_desc)) {
+        std::vector<int> origin_sub_indices;
+        GetSubBlocksIndices(op_desc, &origin_sub_indices);
+        std::vector<int> sub_indices;
+        for (int index : origin_sub_indices) {
+          auto sub_idx = FindMapByValue(pruned_origin_block_id_map, index);
+          PADDLE_ENFORCE_NE(
+              sub_idx,
+              -1,
+              platform::errors::NotFound(
+                  "The origin sub block id should be found in "
+                  "pruned_progin_block_id_map when the op has sub_blocks"));
+          sub_indices.push_back(sub_idx);
+        }
+
+        SetSubBlocksIndices(&op_desc, sub_indices);
       }
     }
   }
@@ -441,6 +537,19 @@ void PruneBackwardImpl(proto::BlockDesc* origin, proto::BlockDesc* pruned) {
       AppendOpInputVarNames(op_desc, &op_input_vars);
       AppendOpOutputVarNames(op_desc, &op_output_vars);
       *op = op_desc;
+
+      // if the type of op is "pylayer", we need to update the ``blocks``
+      // attribute because the backward block will be pruned
+      if (op->type() == kPyLayer && HasSubBlocks(*op)) {
+        std::vector<int> sub_indices;
+        GetSubBlocksIndices(*op, &sub_indices);
+        if (sub_indices.size() > 1) {
+          // sub_indices contains both forward block id and backward block id
+          std::vector<int> new_sub_indices(sub_indices.begin(),
+                                           sub_indices.end() - 1);
+          SetSubBlocksIndices(op, new_sub_indices);
+        }
+      }
     }
   }
 
@@ -471,9 +580,10 @@ std::tuple<framework::ProgramDesc, std::map<int, int>> PruneBackward(
   // Copy original ProgramDesc, origin can't be change
   framework::ProgramDesc origin_clone(origin);
 
-  // Step 1. check if the program contains grad loss operator.
-  // If not, the program need no pruning.
+  // Step 1. check if the program contains grad loss operator or pylayer
+  // operator. If not, the program need no pruning.
   bool has_loss_grad_op = false;
+  bool has_pylayer_op = false;
   std::queue<int> block_contains_loss;
   std::queue<int> block_contains_loss_grad;
   for (size_t i = 0; i < origin_clone.Size(); i++) {
@@ -485,13 +595,15 @@ std::tuple<framework::ProgramDesc, std::map<int, int>> PruneBackward(
                       static_cast<int>(OpRole::kLoss))) {
         op->SetIsTarget(false);
         has_loss_grad_op = true;
-        break;
+      }
+      if (op->Type() == kPyLayer) {
+        has_pylayer_op = true;
       }
     }
   }
 
   std::map<int, int> pruned_progin_block_id_map;
-  if (!has_loss_grad_op) {
+  if (!has_loss_grad_op && !has_pylayer_op) {
     // No pruning, fast return a copy of the origin ProgramDesc with an empty
     // map, means default mapped, i.e.{0:0, 1:1, ..., n:n}.
     return std::make_tuple(framework::ProgramDesc(origin_clone),
@@ -544,12 +656,29 @@ std::tuple<framework::ProgramDesc, std::map<int, int>> PruneBackward(
         int origin_sub_idx = GetSubBlockIndex(op_desc);
         auto sub_idx =
             FindMapByValue(pruned_progin_block_id_map, origin_sub_idx);
-        PADDLE_ENFORCE_NE(sub_idx,
-                          -1,
-                          platform::errors::NotFound(
-                              "The origin sub block id is not found in "
-                              "pruned_progin_block_id_map"));
+        PADDLE_ENFORCE_NE(
+            sub_idx,
+            -1,
+            platform::errors::NotFound(
+                "The origin sub block id is not found in "
+                "pruned_progin_block_id_map when the op has sub_block"));
         SetSubBlockIndex(&op_desc, sub_idx);
+      } else if (HasSubBlocks(op_desc)) {
+        std::vector<int> origin_sub_indices;
+        GetSubBlocksIndices(op_desc, &origin_sub_indices);
+        std::vector<int> sub_indices;
+        for (int index : origin_sub_indices) {
+          auto sub_idx = FindMapByValue(pruned_progin_block_id_map, index);
+          PADDLE_ENFORCE_NE(
+              sub_idx,
+              -1,
+              platform::errors::NotFound(
+                  "The origin sub block id should be found in "
+                  "pruned_progin_block_id_map when the op has sub_blocks"));
+          sub_indices.push_back(sub_idx);
+        }
+
+        SetSubBlocksIndices(&op_desc, sub_indices);
       }
     }
   }
 
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import functools
+import inspect
 
 from paddle.base.framework import Variable
 from paddle.common_ops_import import LayerHelper
@@ -73,9 +74,19 @@ def __init__(self, dyfunc_self):
         )
 
     # NOTE: only support position args and Variables Now
-    def apply(self, *args):
+    def apply(self, *args, **kwargs):
+        # rearrange `position-args + keyword-args` into `position-args`
+        dyfunc_sig = inspect.signature(self.dyfunc_self.forward)
+        bound_args = dyfunc_sig.bind(self.dyfunc_self, *args, **kwargs)
+        bound_args.apply_defaults()
+        input_args = [
+            item
+            for i, item in enumerate(bound_args.arguments.values())
+            if i > 0
+        ]  # index 0 indicate `dyfunc_self` which shouldn't be put into `input_args`
+
         return static_pylayer(
             forward_fn=self.forward_fn_with_ctx,
-            inputs=list(args),
+            inputs=input_args,
             backward_fn=self.backward_fn_with_ctx,
         )
@@ -274,6 +274,20 @@ def normalize_program(program, feed_vars, fetch_vars, **kwargs):
         op.desc.set_is_target(False)
         if op.type == "feed" or op.type == "fetch":
             remove_op_idx.append(i)
+
+        if op.type == "pylayer":
+            sub_blocks_ids = op._blocks_attr_ids("blocks")
+            if len(sub_blocks_ids) > 1:
+                # pylayer op ``blocks`` attr contains forward block id and backward block id
+                backward_block_id = sub_blocks_ids[-1]
+                # remove backward block
+                copy_program.blocks.pop(backward_block_id)
+                # update attrs ``blocks``
+                reserverd_blocks = []
+                for block_id in sub_blocks_ids[:-1]:
+                    reserverd_blocks.append(copy_program.block(block_id))
+                op._update_desc_attr("blocks", reserverd_blocks)
+
     for idx in remove_op_idx[::-1]:
         global_block._remove_op(idx)
     copy_program.desc.flush()