static_global_utils.h 2.8 KB
Newer Older
J
Jiabin Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58
// Copyright (c) 2022 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.

#pragma once
#include <algorithm>
#include <atomic>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "paddle/fluid/framework/op_call_stack.h"
#include "paddle/fluid/framework/op_desc.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/type_defs.h"

namespace paddle {
namespace prim {
class UniqueNameGenerator {
 public:
  explicit UniqueNameGenerator(std::string prefix = "") : prefix_(prefix) {}
  std::string Generate(std::string key = "") {
    return prefix_ + key + "_" + std::to_string(id_++);
  }

 private:
  std::atomic<int> id_{0};
  std::string prefix_;
};

class StaticCompositeContext {
 public:
  static StaticCompositeContext& Instance() {
    return *static_composite_context_;
  }

  framework::BlockDesc* GetBlock() { return current_block_desc_; }

  void SetBlock(framework::BlockDesc* new_block) {
    current_block_desc_ = new_block;
  }

  std::string GenerateUniqueName(std::string key = "composite_tmp") {
    return generator_->Generate(key);
  }

59
  void SetBwdPrimEnabled(bool enable_prim) { enable_bwd_prim_ = enable_prim; }
J
Jiabin Yang 已提交
60

61 62 63 64 65 66 67 68 69 70
  bool IsBwdPrimEnabled() { return enable_bwd_prim_; }

  void SetFwdPrimEnabled(bool enable_prim) { enable_fwd_prim_ = enable_prim; }

  bool IsFwdPrimEnabled() { return enable_fwd_prim_; }

  void SetAllPrimEnabled(bool enable_prim) {
    enable_fwd_prim_ = enable_prim;
    enable_bwd_prim_ = enable_prim;
  }
J
Jiabin Yang 已提交
71

72 73 74 75 76 77 78 79
  void SetTargetGradName(const std::map<std::string, std::string>& m) {
    target_grad_name_ = m;
  }

  std::map<std::string, std::string> GetTargetGradName() {
    return target_grad_name_;
  }

J
Jiabin Yang 已提交
80 81 82 83 84 85
 private:
  StaticCompositeContext()
      : current_block_desc_(nullptr), generator_(new UniqueNameGenerator()) {}

  framework::BlockDesc* current_block_desc_;
  std::unique_ptr<UniqueNameGenerator> generator_;
86
  std::map<std::string, std::string> target_grad_name_;
87 88
  static thread_local bool enable_bwd_prim_;
  static thread_local bool enable_fwd_prim_;
J
Jiabin Yang 已提交
89 90 91 92 93 94
  static StaticCompositeContext* static_composite_context_;
  DISABLE_COPY_AND_ASSIGN(StaticCompositeContext);
};

}  // namespace prim
}  // namespace paddle