amp_auto_cast.h 3.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
// Copyright (c) 2020 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 <memory>
#include <set>
#include <string>
#include <tuple>
#include <unordered_set>

L
Leo Chen 已提交
22
#include "paddle/fluid/framework/op_registry.h"
23 24 25 26 27
#include "paddle/fluid/imperative/type_defs.h"

namespace paddle {
namespace imperative {

L
Leo Chen 已提交
28 29 30 31 32 33 34 35
// NOTE(zhiqiu): only O1 and O2 are valid now
enum class AmpLevel {
  O0 = 0,  // fp32
  O1,      // amp, mixed fp32-fp16
  O2,      // almost fp16
  O3,      // fp16
};

L
Leo Chen 已提交
36 37 38 39 40
std::tuple<std::unordered_set<std::string>, std::unordered_set<std::string>,
           std::unordered_set<std::string>>
OpSupportedInfos(const std::string& place,
                 framework::proto::VarType::Type dtype);

W
wanghuancoder 已提交
41 42
class Tracer;

L
Leo Chen 已提交
43
// Singleton implementation with C++ 11
44 45 46 47 48 49 50 51
class AmpOperators {
 public:
  ~AmpOperators();
  AmpOperators(const AmpOperators& o) = delete;
  const AmpOperators& operator=(const AmpOperators& o) = delete;

  static AmpOperators& Instance();

52
  std::shared_ptr<std::unordered_set<std::string>> GetMutableAllowOps();
53

54
  std::shared_ptr<std::unordered_set<std::string>> GetMutableBlockOps();
55

56 57 58
  std::shared_ptr<std::unordered_set<std::string>>
  GetMutableUnsupportedFp16Ops();

59 60 61 62 63 64 65 66 67 68
 private:
  AmpOperators();  // forbid calling default constructor

  // The set of ops that support fp16 calculation and are considered numerically
  // safe and performance critical. These ops are always converted to fp16.
  std::shared_ptr<std::unordered_set<std::string>> allow_ops_;

  // The set of ops that support fp16 calculation and are considered numerically
  // dangerous and whose effects may also be observed in downstream ops.
  std::shared_ptr<std::unordered_set<std::string>> block_ops_;
69 70 71

  // The set of ops that has no fp16 CUDA kennel.
  std::shared_ptr<std::unordered_set<std::string>> unsupported_fp16_ops_;
72 73
};

74 75
std::ostream& operator<<(std::ostream& os, AmpOperators& ops);

76 77 78
// NOTE(zhiqiu): AutoCastGuard is used for RAII.
class AutoCastGuard {
 public:
L
Leo Chen 已提交
79
  AutoCastGuard(std::shared_ptr<Tracer> tracer, AmpLevel guard_level);
80

L
Leo Chen 已提交
81
  ~AutoCastGuard();
82 83 84 85 86 87 88

  // forbid copy and operator=
  AutoCastGuard(const AutoCastGuard& guard) = delete;
  AutoCastGuard& operator=(const AutoCastGuard& guard) = delete;

 private:
  std::shared_ptr<Tracer> tracer_;
L
Leo Chen 已提交
89
  AmpLevel pre_amp_level_;
90 91
};

J
Jiabin Yang 已提交
92 93 94 95 96 97
template <typename VarType>
NameVarMap<VarType> AutoCastInputs(const std::string& op_type,
                                   const NameVarMap<VarType>& ins);
template <typename VarType>
NameVarMap<VarType> CastPureFp16Inputs(const std::string& op_type,
                                       const NameVarMap<VarType>& ins);
98

99 100
}  // namespace imperative
}  // namespace paddle