<commit_before>//   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/framework/details/fuse_optimizer_op_pass.h"
#include <algorithm>
#include <unordered_set>
#include "paddle/fluid/framework/ir/graph_helper.h"
#include "paddle/fluid/framework/op_registry.h"

namespace paddle {
namespace framework {
namespace details {

void FuseOptimizerOpPass::ApplyImpl(ir::Graph *graph) const {
  ir::Graph &result = *graph;

  auto &places = Get<const std::vector<platform::Place>>(kPlaces);
  auto &local_scopes = Get<const std::vector<Scope *>>(kLocalScopes);

  const std::string fuse_op_type = GetOpType();
  std::vector<std::string> aux_var_names = GetAuxiliaryVarNames();
  aux_var_names.emplace_back(kParam);
  aux_var_names.emplace_back(kGrad);

  // Step 1: Get the specified op and auxiliary variables.
  std::vector<ir::Node *> topo_nodes = ir::TopologySortOperations(result);
  std::unordered_map<std::string, std::vector<std::string>> aux_var_set;
  std::vector<ir::Node *> opt_ops;
  for (auto &node : topo_nodes) {
    GetSpecifiedOpsAndVars(fuse_op_type, aux_var_names, node, &opt_ops,
                           &aux_var_set);
  }

  VLOG(6) << "Find " << fuse_op_type << " operators: " << opt_ops.size();
  if (opt_ops.size() == 0) {
    return;
  }

  if (result.Has(kFusedOptType)) {
    VLOG(6) << "Currently only support fusing one type optimizer op. Has fused "
            << result.Get<FusedOptType>(kFusedOptType);
    return;
  } else {
    result.Set(kFusedOptType, new FusedOptType);
  }
  result.Get<FusedOptType>(kFusedOptType) = fuse_op_type;

  // Step 2: Insert fused_var_name to FusedVars, and the FusedVars need be
  // initialized in scopes before execution.
  if (!result.Has(kFusedVars)) {
    result.Set(kFusedVars, new FusedVars);
  }
  std::unordered_map<std::string, std::string> fused_vars_name;
  fused_vars_name.reserve(aux_var_names.size());
  auto &fused_var_set = result.Get<FusedVars>(kFusedVars);
  const std::string prefix(kFusedVarNamePrefix);
  // NOTE: the fused_var_name should be unique.
  for (auto &var_name : aux_var_names) {
    auto fused_var_name = prefix + "_" + fuse_op_type + "_" + var_name + "_" +
                          aux_var_set[var_name][0];
    VLOG(6) << var_name << ": " << fused_var_name;
    fused_vars_name.emplace(var_name, fused_var_name);
    PADDLE_ENFORCE_EQ(fused_var_set.count(fused_var_name), 0);
    fused_var_set.insert(fused_var_name);
  }

  // Step 3: Get the fused Gradient's name
  bool grad_fused = false;
  if (result.Has(kParamsAndGrads)) {
    auto &params_grads = result.Get<ParamsAndGrads>(kParamsAndGrads);
    PADDLE_ENFORCE_EQ(
        params_grads.size(), aux_var_set.at(kGrad).size(),
        "The number of gradients and optimizer ops is not equal.");
    std::unordered_set<std::string> opt_grad_set(aux_var_set.at(kGrad).begin(),
                                                 aux_var_set.at(kGrad).end());
    size_t same_grad_num = 0;
    for (auto &p_g : params_grads) {
      if (opt_grad_set.count(p_g.second)) {
        ++same_grad_num;
      }
    }

    // NOTE(zcd): the gradient of kParamsAndGrads may be different with the
    // kGrad.
    if (same_grad_num == aux_var_set.at(kGrad).size()) {
      if (!result.Has(kFusedGrads)) {
        PADDLE_THROW(
            "The alloc_continuous_space_for_grad_pass should be called before "
            "this pass.");
      }
      auto &fused_grad = result.Get<FusedGrads>(kFusedGrads);
      auto &fused_vars = result.Get<FusedVars>(kFusedVars);
      auto iter = std::find(fused_vars.begin(), fused_vars.end(), fused_grad);
      PADDLE_ENFORCE(iter != fused_vars.end(), "Not find the fused_grad.");
      fused_vars_name[kGrad] = fused_grad;

      // Sort the parameters and auxiliary variables according
      // to parameters' name to make variables' name correspond correctly.
      SortParametersAndAuxVars(params_grads, &aux_var_set, &opt_ops);
      grad_fused = true;
    }
  }

  // Step 4: Alloc continuous space for Parameters and AuxiliaryVar(e.g.
  // Moment1, Moment2, Beta1Pow, Beta2Pow) of all the optimizer ops separately.
  aux_var_names.pop_back();
  if (!grad_fused) {
    InitFusedGradsAndAllocSpaceForGrads(
        places, local_scopes, aux_var_set.at(kParam), aux_var_set.at(kGrad),
        fused_vars_name.at(kGrad), &result);
  }
  InitFusedVarsAndAllocSpaceForVars(places, local_scopes, aux_var_names,
                                    aux_var_set, fused_vars_name);

  // Step 5: Fuse optimizer Ops and Scale Ops
  FuseOptimizerOps(aux_var_set, fused_vars_name, opt_ops, &result);

  // Step 6: Remove optimizer Ops
  for (auto &opt_op : opt_ops) {
    graph->RemoveNode(opt_op);
  }
}

void FuseOptimizerOpPass::InitFusedGradsAndAllocSpaceForGrads(
    const std::vector<platform::Place> &places,
    const std::vector<Scope *> &local_scopes,
    const std::vector<std::string> &params,
    const std::vector<std::string> &grads, const std::string &fused_grad_name,
    ir::Graph *result) const {
  // Get Var Nodes
  std::unordered_map<std::string, ir::Node *> vars;
  for (ir::Node *node : result->Nodes()) {
    if (node->IsVar() && node->Var()) {
      // Note: The graph may have the same name node. For example, parameter
      // is the input of operator and it also is the output of optimizer;
      vars.emplace(node->Var()->Name(), node);
    }
  }
  // Init Grads
  for (auto it = local_scopes.rbegin(); it != local_scopes.rend(); ++it) {
    auto &scope = *it;
    VLOG(6) << "Init: " << fused_grad_name;
    PADDLE_ENFORCE(scope->FindVar(fused_grad_name) == nullptr,
                   "%s has existed in scope.", fused_grad_name);
    scope->Var(fused_grad_name)->GetMutable<LoDTensor>();

    for (auto &grad_var_name : grads) {
      auto iter = vars.find(grad_var_name);
      PADDLE_ENFORCE(iter != vars.end());
      PADDLE_ENFORCE_NOT_NULL(iter->second->Var());
      PADDLE_ENFORCE_EQ(iter->second->Var()->GetType(),
                        proto::VarType::LOD_TENSOR);
      scope->Var(grad_var_name)->GetMutable<LoDTensor>();
    }
  }
  // Define Ops
  ProgramDesc program_desc;
  auto *global_block = program_desc.MutableBlock(0);
  AppendAllocContinuousSpace(params, grads, fused_grad_name, global_block,
                             false, false);
  // Run Ops
  RunInitOps(places, local_scopes, *global_block);
}<commit_msg>fix fuse optimizer ops (#17102)<commit_after>//   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/framework/details/fuse_optimizer_op_pass.h"
#include <algorithm>
#include <unordered_set>
#include "paddle/fluid/framework/ir/graph_helper.h"
#include "paddle/fluid/framework/op_registry.h"

namespace paddle {
namespace framework {
namespace details {

void FuseOptimizerOpPass::ApplyImpl(ir::Graph *graph) const {
  ir::Graph &result = *graph;

  auto &places = Get<const std::vector<platform::Place>>(kPlaces);
  auto &local_scopes = Get<const std::vector<Scope *>>(kLocalScopes);

  const std::string fuse_op_type = GetOpType();
  std::vector<std::string> aux_var_names = GetAuxiliaryVarNames();
  aux_var_names.emplace_back(kParam);
  aux_var_names.emplace_back(kGrad);

  // Step 1: Get the specified op and auxiliary variables.
  std::vector<ir::Node *> topo_nodes = ir::TopologySortOperations(result);
  std::unordered_map<std::string, std::vector<std::string>> aux_var_set;
  std::vector<ir::Node *> opt_ops;
  for (auto &node : topo_nodes) {
    GetSpecifiedOpsAndVars(fuse_op_type, aux_var_names, node, &opt_ops,
                           &aux_var_set);
  }

  VLOG(6) << "Find " << fuse_op_type << " operators: " << opt_ops.size();
  if (opt_ops.size() == 0) {
    return;
  }

  if (result.Has(kFusedOptType)) {
    VLOG(6) << "Currently only support fusing one type optimizer op. Has fused "
            << result.Get<FusedOptType>(kFusedOptType);
    return;
  } else {
    result.Set(kFusedOptType, new FusedOptType);
  }
  result.Get<FusedOptType>(kFusedOptType) = fuse_op_type;

  // Step 2: Insert fused_var_name to FusedVars, and the FusedVars need be
  // initialized in scopes before execution.
  if (!result.Has(kFusedVars)) {
    result.Set(kFusedVars, new FusedVars);
  }
  std::unordered_map<std::string, std::string> fused_vars_name;
  fused_vars_name.reserve(aux_var_names.size());
  auto &fused_var_set = result.Get<FusedVars>(kFusedVars);
  const std::string prefix(kFusedVarNamePrefix);
  // NOTE: the fused_var_name should be unique.
  for (auto &var_name : aux_var_names) {
    auto fused_var_name = prefix + "_" + fuse_op_type + "_" + var_name + "_" +
                          aux_var_set[var_name][0];
    VLOG(6) << var_name << ": " << fused_var_name;
    fused_vars_name.emplace(var_name, fused_var_name);
    PADDLE_ENFORCE_EQ(fused_var_set.count(fused_var_name), 0);
    fused_var_set.insert(fused_var_name);
  }

  // Step 3: Get the fused Gradient's name
  bool grad_fused = false;
  if (result.Has(kParamsAndGrads)) {
    auto &params_grads = result.Get<ParamsAndGrads>(kParamsAndGrads);
    PADDLE_ENFORCE_EQ(
        params_grads.size(), aux_var_set.at(kGrad).size(),
        "The number of gradients and optimizer ops is not equal.");
    std::unordered_set<std::string> opt_grad_set(aux_var_set.at(kGrad).begin(),
                                                 aux_var_set.at(kGrad).end());
    size_t same_grad_num = 0;
    for (auto &p_g : params_grads) {
      if (opt_grad_set.count(p_g.second)) {
        ++same_grad_num;
      }
    }

    // NOTE(zcd): the gradient of kParamsAndGrads may be different with the
    // kGrad.
    if (same_grad_num == aux_var_set.at(kGrad).size()) {
      if (!result.Has(kFusedGrads)) {
        PADDLE_THROW(
            "The alloc_continuous_space_for_grad_pass should be called before "
            "this pass.");
      }
      auto &fused_grad = result.Get<FusedGrads>(kFusedGrads);
      auto &fused_vars = result.Get<FusedVars>(kFusedVars);
      auto iter = std::find(fused_vars.begin(), fused_vars.end(), fused_grad);
      PADDLE_ENFORCE(iter != fused_vars.end(), "Not find the fused_grad.");
      fused_vars_name[kGrad] = fused_grad;

      // Sort the parameters and auxiliary variables according
      // to parameters' name to make variables' name correspond correctly.
      SortParametersAndAuxVars(params_grads, &aux_var_set, &opt_ops);
      grad_fused = true;
    }
  }

  // Step 4: Alloc continuous space for Parameters and AuxiliaryVar(e.g.
  // Moment1, Moment2, Beta1Pow, Beta2Pow) of all the optimizer ops separately.
  aux_var_names.pop_back();
  if (!grad_fused) {
    InitFusedGradsAndAllocSpaceForGrads(
        places, local_scopes, aux_var_set.at(kParam), aux_var_set.at(kGrad),
        fused_vars_name.at(kGrad), &result);
  }
  InitFusedVarsAndAllocSpaceForVars(places, local_scopes, aux_var_names,
                                    aux_var_set, fused_vars_name);

  // Step 5: Fuse optimizer Ops and Scale Ops
  FuseOptimizerOps(aux_var_set, fused_vars_name, opt_ops, &result);

  // Step 6: Remove optimizer Ops
  for (auto &opt_op : opt_ops) {
    graph->RemoveNode(opt_op);
  }
}

void FuseOptimizerOpPass::InitFusedGradsAndAllocSpaceForGrads(
    const std::vector<platform::Place> &places,
    const std::vector<Scope *> &local_scopes,
    const std::vector<std::string> &params,
    const std::vector<std::string> &grads, const std::string &fused_grad_name,
    ir::Graph *result) const {
  // Get Var Nodes
  std::unordered_map<std::string, ir::Node *> vars;
  for (ir::Node *node : result->Nodes()) {
    if (node->IsVar() && node->Var()) {
      // Note: The graph may have the same name node. For example, parameter
      // is the input of operator and it also is the output of optimizer;
      vars.emplace(node->Var()->Name(), node);
    }
  }

  // Set Gradients as Persistable to prevent this var becoming reusable.
  for (auto &grad_var_name : grads) {
    auto iter = vars.find(grad_var_name);
    PADDLE_ENFORCE(iter != vars.end());
    PADDLE_ENFORCE_NOT_NULL(iter->second->Var());
    PADDLE_ENFORCE(iter->second->Var()->GetType() == proto::VarType::LOD_TENSOR,
                   "Currently the gradient type only should be LoDTensor when "
                   "fusing optimizer ops.");
    iter->second->Var()->SetPersistable(true);
  }

  // Init Grads
  for (auto it = local_scopes.rbegin(); it != local_scopes.rend(); ++it) {
    auto &scope = *it;
    VLOG(6) << "Init: " << fused_grad_name;
    PADDLE_ENFORCE(scope->FindVar(fused_grad_name) == nullptr,
                   "%s has existed in scope.", fused_grad_name);
    scope->Var(fused_grad_name)->GetMutable<LoDTensor>();
    for (auto &grad_var_name : grads) {
      auto iter = vars.find(grad_var_name);
      PADDLE_ENFORCE(iter != vars.end());
      PADDLE_ENFORCE_NOT_NULL(iter->second->Var());
      scope->Var(grad_var_name)->GetMutable<LoDTensor>();
    }
  }
  // Define Ops
  ProgramDesc program_desc;
  auto *global_block = program_desc.MutableBlock(0);
  AppendAllocContinuousSpace(params, grads, fused_grad_name, global_block,
                             false, false);
  // Run Ops
  RunInitOps(places, local_scopes, *global_block);
}
{'fraction_non_alphanumeric': 0.12571748565028698, 'fraction_numerical': 0.004129917401651967, 'mean_word_length': 3.547103755569701, 'pattern_counts': {'":': 2, '<': 99, '<?xml version=': 0, '>': 92, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6984037', 'n_tokens_mistral': 4960, 'n_tokens_neox': 4826, 'n_words': 1069}