kernel_pool.h 3.4 KB
Newer Older
T
tensor-tang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2018 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

T
tensor-tang 已提交
17
#include <memory>  // for unique_ptr
T
tensor-tang 已提交
18 19
#include <string>
#include <unordered_map>
T
tensor-tang 已提交
20
#include <utility>  // for move
T
tensor-tang 已提交
21
#include <vector>
T
tensor-tang 已提交
22 23 24
#include "paddle/fluid/operators/jit/gen_base.h"
#include "paddle/fluid/operators/jit/kernel_base.h"
#include "paddle/fluid/operators/jit/kernel_key.h"
T
tensor-tang 已提交
25
#include "paddle/fluid/platform/place.h"
T
tensor-tang 已提交
26 27 28

namespace paddle {
namespace operators {
T
tensor-tang 已提交
29
namespace jit {
T
tensor-tang 已提交
30 31 32

template <KernelType KT>
class JitCodePool {
T
tensor-tang 已提交
33
  typedef std::unique_ptr<GenBase> GenBasePtr;
34
  typedef std::unordered_map<int64_t, GenBasePtr> JitCodeMap;
T
tensor-tang 已提交
35

T
tensor-tang 已提交
36
 public:
T
tensor-tang 已提交
37
  JitCodePool() = default;
T
tensor-tang 已提交
38 39 40 41 42
  static JitCodePool& Instance() {
    static thread_local JitCodePool<KT> g_jit_codes;
    return g_jit_codes;
  }

T
tensor-tang 已提交
43
  const JitCodeMap& AllKernels() { return codes_; }
T
tensor-tang 已提交
44

45
  bool Has(int64_t key) const { return codes_.find(key) != codes_.end(); }
T
tensor-tang 已提交
46

47
  void Insert(int64_t key, GenBasePtr value) {
T
tensor-tang 已提交
48
    codes_.emplace(key, std::move(value));
T
tensor-tang 已提交
49 50 51
  }

 private:
T
tensor-tang 已提交
52
  JitCodeMap codes_;
T
tensor-tang 已提交
53 54 55
  DISABLE_COPY_AND_ASSIGN(JitCodePool);
};

T
tensor-tang 已提交
56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
class JitCodeCreatorPool {
  typedef std::unique_ptr<const GenCreator> GenCreatorPtr;
  typedef std::unordered_map<KernelKey, std::vector<GenCreatorPtr>,
                             KernelKey::Hash>
      GenCreatorPtrMap;

 public:
  JitCodeCreatorPool() = default;
  static JitCodeCreatorPool& Instance();
  GenCreatorPtrMap& AllCreators() { return creators_; }
  void Insert(const KernelKey& key, GenCreatorPtr value) {
    if (creators_.find(key) == creators_.end()) {
      creators_.emplace(key, std::vector<GenCreatorPtr>());
    }
    creators_.at(key).emplace_back(std::move(value));
  }

 private:
  GenCreatorPtrMap creators_;
  DISABLE_COPY_AND_ASSIGN(JitCodeCreatorPool);
};

T
tensor-tang 已提交
78 79 80 81
typedef std::unique_ptr<const Kernel> KernelPtr;
typedef std::unordered_map<KernelKey, std::vector<KernelPtr>, KernelKey::Hash>
    KernelMap;

T
tensor-tang 已提交
82 83 84
class KernelPool {
 public:
  static KernelPool& Instance();
T
tensor-tang 已提交
85
  KernelPool() = default;
T
tensor-tang 已提交
86 87 88 89 90 91 92 93 94 95 96 97 98
  KernelMap& AllKernels() { return pool_; }
  void Insert(const KernelKey& key, KernelPtr value) {
    if (pool_.find(key) == pool_.end()) {
      pool_.emplace(key, std::vector<KernelPtr>());
    }
    pool_.at(key).emplace_back(std::move(value));
  }

 private:
  KernelMap pool_;
  DISABLE_COPY_AND_ASSIGN(KernelPool);
};

T
tensor-tang 已提交
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
// Every kernel should have refer code and it should be used in unit tests,
// so refer kernels should have it's independent kernel pool
class ReferKernelPool {
 public:
  static ReferKernelPool& Instance();
  ReferKernelPool() = default;
  KernelMap& AllKernels() { return pool_; }
  void Insert(const KernelKey& key, KernelPtr value) {
    if (pool_.find(key) == pool_.end()) {
      pool_.emplace(key, std::vector<KernelPtr>());
    }
    pool_.at(key).emplace_back(std::move(value));
  }

 private:
  KernelMap pool_;
  DISABLE_COPY_AND_ASSIGN(ReferKernelPool);
};

T
tensor-tang 已提交
118
}  // namespace jit
T
tensor-tang 已提交
119 120
}  // namespace operators
}  // namespace paddle