text: // 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
  //
  //     http://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/operators/controlflow/recurrent_op_helper.h"
  
  #include <algorithm>
  #include <string>
  #include <unordered_set>
  #include <utility>
  
  #include "paddle/fluid/framework/program_desc.h"
  #include "paddle/fluid/operators/recurrent_op.h"
  
  namespace paddle {
  namespace operators {
  
  static bool IsMatchedRecurrentOpAndRecurrentGradOp(const OpVariant &fwd_op,
                                                     const OpVariant &grad_op) {
    return fwd_op.Inputs().at(RecurrentBase::kInputs) ==
               grad_op.Inputs().at(RecurrentBase::kInputs) &&
           fwd_op.Outputs().at(RecurrentBase::kOutputs) ==
               grad_op.Inputs().at(RecurrentBase::kOutputs);
  }
  
  // Returns whether the variable is skippable in forward recurrent op
  // The variable is skippable in recurrent_op when the variable used in
  // recurrent_grad is not from grad_block.
  static bool IsSkippableVar(const std::string &name,
                             framework::BlockDesc *grad_block) {
    return name != framework::kEmptyVarName && !grad_block->HasVar(name);
  }
  
  static void ClearSkipVars(const OpVariant &op) {
    auto &attrs = const_cast<framework::AttributeMap &>(op.Attrs());
    std::vector<std::string> &attr_skip_vars =
        boost::get<std::vector<std::string>>(
            attrs[RecurrentBase::kSkipEagerDeletionVars]);
    attr_skip_vars.clear();
  }
  
  // Add skip vars into op's attribute
  template <class Container>
  static void AddSkipVars(const OpVariant &op, const Container &skip_vars) {
    auto &attrs = const_cast<framework::AttributeMap &>(op.Attrs());
    VLOG(2) << "Prepare to add " << skip_vars.size()
            << " skip var(s): " << paddle::string::join_strings(skip_vars, ' ');
    std::vector<std::string> &attr_skip_vars =
        boost::get<std::vector<std::string>>(
            attrs[RecurrentBase::kSkipEagerDeletionVars]);
    attr_skip_vars.insert(attr_skip_vars.end(), skip_vars.cbegin(),
                          skip_vars.cend());
  }
  
  // Find all ops and grad ops with given type name. The ops and grad ops
  // may locate in different blocks so we should traverse all blocks in the
  // program and find them out
  static void FindAllOpAndGradOp(const framework::ProgramDesc &program,
                                 OpAndGradOpPair *op_and_grad_op,
                                 const std::string &type_name,
                                 const std::string &backward_type_name) {
    OpVariantSet &ops = op_and_grad_op->first;
    OpVariantSet &grad_ops = op_and_grad_op->second;
  
    PADDLE_ENFORCE_GE(ops.size(), grad_ops.size(),
                      "There are extra grad ops in the graph or program");
  
    for (size_t i = 1; i < program.Size(); ++i) {
      auto &block = program.Block(i);
      for (size_t j = 0; j < block.OpSize(); ++j) {
        auto *op = block.Op(j);
        if (op->Type() == type_name) {
          ops.emplace(op);
        } else if (op->Type() == backward_type_name) {
          grad_ops.emplace(op);
        }
      }
    }
  
    PADDLE_ENFORCE_GE(ops.size(), grad_ops.size(),
                      "There are extra grad ops in the graph or program");
  }
  
  // Returns GradVarName of input var names
  static std::vector<std::string> GradVarLists(
      const std::vector<std::string> &var_names) {
    std::vector<std::string> retv;
    retv.reserve(var_names.size());
    std::transform(var_names.begin(), var_names.end(), std::back_inserter(retv),
                   framework::GradVarName);
    return retv;
  }
  
  // Add memory vars in recurrent op as skip vars.
  static void AddOpMemVarsAsSkip(const OpVariant &op, bool set_grad_mem_vars) {
    bool has_state = op.Attr<bool>(RecurrentBase::kHasStates);
    if (has_state) {
      std::unordered_set<std::string> skip_vars;
  
      auto &mem_vars = op.Attr<std::vector<std::string>>(RecurrentBase::kStates);
      skip_vars.insert(mem_vars.begin(), mem_vars.end());
  
      auto &pre_mem_vars =
          op.Attr<std::vector<std::string>>(RecurrentBase::kExStates);
      skip_vars.insert(pre_mem_vars.begin(), pre_mem_vars.end());
  
      if (set_grad_mem_vars) {
        auto mem_grad_vars = GradVarLists(mem_vars);
        skip_vars.insert(mem_grad_vars.begin(), mem_grad_vars.end());
        auto pre_mem_grad_vars = GradVarLists(pre_mem_vars);
        skip_vars.insert(pre_mem_grad_vars.begin(), pre_mem_grad_vars.end());
      }
      AddSkipVars(op, skip_vars);
    }
  }
  
  // Set outputs and memory vars of the input forward op as skip vars
  static void SetRecurrentForwardOpOnlySkipVarAttr(const OpVariant &fwd_op) {
    ClearSkipVars(fwd_op);
  
    AddOpMemVarsAsSkip(fwd_op, /* set_grad_mem_vars = */ false);
    auto &output_vars = fwd_op.Outputs().at(RecurrentBase::kOutputs);
    AddSkipVars(fwd_op, output_vars);
  }
  
  // Set skip vars of matched recurrent op and recurrent_grad op
  static void SetRecurrentOpAndRecurrentGradOpSkipVarAttr(
      const OpVariant &fwd_op, const OpVariant &bwd_op) {
    // Find all skippable variables in forward recurrent_op
    ClearSkipVars(fwd_op);
    AddOpMemVarsAsSkip(fwd_op, /* set_grad_mem_vars = */ false);
  
    auto *grad_block =
        bwd_op.Attr<framework::BlockDesc *>(RecurrentBase::kStepBlock);
    std::unordered_set<std::string> fwd_skip_vars;
    for (auto *op_desc : grad_block->AllOps()) {
      for (auto &in_arg_name : op_desc->InputArgumentNames()) {
        if (IsSkippableVar(in_arg_name, grad_block)) {
          fwd_skip_vars.insert(in_arg_name);
        }
      }
      for (auto &out_arg_name : op_desc->OutputArgumentNames()) {
        if (IsSkippableVar(out_arg_name, grad_block)) {
          fwd_skip_vars.insert(out_arg_name);
        }
      }
    }
    AddSkipVars(fwd_op, fwd_skip_vars);
  
    // Find all skippable variables in recurrent_grad_op
    // The skippable variables are those which would be used across time steps
    ClearSkipVars(bwd_op);
    AddOpMemVarsAsSkip(bwd_op, /* set_grad_mem_vars = */ true);
    std::unordered_set<std::string> bwd_skip_vars;
  
    auto &fwd_input = fwd_op.Inputs().at(RecurrentBase::kInputs);
    auto &in_grads =
        bwd_op.Outputs().at(framework::GradVarName(RecurrentBase::kInputs));
  
    PADDLE_ENFORCE_EQ(
        fwd_input.size(), in_grads.size(),
        "Backward input gradient number does not match forward input number.");
    for (size_t i = 0; i < in_grads.size(); ++i) {
      if (in_grads[i] == framework::kEmptyVarName) {
        continue;
      }
      bwd_skip_vars.insert(in_grads[i]);
      bwd_skip_vars.insert(framework::GradVarName(fwd_input[i]));
    }
  
    auto &fwd_param = fwd_op.Inputs().at(RecurrentBase::kParameters);
    auto &param_grads =
        bwd_op.Outputs().at(framework::GradVarName(RecurrentBase::kParameters));
    PADDLE_ENFORCE_EQ(fwd_param.size(), param_grads.size(),
                      "Backward parameter gradient number does not match forward "
                      "parameter number.");
    for (size_t i = 0; i < fwd_param.size(); ++i) {
      if (param_grads[i] == framework::kEmptyVarName) {
        continue;
      }
      bwd_skip_vars.insert(param_grads[i]);
      bwd_skip_vars.insert(framework::GradVarName(fwd_param[i]));
    }
  
    AddSkipVars(bwd_op, bwd_skip_vars);
  }
  
  void PrepareSafeEagerDeletionOnRecurrentOpAndRecurrentGradOp(
      const framework::ProgramDesc &program, int block_id,
      const std::vector<std::unique_ptr<paddle::framework::OperatorBase>>
          &all_ops) {
    // If block_id is not 0, returns
    // This is because all recurrent_ops and recurrent_grad_ops in the whole
    // program would be processed when block_id is 0 (i.e. when Executor::Run()
    // or ParallelExecutor constructs).
  
    // What's more, all recurrent_ops and recurrent_grad_ops must be processed
    // when block_id is zero. If not, recurrent_op may run first and erase
    // variables
    // used in recurrent_grad_op, and in this moment, recurrent_grad_ops may be
    // not constructed yet.
    if (block_id != 0) return;
  
    OpAndGradOpPair op_pair;
    for (auto &op : all_ops) {
      if (op->Type() == "recurrent") {
        op_pair.first.emplace(op.get());
      } else if (op->Type() == "recurrent_grad") {
        op_pair.second.emplace(op.get());
      }
    }
    PrepareSafeEagerDeletionOnRecurrentOpAndRecurrentGradOp(program, &op_pair);
  }
  
  void PrepareSafeEagerDeletionOnRecurrentOpAndRecurrentGradOp(
      const framework::ProgramDesc &program, OpAndGradOpPair *op_pair) {
    // Find all ops and grad ops at all blocks
    FindAllOpAndGradOp(program, op_pair, "recurrent", "recurrent_grad");
  
    OpVariantSet &recurrent_ops = op_pair->first;
    OpVariantSet &recurrent_grad_ops = op_pair->second;
  
    VLOG(2) << "Found recurrent op num: " << recurrent_ops.size()
            << ", recurrent grad op num: " << recurrent_grad_ops.size();
  
    if (recurrent_ops.empty()) {
      return;
    }
  
    for (auto &bwd_op : recurrent_grad_ops) {
      const OpVariant *matched_fwd_op = nullptr;
      for (auto &fwd_op : recurrent_ops) {
        if (IsMatchedRecurrentOpAndRecurrentGradOp(fwd_op, bwd_op)) {
          PADDLE_ENFORCE(matched_fwd_op == nullptr,
                         "Found multiple matched recurrent op");
          matched_fwd_op = &fwd_op;
        }
      }
      PADDLE_ENFORCE_NOT_NULL(matched_fwd_op, "Cannot find matched forward op");
      SetRecurrentOpAndRecurrentGradOpSkipVarAttr(*matched_fwd_op, bwd_op);
      recurrent_ops.erase(*matched_fwd_op);
    }
  
    for (auto &fwd_op : recurrent_ops) {
      SetRecurrentForwardOpOnlySkipVarAttr(fwd_op);
    }
  }
  
  }  // namespace operators
  }  // namespace paddle
  