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