xref: /aosp_15_r20/external/ruy/ruy/prepacked_cache_test.cc (revision bb86c7ed5fb1b98a7eac808e443a46cc8b90dfc0)
1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include "ruy/prepacked_cache.h"
17 
18 #include <thread>  // NOLINT(build/c++11)
19 
20 #include "ruy/context.h"
21 #include "ruy/context_get_ctx.h"
22 #include "ruy/gtest_wrapper.h"
23 #include "ruy/mat.h"
24 #include "ruy/matrix.h"
25 #include "ruy/ruy.h"
26 #include "ruy/time.h"
27 
28 namespace ruy {
29 namespace {
30 
MakeDummyPEMat(Type data_type,int rows,int cols)31 PEMat MakeDummyPEMat(Type data_type, int rows, int cols) {
32   PEMat ret;
33   ret.data_type = data_type;
34   if (!data_type.is_floating_point) {
35     ret.sums_type = Type::Create<std::int32_t>();
36   }
37   ret.layout.rows = rows;
38   ret.layout.cols = cols;
39   ret.layout.stride = rows;
40   ret.layout.order = Order::kColMajor;
41   // The kernel block layout is not relevant to this test, so we leave it
42   // trivial 1x1.
43   ret.layout.kernel.rows = 1;
44   ret.layout.kernel.cols = 1;
45   return ret;
46 }
47 
48 template <typename T>
DummyPack(const std::vector<T> & data,PEMat * packed_matrix)49 void DummyPack(const std::vector<T>& data, PEMat* packed_matrix) {
50   EXPECT_EQ(data.size(), FlatSize(packed_matrix->layout));
51   memcpy(packed_matrix->data, data.data(), data.size() * sizeof(T));
52 }
53 
TEST(PrepackedCacheTest,TestCacheBasic)54 TEST(PrepackedCacheTest, TestCacheBasic) {
55   PrepackedCache prepacked_cache(307);
56   // Allocate the prepacked matrix.
57   // DataBytes=200, SumsBytes=20*4=80, Total: 280 bytes
58   std::vector<std::uint8_t> data1(10 * 20);
59   PEMat mat1 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
60   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
61               PrepackedCache::Action::kInsertedNewEntry);
62   DummyPack(data1, &mat1);
63 
64   // DataBytes=15, SumsBytes=3*4=12, Total: 27 bytes
65   std::vector<std::uint8_t> data2(5 * 3);
66   PEMat mat2 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 5, 3);
67   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
68               PrepackedCache::Action::kInsertedNewEntry);
69   DummyPack(data2, &mat2);
70 
71   // Both should now be in cache.
72   EXPECT_EQ(prepacked_cache.MatrixCount(), 2);
73   EXPECT_EQ(prepacked_cache.BuffersBytes(), 307);
74   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
75               PrepackedCache::Action::kGotExistingEntry);
76   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
77               PrepackedCache::Action::kGotExistingEntry);
78 }
79 
TEST(PrepackedCacheTest,TestCacheBasicFloat)80 TEST(PrepackedCacheTest, TestCacheBasicFloat) {
81   PrepackedCache prepacked_cache(860);
82   // Allocate the prepacked matrix.
83   // DataBytes=200*4, SumsBytes=0 because float, Total: 800 bytes
84   std::vector<float> data1(10 * 20);
85   PEMat mat1 = MakeDummyPEMat(Type::Create<float>(), 10, 20);
86   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
87               PrepackedCache::Action::kInsertedNewEntry);
88   DummyPack(data1, &mat1);
89 
90   // DataBytes=15*4, SumsBytes=0 because float, Total: 60 bytes
91   std::vector<float> data2(5 * 3);
92   PEMat mat2 = MakeDummyPEMat(Type::Create<float>(), 5, 3);
93   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
94               PrepackedCache::Action::kInsertedNewEntry);
95   DummyPack(data2, &mat2);
96 
97   // Both should now be in cache.
98   EXPECT_EQ(prepacked_cache.MatrixCount(), 2);
99   EXPECT_EQ(prepacked_cache.BuffersBytes(), 860);
100   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
101               PrepackedCache::Action::kGotExistingEntry);
102   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
103               PrepackedCache::Action::kGotExistingEntry);
104 }
105 
TEST(PrepackedCacheTest,TestCacheEjection)106 TEST(PrepackedCacheTest, TestCacheEjection) {
107   PrepackedCache prepacked_cache(306);
108   // Allocate the prepacked matrix.
109   // DataBytes=200, SumsBytes=20*4=80, Total: 280 bytes
110   std::vector<std::uint8_t> data1(10 * 20);
111   PEMat mat1 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
112   prepacked_cache.Get(data1.data(), &mat1);
113   DummyPack(data1, &mat1);
114 
115   // DataBytes=15, SumsBytes=3*4=12, Total: 27 bytes
116   std::vector<std::uint8_t> data2(5 * 3);
117   PEMat mat2 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 5, 3);
118   prepacked_cache.Get(data2.data(), &mat2);
119   DummyPack(data2, &mat2);
120 
121   // The first matrix should have been ejected from the cache.
122   // Only the second matrix should now be in cache.
123   EXPECT_EQ(prepacked_cache.MatrixCount(), 1);
124   EXPECT_EQ(prepacked_cache.BuffersBytes(), 27);
125   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
126               PrepackedCache::Action::kGotExistingEntry);
127   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
128               PrepackedCache::Action::kInsertedNewEntry);
129 
130   // The second matrix should have been ejected from the cache.
131   // Only the first matrix should now be in cache.
132   EXPECT_EQ(prepacked_cache.MatrixCount(), 1);
133   EXPECT_EQ(prepacked_cache.BuffersBytes(), 280);
134   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
135               PrepackedCache::Action::kGotExistingEntry);
136   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
137               PrepackedCache::Action::kInsertedNewEntry);
138 }
139 
TEST(PrepackedCacheTest,TestCacheEjection2)140 TEST(PrepackedCacheTest, TestCacheEjection2) {
141   PrepackedCache prepacked_cache(1000);
142   // Allocate the prepacked matrix 1.
143   // DataBytes=200, SumsBytes=20*4=80, Total: 280 bytes
144   std::vector<std::uint8_t> data1(10 * 20);
145   PEMat mat1 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
146   prepacked_cache.Get(data1.data(), &mat1);
147   DummyPack(data1, &mat1);
148 
149   // Allocate the prepacked matrix 2.
150   // DataBytes=200, SumsBytes=20*4=80, Total: 280 bytes
151   std::vector<std::uint8_t> data2(10 * 20);
152   PEMat mat2 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
153   prepacked_cache.Get(data2.data(), &mat2);
154   DummyPack(data2, &mat2);
155 
156   // Allocate the prepacked matrix 3.
157   // DataBytes=200, SumsBytes=20*4=80, Total: 280 bytes
158   std::vector<std::uint8_t> data3(10 * 20);
159   PEMat mat3 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
160   prepacked_cache.Get(data3.data(), &mat3);
161   DummyPack(data3, &mat3);
162 
163   // The next insertion will cause the cache size to go over the ejection
164   // threshold. Touch matrix 1 and matrix 3 to make matrix 2 the oldest
165   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
166               PrepackedCache::Action::kGotExistingEntry);
167   EXPECT_TRUE(prepacked_cache.Get(data3.data(), &mat3) ==
168               PrepackedCache::Action::kGotExistingEntry);
169 
170   // Allocate the prepacked matrix 4.
171   // DataBytes=200, SumsBytes=20*4=80, Total: 280 bytes
172   std::vector<std::uint8_t> data4(10 * 20);
173   PEMat mat4 = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
174   prepacked_cache.Get(data4.data(), &mat4);
175   DummyPack(data4, &mat4);
176 
177   // Ensure that mat2 was ejected, but mat1, mat3, and mat4 were not.
178   EXPECT_EQ(prepacked_cache.MatrixCount(), 3);
179   EXPECT_TRUE(prepacked_cache.Get(data1.data(), &mat1) ==
180               PrepackedCache::Action::kGotExistingEntry);
181   EXPECT_TRUE(prepacked_cache.Get(data3.data(), &mat3) ==
182               PrepackedCache::Action::kGotExistingEntry);
183   EXPECT_TRUE(prepacked_cache.Get(data4.data(), &mat4) ==
184               PrepackedCache::Action::kGotExistingEntry);
185   EXPECT_TRUE(prepacked_cache.Get(data2.data(), &mat2) ==
186               PrepackedCache::Action::kInsertedNewEntry);
187 }
188 
TEST(PrepackedCacheTest,TestDistinguishSubtlyDifferentMatrices)189 TEST(PrepackedCacheTest, TestDistinguishSubtlyDifferentMatrices) {
190   PrepackedCache prepacked_cache;
191 
192   std::vector<std::uint8_t> data(10 * 20);
193   PEMat mat = MakeDummyPEMat(Type::Create<std::uint8_t>(), 10, 20);
194   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
195             PrepackedCache::Action::kInsertedNewEntry);
196 
197   // Same layout, different source data pointer
198   EXPECT_EQ(prepacked_cache.Get(data.data() + 1, &mat),
199             PrepackedCache::Action::kInsertedNewEntry);
200 
201   // Layout tweaks
202   mat.layout.rows = 9;
203   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
204             PrepackedCache::Action::kInsertedNewEntry);
205 
206   mat.layout.cols = 19;
207   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
208             PrepackedCache::Action::kInsertedNewEntry);
209 
210   mat.layout.order = Order::kRowMajor;
211   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
212             PrepackedCache::Action::kInsertedNewEntry);
213 
214   mat.layout.kernel.rows = 2;
215   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
216             PrepackedCache::Action::kInsertedNewEntry);
217 
218   mat.layout.kernel.cols = 2;
219   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
220             PrepackedCache::Action::kInsertedNewEntry);
221 
222   mat.layout.kernel.order = Order::kRowMajor;
223   EXPECT_EQ(prepacked_cache.Get(data.data(), &mat),
224             PrepackedCache::Action::kInsertedNewEntry);
225 
226   EXPECT_EQ(prepacked_cache.MatrixCount(), 8);
227 }
228 
TestCachePolicies(CachePolicy cache_policy,bool expected_cached)229 void TestCachePolicies(CachePolicy cache_policy, bool expected_cached) {
230   ruy::Context context;
231   ruy::Ctx* ctx = get_ctx(&context);
232   PrepackedCache* cache = ctx->GetPrepackedCache();
233   EXPECT_EQ(cache->MatrixCount(), 0);
234 
235   const float lhs_data[] = {1, 2, 3, 4};
236   const float rhs_data[] = {1, 2};
237   float dst_data[4];
238 
239   ruy::Matrix<float> lhs;
240   ruy::MakeSimpleLayout(2, 2, ruy::Order::kRowMajor, lhs.mutable_layout());
241   lhs.set_data(lhs_data);
242   ruy::Matrix<float> rhs;
243   ruy::MakeSimpleLayout(2, 1, ruy::Order::kColMajor, rhs.mutable_layout());
244   rhs.set_data(rhs_data);
245   ruy::Matrix<float> dst;
246   ruy::MakeSimpleLayout(2, 1, ruy::Order::kColMajor, dst.mutable_layout());
247   dst.set_data(dst_data);
248 
249   ruy::MulParams<float, float> mul_params;
250   // Perform the multiplication and confirm no caching occurred.
251   ruy::Mul<ruy::kAllPaths>(lhs, rhs, mul_params, &context, &dst);
252   EXPECT_EQ(cache->MatrixCount(), 0);
253 
254   // Set cache policy for the LHS, repeat the multiplication, and see
255   // that caching did occur.
256   lhs.set_cache_policy(cache_policy);
257   ruy::Mul<ruy::kAllPaths>(lhs, rhs, mul_params, &context, &dst);
258   const bool actual_cached = cache->MatrixCount() == 1;
259   EXPECT_EQ(actual_cached, expected_cached);
260 }
261 
TEST(PrepackedCacheTest,TestCachePolicies)262 TEST(PrepackedCacheTest, TestCachePolicies) {
263   for (CachePolicy cache_policy :
264        {CachePolicy::kNeverCache, CachePolicy::kCacheIfLargeSpeedup,
265         CachePolicy::kCacheIfSignificantSpeedup, CachePolicy::kAlwaysCache}) {
266     TestCachePolicies(cache_policy,
267                          cache_policy != CachePolicy::kNeverCache);
268   }
269 }
270 
TEST(PrepackedCacheTest,TestClearCache)271 TEST(PrepackedCacheTest, TestClearCache) {
272   ruy::Context context;
273   PrepackedCache* cache = get_ctx(&context)->GetPrepackedCache();
274   EXPECT_EQ(cache->MatrixCount(), 0);
275 
276   const float lhs_data[] = {1, 2, 3, 4};
277   const float rhs_data[] = {1, 2};
278   float dst_data[4];
279 
280   ruy::Matrix<float> lhs;
281   ruy::MakeSimpleLayout(2, 2, ruy::Order::kRowMajor, lhs.mutable_layout());
282   lhs.set_data(lhs_data);
283   ruy::Matrix<float> rhs;
284   ruy::MakeSimpleLayout(2, 1, ruy::Order::kColMajor, rhs.mutable_layout());
285   rhs.set_data(rhs_data);
286   ruy::Matrix<float> dst;
287   ruy::MakeSimpleLayout(2, 1, ruy::Order::kColMajor, dst.mutable_layout());
288   dst.set_data(dst_data);
289 
290   ruy::MulParams<float, float> mul_params;
291   // Set cache policy for the LHS and see that caching occurs.
292   lhs.set_cache_policy(CachePolicy::kAlwaysCache);
293   ruy::Mul<ruy::kAllPaths>(lhs, rhs, mul_params, &context, &dst);
294   EXPECT_NE(cache->MatrixCount(), 0);
295 
296   // Clear the cache via the Context.
297   context.ClearPrepackedCache();
298   // Verify that the cache is now empty.
299   cache = get_ctx(&context)->GetPrepackedCache();
300   EXPECT_EQ(cache->MatrixCount(), 0);
301 }
302 
303 }  // namespace
304 }  // namespace ruy
305 
main(int argc,char ** argv)306 int main(int argc, char** argv) {
307   ::testing::InitGoogleTest(&argc, argv);
308   return RUN_ALL_TESTS();
309 }
310