1/**
2 * Copyright (c) Glow Contributors. See CONTRIBUTORS file.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16#ifndef GLOW_TENSORPOOL_H
17#define GLOW_TENSORPOOL_H
18
19#include "glow/Base/Tensor.h"
20#include "llvm/ADT/Optional.h"
21
22#include <atomic>
23#include <iostream>
24#include <mutex>
25#include <unordered_map>
26#include <vector>
27
28namespace glow {
29
30class TensorPool final {
31private:
32 struct TypeHash {
33 size_t operator()(const Type &t) const { return t.equals_hash(); }
34 };
35
36 struct TypeEquals {
37 bool operator()(const Type &a, const Type &b) const { return a.isEqual(b); }
38 };
39 /// A stack of available Tensors per Type.
40 std::unordered_map<Type, std::vector<Tensor>, TypeHash, TypeEquals> pools_;
41
42 /// Mutex around pools_;
43 std::mutex lock_;
44
45 /// Whether or not to allow allocation of new buffers if the pool is empty.
46 const bool preventInlineAllocs_{false};
47
48public:
49 /// Statistics relating to the usage of the pool.
50 struct Stats {
51 /// The total number of Types that has ever been available in this pool.
52 std::atomic<uint64_t> totalTypes{0};
53 /// The number of Tensors currently allocated and available.
54 std::atomic<uint64_t> currentBuffers{0};
55 /// The number of Tensor allocations ever done by the pool.
56 std::atomic<uint64_t> totalAllocs{0};
57 /// The number of Tensor allocations that were done inline to get (as
58 /// opposed to reserve).
59 std::atomic<uint64_t> inlineAllocs{0};
60 /// The total number of times a Tensor was retrieved from the pool.
61 std::atomic<uint64_t> totalGets{0};
62 /// The total number of times a Tensor was returned to the pool.
63 std::atomic<uint64_t> totalReclaims{0};
64 /// The total number of times a Tensor was freed (e.g. via clear()).
65 std::atomic<uint64_t> totalFrees{0};
66 } stats_;
67
68 TensorPool(bool preventAllocs = false)
69 : preventInlineAllocs_{preventAllocs} {}
70
71 ~TensorPool() { clear(); }
72
73 /// Retrieve a Tensor with type \p ty from the pool - this type must have
74 /// previously been added by initialize. If the pool is empty this will
75 /// allocate a new Tensor unless preventAllocs was set true at construction
76 /// time.
77 llvm::Optional<Tensor> get(TypeRef ty);
78
79 /// Return a Tensor \p t to the pool. This Tensor must have been previously
80 /// allocated by this TensorPool.
81 void reclaim(Tensor &&t);
82
83 /// Add \p count elements of the provided type \p ty to the pool.
84 void reserve(TypeRef ty, size_t count);
85
86 /// Clear the pool and all allocated Tensors.
87 /// Note: this does not delete tensors that were allocated by the pool but
88 /// were not reclaimed.
89 void clear();
90
91 /// Return statistics about the TensorPool.
92 const Stats &getStats() { return stats_; }
93};
94} // namespace glow
95
96#endif
97