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 | |
28 | namespace glow { |
29 | |
30 | class TensorPool final { |
31 | private: |
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 | |
48 | public: |
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 | |