1 /*
2 * Copyright (c) 2017-2019 Arm Limited.
3 *
4 * SPDX-License-Identifier: MIT
5 *
6 * Permission is hereby granted, free of charge, to any person obtaining a copy
7 * of this software and associated documentation files (the "Software"), to
8 * deal in the Software without restriction, including without limitation the
9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10 * sell copies of the Software, and to permit persons to whom the Software is
11 * furnished to do so, subject to the following conditions:
12 *
13 * The above copyright notice and this permission notice shall be included in all
14 * copies or substantial portions of the Software.
15 *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22 * SOFTWARE.
23 */
24 #ifndef ARM_COMPUTE_TEST_TENSOR_CACHE_H
25 #define ARM_COMPUTE_TEST_TENSOR_CACHE_H
26
27 #include "RawTensor.h"
28
29 #include "support/Mutex.h"
30
31 #include <map>
32 #include <mutex>
33 #include <utility>
34
35 namespace arm_compute
36 {
37 namespace test
38 {
39 /** Stores @ref RawTensor categorised by the image they are created from
40 * including name, format and channel.
41 */
42 class TensorCache
43 {
44 public:
45 /** Search the cache for a tensor of created from the specified image and
46 * format.
47 *
48 * @param[in] key Key to look up the tensor. Consists of image name and format.
49 *
50 * @return The cached tensor matching the image name and format if found. A
51 * nullptr otherwise.
52 */
53 RawTensor *find(std::tuple<const std::string &, Format> key);
54
55 /** Search the cache for a tensor of created from the specified image,
56 * format and channel.
57 *
58 * @param[in] key Key to look up the tensor. Consists of image name, format and channel.
59 *
60 * @return The cached tensor matching the image name and format if found. A
61 * nullptr otherwise.
62 */
63 RawTensor *find(std::tuple<const std::string &, Format, Channel> key);
64
65 /** Add the given tensor to the cache. Can later be found under the given
66 * image name and format.
67 *
68 * @param[in] key Key under which to store the tensor. Consists of image name and format.
69 * @param[in] raw Raw tensor to be stored.
70 *
71 * @return A reference to the cached tensor.
72 */
73 RawTensor &add(std::tuple<const std::string &, Format> key, RawTensor raw);
74
75 /** Add the given tensor to the cache. Can later be found under the given
76 * image name, format and channel.
77 *
78 * @param[in] key Key under which to store the tensor. Consists of image name, format and channel.
79 * @param[in] raw Raw tensor to be stored.
80 *
81 * @return A reference to the cached tensor.
82 */
83 RawTensor &add(std::tuple<const std::string &, Format, Channel> key, RawTensor raw);
84
85 private:
86 using FormatMap = std::map<std::tuple<std::string, Format>, RawTensor>;
87 using ChannelMap = std::map<std::tuple<std::string, Format, Channel>, RawTensor>;
88
89 FormatMap _raw_tensor_cache{};
90 ChannelMap _raw_tensor_channel_cache{};
91 arm_compute::Mutex _raw_tensor_cache_mutex{};
92 arm_compute::Mutex _raw_tensor_channel_cache_mutex{};
93 };
94
find(std::tuple<const std::string &,Format> key)95 inline RawTensor *TensorCache::find(std::tuple<const std::string &, Format> key)
96 {
97 const auto it = _raw_tensor_cache.find(key);
98 return it == _raw_tensor_cache.end() ? nullptr : &it->second;
99 }
100
find(std::tuple<const std::string &,Format,Channel> key)101 inline RawTensor *TensorCache::find(std::tuple<const std::string &, Format, Channel> key)
102 {
103 const auto it = _raw_tensor_channel_cache.find(key);
104 return it == _raw_tensor_channel_cache.end() ? nullptr : &it->second;
105 }
106
add(std::tuple<const std::string &,Format> key,RawTensor raw)107 inline RawTensor &TensorCache::add(std::tuple<const std::string &, Format> key, RawTensor raw)
108 {
109 std::lock_guard<arm_compute::Mutex> lock(_raw_tensor_cache_mutex);
110 return std::get<0>(_raw_tensor_cache.emplace(std::move(key), std::move(raw)))->second;
111 }
112
add(std::tuple<const std::string &,Format,Channel> key,RawTensor raw)113 inline RawTensor &TensorCache::add(std::tuple<const std::string &, Format, Channel> key, RawTensor raw)
114 {
115 std::lock_guard<arm_compute::Mutex> lock(_raw_tensor_channel_cache_mutex);
116 return std::get<0>(_raw_tensor_channel_cache.emplace(std::move(key), std::move(raw)))->second;
117 }
118 } // namespace test
119 } // namespace arm_compute
120 #endif /* ARM_COMPUTE_TEST_TENSOR_CACHE_H */
121