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