xref: /aosp_15_r20/external/ComputeLibrary/tests/TensorCache.h (revision c217d954acce2dbc11938adb493fc0abd69584f3)
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