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