1 /*
2 //@HEADER
3 // ************************************************************************
4 //
5 // Kokkos v. 3.0
6 // Copyright (2020) National Technology & Engineering
7 // Solutions of Sandia, LLC (NTESS).
8 //
9 // Under the terms of Contract DE-NA0003525 with NTESS,
10 // the U.S. Government retains certain rights in this software.
11 //
12 // Redistribution and use in source and binary forms, with or without
13 // modification, are permitted provided that the following conditions are
14 // met:
15 //
16 // 1. Redistributions of source code must retain the above copyright
17 // notice, this list of conditions and the following disclaimer.
18 //
19 // 2. Redistributions in binary form must reproduce the above copyright
20 // notice, this list of conditions and the following disclaimer in the
21 // documentation and/or other materials provided with the distribution.
22 //
23 // 3. Neither the name of the Corporation nor the names of the
24 // contributors may be used to endorse or promote products derived from
25 // this software without specific prior written permission.
26 //
27 // THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
28 // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
29 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
30 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
31 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
32 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
33 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
34 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
35 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
36 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
37 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
38 //
39 // Questions? Contact Christian R. Trott (crtrott@sandia.gov)
40 //
41 // ************************************************************************
42 //@HEADER
43 */
44
45 #ifndef KOKKOS_UNITTEST_MEMPOOL_HPP
46 #define KOKKOS_UNITTEST_MEMPOOL_HPP
47
48 #include <cstdio>
49 #include <iostream>
50 #include <cmath>
51 #include <algorithm>
52
53 #include <impl/Kokkos_Timer.hpp>
54
55 namespace TestMemoryPool {
56
57 template <typename MemSpace = Kokkos::HostSpace>
test_host_memory_pool_defaults()58 void test_host_memory_pool_defaults() {
59 using Space = typename MemSpace::execution_space;
60 using MemPool = typename Kokkos::MemoryPool<Space>;
61
62 {
63 const size_t MemoryCapacity = 32000;
64 const size_t MinBlockSize = 64;
65 const size_t MaxBlockSize = 1024;
66 const size_t SuperBlockSize = 4096;
67
68 MemPool pool(MemSpace(), MemoryCapacity, MinBlockSize, MaxBlockSize,
69 SuperBlockSize);
70
71 typename MemPool::usage_statistics stats;
72
73 pool.get_usage_statistics(stats);
74
75 ASSERT_LE(MemoryCapacity, stats.capacity_bytes);
76 ASSERT_LE(MinBlockSize, stats.min_block_bytes);
77 ASSERT_LE(MaxBlockSize, stats.max_block_bytes);
78 ASSERT_LE(SuperBlockSize, stats.superblock_bytes);
79 }
80
81 {
82 const size_t MemoryCapacity = 10000;
83
84 MemPool pool(MemSpace(), MemoryCapacity);
85
86 typename MemPool::usage_statistics stats;
87
88 pool.get_usage_statistics(stats);
89
90 ASSERT_LE(MemoryCapacity, stats.capacity_bytes);
91 ASSERT_LE(64u /* default */, stats.min_block_bytes);
92 ASSERT_LE(stats.min_block_bytes, stats.max_block_bytes);
93 ASSERT_LE(stats.max_block_bytes, stats.superblock_bytes);
94 ASSERT_LE(stats.superblock_bytes, stats.capacity_bytes);
95 }
96
97 {
98 const size_t MemoryCapacity = 10000;
99 const size_t MinBlockSize = 32; // power of two is exact
100
101 MemPool pool(MemSpace(), MemoryCapacity, MinBlockSize);
102
103 typename MemPool::usage_statistics stats;
104
105 pool.get_usage_statistics(stats);
106
107 ASSERT_LE(MemoryCapacity, stats.capacity_bytes);
108 ASSERT_EQ(MinBlockSize, stats.min_block_bytes);
109 ASSERT_LE(stats.min_block_bytes, stats.max_block_bytes);
110 ASSERT_LE(stats.max_block_bytes, stats.superblock_bytes);
111 ASSERT_LE(stats.superblock_bytes, stats.capacity_bytes);
112 }
113
114 {
115 const size_t MemoryCapacity = 32000;
116 const size_t MinBlockSize = 32; // power of two is exact
117 const size_t MaxBlockSize = 1024; // power of two is exact
118
119 MemPool pool(MemSpace(), MemoryCapacity, MinBlockSize, MaxBlockSize);
120
121 typename MemPool::usage_statistics stats;
122
123 pool.get_usage_statistics(stats);
124
125 ASSERT_LE(MemoryCapacity, stats.capacity_bytes);
126 ASSERT_EQ(MinBlockSize, stats.min_block_bytes);
127 ASSERT_EQ(MaxBlockSize, stats.max_block_bytes);
128 ASSERT_LE(stats.max_block_bytes, stats.superblock_bytes);
129 ASSERT_LE(stats.superblock_bytes, stats.capacity_bytes);
130 }
131 }
132
133 template <typename MemSpace = Kokkos::HostSpace>
test_host_memory_pool_stats()134 void test_host_memory_pool_stats() {
135 using Space = typename MemSpace::execution_space;
136 using MemPool = typename Kokkos::MemoryPool<Space>;
137
138 const size_t MemoryCapacity = 32000;
139 const size_t MinBlockSize = 64;
140 const size_t MaxBlockSize = 1024;
141 const size_t SuperBlockSize = 4096;
142
143 MemPool pool(MemSpace(), MemoryCapacity, MinBlockSize, MaxBlockSize,
144 SuperBlockSize);
145
146 {
147 typename MemPool::usage_statistics stats;
148
149 pool.get_usage_statistics(stats);
150
151 ASSERT_LE(MemoryCapacity, stats.capacity_bytes);
152 ASSERT_LE(MinBlockSize, stats.min_block_bytes);
153 ASSERT_LE(MaxBlockSize, stats.max_block_bytes);
154 ASSERT_LE(SuperBlockSize, stats.superblock_bytes);
155 }
156
157 void* p0064 = pool.allocate(64);
158 void* p0128 = pool.allocate(128);
159 void* p0256 = pool.allocate(256);
160 void* p1024 = pool.allocate(1024);
161
162 // Aborts because exceeds max block size:
163 // void * p2048 = pool.allocate(2048);
164
165 ASSERT_NE(p0064, nullptr);
166 ASSERT_NE(p0128, nullptr);
167 ASSERT_NE(p0256, nullptr);
168 ASSERT_NE(p1024, nullptr);
169
170 pool.deallocate(p0064, 64);
171 pool.deallocate(p0128, 128);
172 pool.deallocate(p0256, 256);
173 pool.deallocate(p1024, 1024);
174 }
175
176 //----------------------------------------------------------------------------
177 //----------------------------------------------------------------------------
178
179 template <class DeviceType>
180 struct TestMemoryPool_Functor {
181 using ptrs_type = Kokkos::View<uintptr_t*, DeviceType>;
182 using pool_type = Kokkos::MemoryPool<DeviceType>;
183
184 pool_type pool;
185 ptrs_type ptrs;
186
TestMemoryPool_FunctorTestMemoryPool::TestMemoryPool_Functor187 TestMemoryPool_Functor(const pool_type& arg_pool, size_t n)
188 : pool(arg_pool), ptrs("ptrs", n) {}
189
190 // Specify reduction argument value_type to avoid
191 // confusion with tag-dispatch.
192
193 using value_type = long;
194
195 struct TagAlloc {};
196
197 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPool_Functor198 void operator()(TagAlloc, int i, long& update) const noexcept {
199 unsigned alloc_size = 32 * (1 + (i % 5));
200 ptrs(i) = (uintptr_t)pool.allocate(alloc_size);
201 if (ptrs(i)) {
202 ++update;
203 }
204 }
205
206 struct TagDealloc {};
207
208 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPool_Functor209 void operator()(TagDealloc, int i, long& update) const noexcept {
210 if (ptrs(i) && (0 == i % 3)) {
211 unsigned alloc_size = 32 * (1 + (i % 5));
212 pool.deallocate((void*)ptrs(i), alloc_size);
213 ptrs(i) = 0;
214 ++update;
215 }
216 }
217
218 struct TagRealloc {};
219
220 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPool_Functor221 void operator()(TagRealloc, int i, long& update) const noexcept {
222 if (0 == ptrs(i)) {
223 unsigned alloc_size = 32 * (1 + (i % 5));
224 ptrs(i) = (uintptr_t)pool.allocate(alloc_size);
225 if (ptrs(i)) {
226 ++update;
227 }
228 }
229 }
230
231 struct TagMixItUp {};
232
233 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPool_Functor234 void operator()(TagMixItUp, int i, long& update) const noexcept {
235 if (ptrs(i) && (0 == i % 3)) {
236 unsigned alloc_size = 32 * (1 + (i % 5));
237
238 pool.deallocate((void*)ptrs(i), alloc_size);
239
240 ptrs(i) = (uintptr_t)pool.allocate(alloc_size);
241
242 if (ptrs(i)) {
243 ++update;
244 }
245 }
246 }
247 };
248
249 template <class PoolType>
print_memory_pool_stats(typename PoolType::usage_statistics const & stats)250 void print_memory_pool_stats(typename PoolType::usage_statistics const& stats) {
251 std::cout << "MemoryPool {" << std::endl
252 << " bytes capacity = " << stats.capacity_bytes << std::endl
253 << " bytes used = " << stats.consumed_bytes << std::endl
254 << " bytes reserved = " << stats.reserved_bytes << std::endl
255 << " bytes free = "
256 << (stats.capacity_bytes -
257 (stats.consumed_bytes + stats.reserved_bytes))
258 << std::endl
259 << " block used = " << stats.consumed_blocks << std::endl
260 << " block reserved = " << stats.reserved_blocks << std::endl
261 << " super used = " << stats.consumed_superblocks << std::endl
262 << " super reserved = "
263 << (stats.capacity_superblocks - stats.consumed_superblocks)
264 << std::endl
265 << "}" << std::endl;
266 }
267
268 template <class DeviceType>
test_memory_pool_v2(const bool print_statistics,const bool print_superblocks)269 void test_memory_pool_v2(const bool print_statistics,
270 const bool print_superblocks) {
271 using memory_space = typename DeviceType::memory_space;
272 using execution_space = typename DeviceType::execution_space;
273 using pool_type = Kokkos::MemoryPool<DeviceType>;
274 using functor_type = TestMemoryPool_Functor<DeviceType>;
275
276 using TagAlloc = typename functor_type::TagAlloc;
277 using TagDealloc = typename functor_type::TagDealloc;
278 using TagRealloc = typename functor_type::TagRealloc;
279 using TagMixItUp = typename functor_type::TagMixItUp;
280
281 const size_t total_alloc_size = 10000000;
282 const unsigned min_block_size = 64;
283 const unsigned max_block_size = 256;
284 const long nfill = 70000;
285
286 for (uint32_t k = 0, min_superblock_size = 10000; k < 3;
287 ++k, min_superblock_size *= 10) {
288 typename pool_type::usage_statistics stats;
289
290 pool_type pool(memory_space(), total_alloc_size, min_block_size,
291 max_block_size, min_superblock_size);
292
293 functor_type functor(pool, nfill);
294
295 long result = 0;
296 long ndel = 0;
297
298 Kokkos::parallel_reduce(
299 Kokkos::RangePolicy<execution_space, TagAlloc>(0, nfill), functor,
300 result);
301
302 pool.get_usage_statistics(stats);
303
304 const int fill_error =
305 (nfill != result) || (nfill != long(stats.consumed_blocks));
306
307 if (fill_error || print_statistics)
308 print_memory_pool_stats<pool_type>(stats);
309 if (fill_error || print_superblocks) pool.print_state(std::cout);
310
311 ASSERT_EQ(nfill, result);
312 ASSERT_EQ(nfill, long(stats.consumed_blocks));
313
314 Kokkos::parallel_reduce(
315 Kokkos::RangePolicy<execution_space, TagDealloc>(0, nfill), functor,
316 ndel);
317
318 pool.get_usage_statistics(stats);
319
320 const int del_error = (nfill - ndel) != long(stats.consumed_blocks);
321
322 if (del_error || print_statistics)
323 print_memory_pool_stats<pool_type>(stats);
324 if (del_error || print_superblocks) pool.print_state(std::cout);
325
326 ASSERT_EQ((nfill - ndel), long(stats.consumed_blocks));
327
328 Kokkos::parallel_reduce(
329 Kokkos::RangePolicy<execution_space, TagRealloc>(0, nfill), functor,
330 result);
331
332 pool.get_usage_statistics(stats);
333
334 const int refill_error =
335 (ndel != result) || (nfill != long(stats.consumed_blocks));
336
337 if (refill_error || print_statistics)
338 print_memory_pool_stats<pool_type>(stats);
339 if (refill_error || print_superblocks) pool.print_state(std::cout);
340
341 ASSERT_EQ(ndel, result);
342 ASSERT_EQ(nfill, long(stats.consumed_blocks));
343
344 Kokkos::parallel_reduce(
345 Kokkos::RangePolicy<execution_space, TagMixItUp>(0, nfill), functor,
346 result);
347
348 pool.get_usage_statistics(stats);
349
350 const int mix_error =
351 (ndel != result) || (nfill != long(stats.consumed_blocks));
352
353 if (mix_error || print_statistics)
354 print_memory_pool_stats<pool_type>(stats);
355 if (mix_error || print_superblocks) pool.print_state(std::cout);
356
357 ASSERT_EQ(ndel, result);
358 ASSERT_EQ(nfill, long(stats.consumed_blocks));
359 }
360 }
361
362 //----------------------------------------------------------------------------
363 //----------------------------------------------------------------------------
364
365 template <class DeviceType>
366 struct TestMemoryPoolCorners {
367 using ptrs_type = Kokkos::View<uintptr_t*, DeviceType>;
368 using pool_type = Kokkos::MemoryPool<DeviceType>;
369
370 pool_type pool;
371 ptrs_type ptrs;
372 uint32_t size;
373 uint32_t stride;
374
TestMemoryPoolCornersTestMemoryPool::TestMemoryPoolCorners375 TestMemoryPoolCorners(const pool_type& arg_pool, const ptrs_type& arg_ptrs,
376 const uint32_t arg_base, const uint32_t arg_stride)
377 : pool(arg_pool), ptrs(arg_ptrs), size(arg_base), stride(arg_stride) {}
378
379 // Specify reduction argument value_type to
380 // avoid confusion with tag-dispatch.
381
382 using value_type = long;
383
384 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPoolCorners385 void operator()(int i, long& err) const noexcept {
386 unsigned alloc_size = size << (i % stride);
387 if (0 == ptrs(i)) {
388 ptrs(i) = (uintptr_t)pool.allocate(alloc_size);
389 if (ptrs(i) && !alloc_size) {
390 ++err;
391 }
392 }
393 }
394
395 struct TagDealloc {};
396
397 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPoolCorners398 void operator()(int i) const noexcept {
399 unsigned alloc_size = size << (i % stride);
400 if (ptrs(i)) {
401 pool.deallocate((void*)ptrs(i), alloc_size);
402 }
403 ptrs(i) = 0;
404 }
405 };
406
407 template <class DeviceType>
test_memory_pool_corners(const bool print_statistics,const bool print_superblocks)408 void test_memory_pool_corners(const bool print_statistics,
409 const bool print_superblocks) {
410 using memory_space = typename DeviceType::memory_space;
411 using execution_space = typename DeviceType::execution_space;
412 using pool_type = Kokkos::MemoryPool<DeviceType>;
413 using functor_type = TestMemoryPoolCorners<DeviceType>;
414 using ptrs_type = typename functor_type::ptrs_type;
415
416 {
417 // superblock size 1 << 14
418 const size_t min_superblock_size = 1u << 14;
419
420 // four superblocks
421 const size_t total_alloc_size = min_superblock_size * 4;
422
423 // block sizes { 64 , 128 , 256 , 512 }
424 // block counts { 256 , 128 , 64 , 32 }
425 const unsigned min_block_size = 64;
426 const unsigned max_block_size = 512;
427 const unsigned num_blocks = 480;
428
429 pool_type pool(memory_space(), total_alloc_size, min_block_size,
430 max_block_size, min_superblock_size);
431
432 // Allocate one block from each superblock to lock that
433 // superblock into the block size.
434
435 ptrs_type ptrs("ptrs", num_blocks);
436
437 long err = 0;
438
439 Kokkos::parallel_reduce(Kokkos::RangePolicy<execution_space>(0, 4),
440 functor_type(pool, ptrs, 64, 4), err);
441
442 if (print_statistics || err) {
443 typename pool_type::usage_statistics stats;
444
445 pool.get_usage_statistics(stats);
446
447 print_memory_pool_stats<pool_type>(stats);
448 }
449
450 if (print_superblocks || err) {
451 pool.print_state(std::cout);
452 }
453
454 // Now fill remaining allocations with small size
455
456 Kokkos::parallel_reduce(Kokkos::RangePolicy<execution_space>(0, num_blocks),
457 functor_type(pool, ptrs, 64, 1), err);
458
459 if (print_statistics || err) {
460 typename pool_type::usage_statistics stats;
461
462 pool.get_usage_statistics(stats);
463
464 print_memory_pool_stats<pool_type>(stats);
465 }
466
467 if (print_superblocks || err) {
468 pool.print_state(std::cout);
469 }
470 }
471 }
472
473 //----------------------------------------------------------------------------
474 //----------------------------------------------------------------------------
475
476 template <class DeviceType, class Enable = void>
477 struct TestMemoryPoolHuge {
478 enum : size_t { num_superblock = 0 };
479
480 using value_type = long;
481
482 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPoolHuge483 void operator()(int /*i*/, long& /*err*/) const noexcept {}
484
485 KOKKOS_INLINE_FUNCTION
operator ()TestMemoryPool::TestMemoryPoolHuge486 void operator()(int /*i*/) const noexcept {}
487 };
488
489 template <class DeviceType>
490 struct TestMemoryPoolHuge<
491 DeviceType,
492 typename std::enable_if<std::is_same<
493 Kokkos::HostSpace, typename DeviceType::memory_space>::value>::type> {
494 using ptrs_type = Kokkos::View<uintptr_t*, DeviceType>;
495 using pool_type = Kokkos::MemoryPool<DeviceType>;
496 using memory_space = typename DeviceType::memory_space;
497
498 pool_type pool;
499 ptrs_type ptrs;
500
501 enum : size_t {
502 min_block_size = 512,
503 max_block_size = 1lu << 31,
504 min_superblock_size = max_block_size,
505 num_superblock = 4,
506 total_alloc_size = num_superblock * max_block_size
507 };
508
TestMemoryPoolHugeTestMemoryPool::TestMemoryPoolHuge509 TestMemoryPoolHuge()
510 : pool(memory_space(), total_alloc_size, min_block_size, max_block_size,
511 min_superblock_size),
512 ptrs("ptrs", num_superblock) {}
513
514 // Specify reduction argument value_type to
515 // avoid confusion with tag-dispatch.
516
517 using value_type = long;
518
operator ()TestMemoryPool::TestMemoryPoolHuge519 void operator()(int i, long& err) const noexcept {
520 if (i < int(num_superblock)) {
521 ptrs(i) = (uintptr_t)pool.allocate(max_block_size);
522 #if 0
523 printf("TestMemoryPoolHuge size(0x%lx) ptr(0x%lx)\n"
524 , max_block_size
525 , ptrs(i) );
526 #endif
527 if (!ptrs(i)) {
528 Kokkos::abort("TestMemoryPoolHuge");
529 ++err;
530 }
531 }
532 }
533
operator ()TestMemoryPool::TestMemoryPoolHuge534 void operator()(int i) const noexcept {
535 if (i < int(num_superblock)) {
536 pool.deallocate((void*)ptrs(i), max_block_size);
537 ptrs(i) = 0;
538 }
539 }
540 };
541
542 template <class DeviceType>
test_memory_pool_huge()543 void test_memory_pool_huge() {
544 using execution_space = typename DeviceType::execution_space;
545 using functor_type = TestMemoryPoolHuge<DeviceType>;
546 using policy_type = Kokkos::RangePolicy<execution_space>;
547
548 functor_type f;
549 policy_type policy(0, functor_type::num_superblock);
550
551 long err = 0;
552
553 Kokkos::parallel_reduce(policy, f, err);
554 Kokkos::parallel_for(policy, f);
555 }
556
557 //----------------------------------------------------------------------------
558 //----------------------------------------------------------------------------
559
560 } // namespace TestMemoryPool
561
562 namespace Test {
563
TEST(TEST_CATEGORY,memory_pool)564 TEST(TEST_CATEGORY, memory_pool) {
565 TestMemoryPool::test_host_memory_pool_defaults<>();
566 TestMemoryPool::test_host_memory_pool_stats<>();
567 TestMemoryPool::test_memory_pool_v2<TEST_EXECSPACE>(false, false);
568 TestMemoryPool::test_memory_pool_corners<TEST_EXECSPACE>(false, false);
569 #ifdef KOKKOS_ENABLE_LARGE_MEM_TESTS
570 TestMemoryPool::test_memory_pool_huge<TEST_EXECSPACE>();
571 #endif
572 }
573
574 } // namespace Test
575
576 #endif
577