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