Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_MemoryPool.hpp
1//@HEADER
2// ************************************************************************
3//
4// Kokkos v. 4.0
5// Copyright (2022) National Technology & Engineering
6// Solutions of Sandia, LLC (NTESS).
7//
8// Under the terms of Contract DE-NA0003525 with NTESS,
9// the U.S. Government retains certain rights in this software.
10//
11// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12// See https://kokkos.org/LICENSE for license information.
13// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14//
15//@HEADER
16
17#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
18#include <Kokkos_Macros.hpp>
19static_assert(false,
20 "Including non-public Kokkos header files is not allowed.");
21#endif
22#ifndef KOKKOS_MEMORYPOOL_HPP
23#define KOKKOS_MEMORYPOOL_HPP
24
25#include <Kokkos_Core_fwd.hpp>
26#include <Kokkos_Parallel.hpp>
27#include <Kokkos_Atomic.hpp>
28#include <impl/Kokkos_ConcurrentBitset.hpp>
29#include <impl/Kokkos_Error.hpp>
30#include <impl/Kokkos_SharedAlloc.hpp>
31
32#include <iostream>
33
34namespace Kokkos {
35namespace Impl {
36/* Report violation of size constraints:
37 * min_block_alloc_size <= max_block_alloc_size
38 * max_block_alloc_size <= min_superblock_size
39 * min_superblock_size <= max_superblock_size
40 * min_superblock_size <= min_total_alloc_size
41 * min_superblock_size <= min_block_alloc_size *
42 * max_block_per_superblock
43 */
44void memory_pool_bounds_verification(size_t min_block_alloc_size,
45 size_t max_block_alloc_size,
46 size_t min_superblock_size,
47 size_t max_superblock_size,
48 size_t max_block_per_superblock,
49 size_t min_total_alloc_size);
50} // namespace Impl
51} // namespace Kokkos
52
53namespace Kokkos {
54
55namespace Impl {
56
57void _print_memory_pool_state(std::ostream &s, uint32_t const *sb_state_ptr,
58 int32_t sb_count, uint32_t sb_size_lg2,
59 uint32_t sb_state_size, uint32_t state_shift,
60 uint32_t state_used_mask);
61
62} // end namespace Impl
63
64template <typename DeviceType>
65class MemoryPool {
66 private:
67 using CB = Kokkos::Impl::concurrent_bitset;
68
69 enum : uint32_t { bits_per_int_lg2 = CB::bits_per_int_lg2 };
70 enum : uint32_t { state_shift = CB::state_shift };
71 enum : uint32_t { state_used_mask = CB::state_used_mask };
72 enum : uint32_t { state_header_mask = CB::state_header_mask };
73 enum : uint32_t { max_bit_count_lg2 = CB::max_bit_count_lg2 };
74 enum : uint32_t { max_bit_count = CB::max_bit_count };
75
76 enum : uint32_t { HINT_PER_BLOCK_SIZE = 2 };
77
78 /* Each superblock has a concurrent bitset state
79 * which is an array of uint32_t integers.
80 * [ { block_count_lg2 : state_shift bits
81 * , used_block_count : ( 32 - state_shift ) bits
82 * }
83 * , { block allocation bit set }* ]
84 *
85 * As superblocks are assigned (allocated) to a block size
86 * and released (deallocated) back to empty the superblock state
87 * is concurrently updated.
88 */
89
90 /* Mapping between block_size <-> block_state
91 *
92 * block_state = ( m_sb_size_lg2 - block_size_lg2 ) << state_shift
93 * block_size = m_sb_size_lg2 - ( block_state >> state_shift )
94 *
95 * Thus A_block_size < B_block_size <=> A_block_state > B_block_state
96 */
97
98 using base_memory_space = typename DeviceType::memory_space;
99
100 enum {
102 base_memory_space>::accessible
103 };
104
105 using Tracker = Kokkos::Impl::SharedAllocationTracker;
106 using Record = Kokkos::Impl::SharedAllocationRecord<base_memory_space>;
107
108 Tracker m_tracker;
109 uint32_t *m_sb_state_array;
110 uint32_t m_sb_state_size;
111 uint32_t m_sb_size_lg2;
112 uint32_t m_max_block_size_lg2;
113 uint32_t m_min_block_size_lg2;
114 int32_t m_sb_count;
115 int32_t m_hint_offset; // Offset to K * #block_size array of hints
116 int32_t m_data_offset; // Offset to 0th superblock data
117 int32_t m_unused_padding;
118
119 public:
120 using memory_space = typename DeviceType::memory_space;
121
123 enum : uint32_t { max_superblock_size = 1LU << 31 /* 2 gigabytes */ };
124 enum : uint32_t { max_block_per_superblock = max_bit_count };
125
126 //--------------------------------------------------------------------------
127
128 KOKKOS_INLINE_FUNCTION
129 bool operator==(MemoryPool const &other) const {
130 return m_sb_state_array == other.m_sb_state_array;
131 }
132
133 KOKKOS_INLINE_FUNCTION
134 size_t capacity() const noexcept {
135 return size_t(m_sb_count) << m_sb_size_lg2;
136 }
137
138 KOKKOS_INLINE_FUNCTION
139 size_t min_block_size() const noexcept {
140 return (1LU << m_min_block_size_lg2);
141 }
142
143 KOKKOS_INLINE_FUNCTION
144 size_t max_block_size() const noexcept {
145 return (1LU << m_max_block_size_lg2);
146 }
147
148 struct usage_statistics {
149 size_t capacity_bytes;
150 size_t superblock_bytes;
151 size_t max_block_bytes;
152 size_t min_block_bytes;
153 size_t capacity_superblocks;
154 size_t consumed_superblocks;
155 size_t consumed_blocks;
156 size_t consumed_bytes;
157 size_t reserved_blocks;
158 size_t reserved_bytes;
159 };
160
161 void get_usage_statistics(usage_statistics &stats) const {
163
164 const size_t alloc_size = m_hint_offset * sizeof(uint32_t);
165
166 uint32_t *const sb_state_array =
167 accessible ? m_sb_state_array : (uint32_t *)host.allocate(alloc_size);
168
169 if (!accessible) {
170 Kokkos::Impl::DeepCopy<Kokkos::HostSpace, base_memory_space>(
171 sb_state_array, m_sb_state_array, alloc_size);
172 Kokkos::fence(
173 "MemoryPool::get_usage_statistics(): fence after copying state "
174 "array to HostSpace");
175 }
176
177 stats.superblock_bytes = (1LU << m_sb_size_lg2);
178 stats.max_block_bytes = (1LU << m_max_block_size_lg2);
179 stats.min_block_bytes = (1LU << m_min_block_size_lg2);
180 stats.capacity_bytes = stats.superblock_bytes * m_sb_count;
181 stats.capacity_superblocks = m_sb_count;
182 stats.consumed_superblocks = 0;
183 stats.consumed_blocks = 0;
184 stats.consumed_bytes = 0;
185 stats.reserved_blocks = 0;
186 stats.reserved_bytes = 0;
187
188 const uint32_t *sb_state_ptr = sb_state_array;
189
190 for (int32_t i = 0; i < m_sb_count; ++i, sb_state_ptr += m_sb_state_size) {
191 const uint32_t block_count_lg2 = (*sb_state_ptr) >> state_shift;
192
193 if (block_count_lg2) {
194 const uint32_t block_count = 1u << block_count_lg2;
195 const uint32_t block_size_lg2 = m_sb_size_lg2 - block_count_lg2;
196 const uint32_t block_size = 1u << block_size_lg2;
197 const uint32_t block_used = (*sb_state_ptr) & state_used_mask;
198
199 stats.consumed_superblocks++;
200 stats.consumed_blocks += block_used;
201 stats.consumed_bytes += block_used * block_size;
202 stats.reserved_blocks += block_count - block_used;
203 stats.reserved_bytes += (block_count - block_used) * block_size;
204 }
205 }
206
207 if (!accessible) {
208 host.deallocate(sb_state_array, alloc_size);
209 }
210 }
211
212 void print_state(std::ostream &s) const {
214
215 const size_t alloc_size = m_hint_offset * sizeof(uint32_t);
216
217 uint32_t *const sb_state_array =
218 accessible ? m_sb_state_array : (uint32_t *)host.allocate(alloc_size);
219
220 if (!accessible) {
221 Kokkos::Impl::DeepCopy<Kokkos::HostSpace, base_memory_space>(
222 sb_state_array, m_sb_state_array, alloc_size);
223 Kokkos::fence(
224 "MemoryPool::print_state(): fence after copying state array to "
225 "HostSpace");
226 }
227
228 Impl::_print_memory_pool_state(s, sb_state_array, m_sb_count, m_sb_size_lg2,
229 m_sb_state_size, state_shift,
230 state_used_mask);
231
232 if (!accessible) {
233 host.deallocate(sb_state_array, alloc_size);
234 }
235 }
236
237 //--------------------------------------------------------------------------
238
239 KOKKOS_DEFAULTED_FUNCTION MemoryPool(MemoryPool &&) = default;
240 KOKKOS_DEFAULTED_FUNCTION MemoryPool(const MemoryPool &) = default;
241 KOKKOS_DEFAULTED_FUNCTION MemoryPool &operator=(MemoryPool &&) = default;
242 KOKKOS_DEFAULTED_FUNCTION MemoryPool &operator=(const MemoryPool &) = default;
243
244 KOKKOS_INLINE_FUNCTION MemoryPool()
245 : m_tracker(),
246 m_sb_state_array(nullptr),
247 m_sb_state_size(0),
248 m_sb_size_lg2(0),
249 m_max_block_size_lg2(0),
250 m_min_block_size_lg2(0),
251 m_sb_count(0),
252 m_hint_offset(0),
253 m_data_offset(0),
254 m_unused_padding(0) {}
255
270 MemoryPool(const base_memory_space &memspace,
271 const size_t min_total_alloc_size, size_t min_block_alloc_size = 0,
272 size_t max_block_alloc_size = 0, size_t min_superblock_size = 0)
273 : m_tracker(),
274 m_sb_state_array(nullptr),
275 m_sb_state_size(0),
276 m_sb_size_lg2(0),
277 m_max_block_size_lg2(0),
278 m_min_block_size_lg2(0),
279 m_sb_count(0),
280 m_hint_offset(0),
281 m_data_offset(0),
282 m_unused_padding(0) {
283 const uint32_t int_align_lg2 = 3; /* align as int[8] */
284 const uint32_t int_align_mask = (1u << int_align_lg2) - 1;
285 const uint32_t default_min_block_size = 1u << 6; /* 64 bytes */
286 const uint32_t default_max_block_size = 1u << 12; /* 4k bytes */
287 const uint32_t default_min_superblock_size = 1u << 20; /* 1M bytes */
288
289 //--------------------------------------------------
290 // Default block and superblock sizes:
291
292 if (0 == min_block_alloc_size) {
293 // Default all sizes:
294
295 min_superblock_size =
296 std::min(size_t(default_min_superblock_size), min_total_alloc_size);
297
298 min_block_alloc_size =
299 std::min(size_t(default_min_block_size), min_superblock_size);
300
301 max_block_alloc_size =
302 std::min(size_t(default_max_block_size), min_superblock_size);
303 } else if (0 == min_superblock_size) {
304 // Choose superblock size as minimum of:
305 // max_block_per_superblock * min_block_size
306 // max_superblock_size
307 // min_total_alloc_size
308
309 const size_t max_superblock =
310 min_block_alloc_size * max_block_per_superblock;
311
312 min_superblock_size =
313 std::min(max_superblock,
314 std::min(size_t(max_superblock_size), min_total_alloc_size));
315 }
316
317 if (0 == max_block_alloc_size) {
318 max_block_alloc_size = min_superblock_size;
319 }
320
321 //--------------------------------------------------
322
323 /* Enforce size constraints:
324 * min_block_alloc_size <= max_block_alloc_size
325 * max_block_alloc_size <= min_superblock_size
326 * min_superblock_size <= max_superblock_size
327 * min_superblock_size <= min_total_alloc_size
328 * min_superblock_size <= min_block_alloc_size *
329 * max_block_per_superblock
330 */
331
332 Kokkos::Impl::memory_pool_bounds_verification(
333 min_block_alloc_size, max_block_alloc_size, min_superblock_size,
334 max_superblock_size, max_block_per_superblock, min_total_alloc_size);
335
336 //--------------------------------------------------
337 // Block and superblock size is power of two:
338 // Maximum value is 'max_superblock_size'
339
340 m_min_block_size_lg2 =
341 Kokkos::Impl::integral_power_of_two_that_contains(min_block_alloc_size);
342
343 m_max_block_size_lg2 =
344 Kokkos::Impl::integral_power_of_two_that_contains(max_block_alloc_size);
345
346 m_sb_size_lg2 =
347 Kokkos::Impl::integral_power_of_two_that_contains(min_superblock_size);
348
349 {
350 // number of superblocks is multiple of superblock size that
351 // can hold min_total_alloc_size.
352
353 const uint64_t sb_size_mask = (1LU << m_sb_size_lg2) - 1;
354
355 m_sb_count = (min_total_alloc_size + sb_size_mask) >> m_sb_size_lg2;
356 }
357
358 {
359 // Any superblock can be assigned to the smallest size block
360 // Size the block bitset to maximum number of blocks
361
362 const uint32_t max_block_count_lg2 = m_sb_size_lg2 - m_min_block_size_lg2;
363
364 m_sb_state_size =
365 (CB::buffer_bound_lg2(max_block_count_lg2) + int_align_mask) &
366 ~int_align_mask;
367 }
368
369 // Array of all superblock states
370
371 const size_t all_sb_state_size =
372 (m_sb_count * m_sb_state_size + int_align_mask) & ~int_align_mask;
373
374 // Number of block sizes
375
376 const int32_t number_block_sizes =
377 1 + m_max_block_size_lg2 - m_min_block_size_lg2;
378
379 // Array length for possible block sizes
380 // Hint array is one uint32_t per block size
381
382 const int32_t block_size_array_size =
383 (number_block_sizes + int_align_mask) & ~int_align_mask;
384
385 m_hint_offset = all_sb_state_size;
386 m_data_offset = m_hint_offset + block_size_array_size * HINT_PER_BLOCK_SIZE;
387
388 // Allocation:
389
390 const size_t header_size = m_data_offset * sizeof(uint32_t);
391 const size_t alloc_size =
392 header_size + (size_t(m_sb_count) << m_sb_size_lg2);
393
394 Record *rec = Record::allocate(memspace, "Kokkos::MemoryPool", alloc_size);
395
396 m_tracker.assign_allocated_record_to_uninitialized(rec);
397
398 m_sb_state_array = (uint32_t *)rec->data();
399
401
402 uint32_t *const sb_state_array =
403 accessible ? m_sb_state_array : (uint32_t *)host.allocate(header_size);
404
405 for (int32_t i = 0; i < m_data_offset; ++i) sb_state_array[i] = 0;
406
407 // Initial assignment of empty superblocks to block sizes:
408
409 for (int32_t i = 0; i < number_block_sizes; ++i) {
410 const uint32_t block_size_lg2 = i + m_min_block_size_lg2;
411 const uint32_t block_count_lg2 = m_sb_size_lg2 - block_size_lg2;
412 const uint32_t block_state = block_count_lg2 << state_shift;
413 const uint32_t hint_begin = m_hint_offset + i * HINT_PER_BLOCK_SIZE;
414
415 // for block size index 'i':
416 // sb_id_hint = sb_state_array[ hint_begin ];
417 // sb_id_begin = sb_state_array[ hint_begin + 1 ];
418
419 const int32_t jbeg = (i * m_sb_count) / number_block_sizes;
420 const int32_t jend = ((i + 1) * m_sb_count) / number_block_sizes;
421
422 sb_state_array[hint_begin] = uint32_t(jbeg);
423 sb_state_array[hint_begin + 1] = uint32_t(jbeg);
424
425 for (int32_t j = jbeg; j < jend; ++j) {
426 sb_state_array[j * m_sb_state_size] = block_state;
427 }
428 }
429
430 // Write out initialized state:
431
432 if (!accessible) {
433 Kokkos::Impl::DeepCopy<base_memory_space, Kokkos::HostSpace>(
434 m_sb_state_array, sb_state_array, header_size);
435 Kokkos::fence(
436 "MemoryPool::MemoryPool(): fence after copying state array from "
437 "HostSpace");
438
439 host.deallocate(sb_state_array, header_size);
440 } else {
442 }
443 }
444
445 //--------------------------------------------------------------------------
446
447 private:
448 /* Given a size 'n' get the block size in which it can be allocated.
449 * Restrict lower bound to minimum block size.
450 */
451 KOKKOS_FORCEINLINE_FUNCTION
452 uint32_t get_block_size_lg2(uint32_t n) const noexcept {
453 const unsigned i = Kokkos::Impl::integral_power_of_two_that_contains(n);
454
455 return i < m_min_block_size_lg2 ? m_min_block_size_lg2 : i;
456 }
457
458 public:
459 /* Return 0 for invalid block size */
460 KOKKOS_INLINE_FUNCTION
461 uint32_t allocate_block_size(uint64_t alloc_size) const noexcept {
462 return alloc_size <= (1UL << m_max_block_size_lg2)
463 ? (1UL << get_block_size_lg2(uint32_t(alloc_size)))
464 : 0;
465 }
466
467 //--------------------------------------------------------------------------
477 KOKKOS_FUNCTION
478 void *allocate(size_t alloc_size, int32_t attempt_limit = 1) const noexcept {
479 if (size_t(1LU << m_max_block_size_lg2) < alloc_size) {
481 "Kokkos MemoryPool allocation request exceeded specified maximum "
482 "allocation size");
483 }
484
485 if (0 == alloc_size) return nullptr;
486
487 void *p = nullptr;
488
489 const uint32_t block_size_lg2 = get_block_size_lg2(alloc_size);
490
491 // Allocation will fit within a superblock
492 // that has block sizes ( 1 << block_size_lg2 )
493
494 const uint32_t block_count_lg2 = m_sb_size_lg2 - block_size_lg2;
495 const uint32_t block_state = block_count_lg2 << state_shift;
496 const uint32_t block_count = 1u << block_count_lg2;
497
498 // Superblock hints for this block size:
499 // hint_sb_id_ptr[0] is the dynamically changing hint
500 // hint_sb_id_ptr[1] is the static start point
501
502 volatile uint32_t *const hint_sb_id_ptr =
503 m_sb_state_array /* memory pool state array */
504 + m_hint_offset /* offset to hint portion of array */
505 + HINT_PER_BLOCK_SIZE /* number of hints per block size */
506 * (block_size_lg2 - m_min_block_size_lg2); /* block size id */
507
508 const int32_t sb_id_begin = int32_t(hint_sb_id_ptr[1]);
509
510 // Fast query clock register 'tic' to pseudo-randomize
511 // the guess for which block within a superblock should
512 // be claimed. If not available then a search occurs.
513#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOS_ARCH_INTEL_GPU)
514 const uint32_t block_id_hint = alloc_size;
515#else
516 const uint32_t block_id_hint =
517 (uint32_t)(Kokkos::Impl::clock_tic()
518#ifdef __CUDA_ARCH__ // FIXME_CUDA
519 // Spread out potentially concurrent access
520 // by threads within a warp or thread block.
521 + (threadIdx.x + blockDim.x * threadIdx.y)
522#endif
523 );
524#endif
525
526 // expected state of superblock for allocation
527 uint32_t sb_state = block_state;
528
529 int32_t sb_id = -1;
530
531 volatile uint32_t *sb_state_array = nullptr;
532
533 while (attempt_limit) {
534 int32_t hint_sb_id = -1;
535
536 if (sb_id < 0) {
537 // No superblock specified, try the hint for this block size
538
539 sb_id = hint_sb_id = int32_t(*hint_sb_id_ptr);
540
541 sb_state_array = m_sb_state_array + (sb_id * m_sb_state_size);
542 }
543
544 // Require:
545 // 0 <= sb_id
546 // sb_state_array == m_sb_state_array + m_sb_state_size * sb_id
547
548 if (sb_state == (state_header_mask & *sb_state_array)) {
549 // This superblock state is as expected, for the moment.
550 // Attempt to claim a bit. The attempt updates the state
551 // so have already made sure the state header is as expected.
552
553 const uint32_t count_lg2 = sb_state >> state_shift;
554 const uint32_t mask = (1u << count_lg2) - 1;
555
556 const Kokkos::pair<int, int> result = CB::acquire_bounded_lg2(
557 sb_state_array, count_lg2, block_id_hint & mask, sb_state);
558
559 // If result.first < 0 then failed to acquire
560 // due to either full or buffer was wrong state.
561 // Could be wrong state if a deallocation raced the
562 // superblock to empty before the acquire could succeed.
563
564 if (0 <= result.first) { // acquired a bit
565
566 const uint32_t size_lg2 = m_sb_size_lg2 - count_lg2;
567
568 // Set the allocated block pointer
569
570 p = ((char *)(m_sb_state_array + m_data_offset)) +
571 (uint64_t(sb_id) << m_sb_size_lg2) // superblock memory
572 + (uint64_t(result.first) << size_lg2); // block memory
573
574 break; // Success
575 }
576 }
577 //------------------------------------------------------------------
578 // Arrive here if failed to acquire a block.
579 // Must find a new superblock.
580
581 // Start searching at designated index for this block size.
582 // Look for superblock that, in preferential order,
583 // 1) part-full superblock of this block size
584 // 2) empty superblock to claim for this block size
585 // 3) part-full superblock of the next larger block size
586
587 sb_state = block_state; // Expect to find the desired state
588 sb_id = -1;
589
590 bool update_hint = false;
591 int32_t sb_id_empty = -1;
592 int32_t sb_id_large = -1;
593 uint32_t sb_state_large = 0;
594
595 sb_state_array = m_sb_state_array + sb_id_begin * m_sb_state_size;
596
597 for (int32_t i = 0, id = sb_id_begin; i < m_sb_count; ++i) {
598 // Query state of the candidate superblock.
599 // Note that the state may change at any moment
600 // as concurrent allocations and deallocations occur.
601
602 const uint32_t full_state = *sb_state_array;
603 const uint32_t used = full_state & state_used_mask;
604 const uint32_t state = full_state & state_header_mask;
605
606 if (state == block_state) {
607 // Superblock is assigned to this block size
608
609 if (used < block_count) {
610 // There is room to allocate one block
611
612 sb_id = id;
613
614 // Is there room to allocate more than one block?
615
616 update_hint = used + 1 < block_count;
617
618 break;
619 }
620 } else if (0 == used) {
621 // Superblock is empty
622
623 if (-1 == sb_id_empty) {
624 // Superblock is not assigned to this block size
625 // and is the first empty superblock encountered.
626 // Save this id to use if a partfull superblock is not found.
627
628 sb_id_empty = id;
629 }
630 } else if ((-1 == sb_id_empty /* have not found an empty */) &&
631 (-1 == sb_id_large /* have not found a larger */) &&
632 (state < block_state /* a larger block */) &&
633 // is not full:
634 (used < (1u << (state >> state_shift)))) {
635 // First superblock encountered that is
636 // larger than this block size and
637 // has room for an allocation.
638 // Save this id to use of partfull or empty superblock not found
639 sb_id_large = id;
640 sb_state_large = state;
641 }
642
643 // Iterate around the superblock array:
644
645 if (++id < m_sb_count) {
646 sb_state_array += m_sb_state_size;
647 } else {
648 id = 0;
649 sb_state_array = m_sb_state_array;
650 }
651 }
652
653 // printf(" search m_sb_count(%d) sb_id(%d) sb_id_empty(%d)
654 // sb_id_large(%d)\n" , m_sb_count , sb_id , sb_id_empty , sb_id_large);
655
656 if (sb_id < 0) {
657 // Did not find a partfull superblock for this block size.
658
659 if (0 <= sb_id_empty) {
660 // Found first empty superblock following designated superblock
661 // Attempt to claim it for this block size.
662 // If the claim fails assume that another thread claimed it
663 // for this block size and try to use it anyway,
664 // but do not update hint.
665
666 sb_id = sb_id_empty;
667
668 sb_state_array = m_sb_state_array + (sb_id * m_sb_state_size);
669
670 // If successfully changed assignment of empty superblock 'sb_id'
671 // to this block_size then update the hint.
672
673 const uint32_t state_empty = state_header_mask & *sb_state_array;
674
675 // If this thread claims the empty block then update the hint
676 update_hint =
677 state_empty == Kokkos::atomic_compare_exchange(
678 sb_state_array, state_empty, block_state);
679 } else if (0 <= sb_id_large) {
680 // Found a larger superblock with space available
681
682 sb_id = sb_id_large;
683 sb_state = sb_state_large;
684
685 sb_state_array = m_sb_state_array + (sb_id * m_sb_state_size);
686 } else {
687 // Did not find a potentially usable superblock
688 --attempt_limit;
689 }
690 }
691
692 if (update_hint) {
693 Kokkos::atomic_compare_exchange(hint_sb_id_ptr, uint32_t(hint_sb_id),
694 uint32_t(sb_id));
695 }
696 } // end allocation attempt loop
697 //--------------------------------------------------------------------
698
699 return p;
700 }
701 // end allocate
702 //--------------------------------------------------------------------------
703
710 KOKKOS_INLINE_FUNCTION
711 void deallocate(void *p, size_t /* alloc_size */) const noexcept {
712 if (nullptr == p) return;
713
714 // Determine which superblock and block
715 const ptrdiff_t d =
716 static_cast<char *>(p) -
717 reinterpret_cast<char *>(m_sb_state_array + m_data_offset);
718
719 // Verify contained within the memory pool's superblocks:
720 const int ok_contains =
721 (0 <= d) && (size_t(d) < (size_t(m_sb_count) << m_sb_size_lg2));
722
723 int ok_block_aligned = 0;
724 int ok_dealloc_once = 0;
725
726 if (ok_contains) {
727 const int sb_id = d >> m_sb_size_lg2;
728
729 // State array for the superblock.
730 volatile uint32_t *const sb_state_array =
731 m_sb_state_array + (sb_id * m_sb_state_size);
732
733 const uint32_t block_state = (*sb_state_array) & state_header_mask;
734 const uint32_t block_size_lg2 =
735 m_sb_size_lg2 - (block_state >> state_shift);
736
737 ok_block_aligned = 0 == (d & ((1UL << block_size_lg2) - 1));
738
739 if (ok_block_aligned) {
740 // Map address to block's bit
741 // mask into superblock and then shift down for block index
742
743 const uint32_t bit =
744 (d & (ptrdiff_t(1LU << m_sb_size_lg2) - 1)) >> block_size_lg2;
745
746 const int result = CB::release(sb_state_array, bit, block_state);
747
748 ok_dealloc_once = 0 <= result;
749 }
750 }
751
752 if (!ok_contains || !ok_block_aligned || !ok_dealloc_once) {
753 Kokkos::abort("Kokkos MemoryPool::deallocate given erroneous pointer");
754 }
755 }
756 // end deallocate
757 //--------------------------------------------------------------------------
758
759 KOKKOS_INLINE_FUNCTION
760 int number_of_superblocks() const noexcept { return m_sb_count; }
761
762 KOKKOS_INLINE_FUNCTION
763 void superblock_state(int sb_id, int &block_size, int &block_count_capacity,
764 int &block_count_used) const noexcept {
765 block_size = 0;
766 block_count_capacity = 0;
767 block_count_used = 0;
768
769 bool can_access_state_array = []() {
770 KOKKOS_IF_ON_HOST(
771 (return SpaceAccessibility<DefaultHostExecutionSpace,
772 base_memory_space>::accessible;))
773 KOKKOS_IF_ON_DEVICE(
774 (return SpaceAccessibility<DefaultExecutionSpace,
775 base_memory_space>::accessible;))
776 }();
777
778 if (can_access_state_array) {
779 // Can access the state array
780
781 const uint32_t state =
782 ((uint32_t volatile *)m_sb_state_array)[sb_id * m_sb_state_size];
783
784 const uint32_t block_count_lg2 = state >> state_shift;
785 const uint32_t block_used = state & state_used_mask;
786
787 block_size = 1LU << (m_sb_size_lg2 - block_count_lg2);
788 block_count_capacity = 1LU << block_count_lg2;
789 block_count_used = block_used;
790 }
791 }
792};
793
794} // namespace Kokkos
795
796#endif /* #ifndef KOKKOS_MEMORYPOOL_HPP */
Atomic functions.
Declaration of parallel operators.
A thread safe view to a bitset.
Memory management for host memory.
void deallocate(void *const arg_alloc_ptr, const size_t arg_alloc_size) const
Deallocate untracked memory in the space.
void * allocate(const size_t arg_alloc_size) const
Allocate untracked memory in the space.
Access relationship between DstMemorySpace and SrcMemorySpace.