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. 2.0
6 // Copyright (2014) Sandia Corporation
7 //
8 // Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
9 // the U.S. Government retains certain rights in this software.
10 //
11 // Redistribution and use in source and binary forms, with or without
12 // modification, are permitted provided that the following conditions are
13 // met:
14 //
15 // 1. Redistributions of source code must retain the above copyright
16 // notice, this list of conditions and the following disclaimer.
17 //
18 // 2. Redistributions in binary form must reproduce the above copyright
19 // notice, this list of conditions and the following disclaimer in the
20 // documentation and/or other materials provided with the distribution.
21 //
22 // 3. Neither the name of the Corporation nor the names of the
23 // contributors may be used to endorse or promote products derived from
24 // this software without specific prior written permission.
25 //
26 // THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
27 // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
28 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
29 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
30 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
31 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
32 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
33 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
34 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
35 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
36 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
37 //
38 // Questions? Contact Christian R. Trott (crtrott@sandia.gov)
39 //
40 // ************************************************************************
41 //@HEADER
42 */
43 
44 #ifndef KOKKOS_MEMORYPOOL_HPP
45 #define KOKKOS_MEMORYPOOL_HPP
46 
47 #include <Kokkos_Core_fwd.hpp>
48 #include <Kokkos_Parallel.hpp>
49 #include <Kokkos_Atomic.hpp>
50 #include <impl/Kokkos_ConcurrentBitset.hpp>
51 #include <impl/Kokkos_Error.hpp>
52 #include <impl/Kokkos_SharedAlloc.hpp>
53 
54 namespace Kokkos {
55 namespace Impl {
56 /* Report violation of size constraints:
57  * min_block_alloc_size <= max_block_alloc_size
58  * max_block_alloc_size <= min_superblock_size
59  * min_superblock_size <= max_superblock_size
60  * min_superblock_size <= min_total_alloc_size
61  * min_superblock_size <= min_block_alloc_size *
62  * max_block_per_superblock
63  */
64 void memory_pool_bounds_verification
65  ( size_t min_block_alloc_size
66  , size_t max_block_alloc_size
67  , size_t min_superblock_size
68  , size_t max_superblock_size
69  , size_t max_block_per_superblock
70  , size_t min_total_alloc_size
71  );
72 }
73 }
74 
75 namespace Kokkos {
76 
77 template< typename DeviceType >
78 class MemoryPool {
79 private:
80 
81  typedef typename Kokkos::Impl::concurrent_bitset CB ;
82 
83  enum : uint32_t { bits_per_int_lg2 = CB::bits_per_int_lg2 };
84  enum : uint32_t { state_shift = CB::state_shift };
85  enum : uint32_t { state_used_mask = CB::state_used_mask };
86  enum : uint32_t { state_header_mask = CB::state_header_mask };
87  enum : uint32_t { max_bit_count_lg2 = CB::max_bit_count_lg2 };
88  enum : uint32_t { max_bit_count = CB::max_bit_count };
89 
90  enum : uint32_t { HINT_PER_BLOCK_SIZE = 2 };
91 
92  /* Each superblock has a concurrent bitset state
93  * which is an array of uint32_t integers.
94  * [ { block_count_lg2 : state_shift bits
95  * , used_block_count : ( 32 - state_shift ) bits
96  * }
97  * , { block allocation bit set }* ]
98  *
99  * As superblocks are assigned (allocated) to a block size
100  * and released (deallocated) back to empty the superblock state
101  * is concurrently updated.
102  */
103 
104  /* Mapping between block_size <-> block_state
105  *
106  * block_state = ( m_sb_size_lg2 - block_size_lg2 ) << state_shift
107  * block_size = m_sb_size_lg2 - ( block_state >> state_shift )
108  *
109  * Thus A_block_size < B_block_size <=> A_block_state > B_block_state
110  */
111 
112  typedef typename DeviceType::memory_space base_memory_space ;
113 
114  enum { accessible =
116  , base_memory_space >::accessible };
117 
118  typedef Kokkos::Impl::SharedAllocationTracker Tracker ;
119  typedef Kokkos::Impl::SharedAllocationRecord
120  < base_memory_space > Record ;
121 
122  Tracker m_tracker ;
123  uint32_t * m_sb_state_array ;
124  uint32_t m_sb_state_size ;
125  uint32_t m_sb_size_lg2 ;
126  uint32_t m_max_block_size_lg2 ;
127  uint32_t m_min_block_size_lg2 ;
128  int32_t m_sb_count ;
129  int32_t m_hint_offset ; // Offset to K * #block_size array of hints
130  int32_t m_data_offset ; // Offset to 0th superblock data
131  int32_t m_unused_padding ;
132 
133 public:
134 
136  enum : uint32_t { max_superblock_size = 1LU << 31 /* 2 gigabytes */ };
137  enum : uint32_t { max_block_per_superblock = max_bit_count };
138 
139  //--------------------------------------------------------------------------
140 
141  KOKKOS_INLINE_FUNCTION
142  size_t capacity() const noexcept
143  { return size_t(m_sb_count) << m_sb_size_lg2 ; }
144 
145  KOKKOS_INLINE_FUNCTION
146  size_t min_block_size() const noexcept
147  { return ( 1LU << m_min_block_size_lg2 ); }
148 
149  KOKKOS_INLINE_FUNCTION
150  size_t max_block_size() const noexcept
151  { return ( 1LU << m_max_block_size_lg2 ); }
152 
153  struct usage_statistics {
154  size_t capacity_bytes ;
155  size_t superblock_bytes ;
156  size_t max_block_bytes ;
157  size_t min_block_bytes ;
158  size_t capacity_superblocks ;
159  size_t consumed_superblocks ;
160  size_t consumed_blocks ;
161  size_t consumed_bytes ;
162  size_t reserved_blocks ;
163  size_t reserved_bytes ;
164  };
165 
166  void get_usage_statistics( usage_statistics & stats ) const
167  {
168  Kokkos::HostSpace host ;
169 
170  const size_t alloc_size = m_hint_offset * sizeof(uint32_t);
171 
172  uint32_t * const sb_state_array =
173  accessible ? m_sb_state_array : (uint32_t *) host.allocate(alloc_size);
174 
175  if ( ! accessible ) {
176  Kokkos::Impl::DeepCopy< Kokkos::HostSpace , base_memory_space >
177  ( sb_state_array , m_sb_state_array , alloc_size );
178  }
179 
180  stats.superblock_bytes = ( 1LU << m_sb_size_lg2 );
181  stats.max_block_bytes = ( 1LU << m_max_block_size_lg2 );
182  stats.min_block_bytes = ( 1LU << m_min_block_size_lg2 );
183  stats.capacity_bytes = stats.superblock_bytes * m_sb_count ;
184  stats.capacity_superblocks = m_sb_count ;
185  stats.consumed_superblocks = 0 ;
186  stats.consumed_blocks = 0 ;
187  stats.consumed_bytes = 0 ;
188  stats.reserved_blocks = 0 ;
189  stats.reserved_bytes = 0 ;
190 
191  const uint32_t * sb_state_ptr = sb_state_array ;
192 
193  for ( int32_t i = 0 ; i < m_sb_count
194  ; ++i , sb_state_ptr += m_sb_state_size ) {
195 
196  const uint32_t block_count_lg2 = (*sb_state_ptr) >> state_shift ;
197 
198  if ( block_count_lg2 ) {
199  const uint32_t block_count = 1u << block_count_lg2 ;
200  const uint32_t block_size_lg2 = m_sb_size_lg2 - block_count_lg2 ;
201  const uint32_t block_size = 1u << block_size_lg2 ;
202  const uint32_t block_used = (*sb_state_ptr) & state_used_mask ;
203 
204  stats.consumed_superblocks++ ;
205  stats.consumed_blocks += block_used ;
206  stats.consumed_bytes += block_used * block_size ;
207  stats.reserved_blocks += block_count - block_used ;
208  stats.reserved_bytes += (block_count - block_used ) * block_size ;
209  }
210  }
211 
212  if ( ! accessible ) {
213  host.deallocate( sb_state_array, alloc_size );
214  }
215  }
216 
217  void print_state( std::ostream & s ) const
218  {
219  Kokkos::HostSpace host ;
220 
221  const size_t alloc_size = m_hint_offset * sizeof(uint32_t);
222 
223  uint32_t * const sb_state_array =
224  accessible ? m_sb_state_array : (uint32_t *) host.allocate(alloc_size);
225 
226  if ( ! accessible ) {
227  Kokkos::Impl::DeepCopy< Kokkos::HostSpace , base_memory_space >
228  ( sb_state_array , m_sb_state_array , alloc_size );
229  }
230 
231  const uint32_t * sb_state_ptr = sb_state_array ;
232 
233  s << "pool_size(" << ( size_t(m_sb_count) << m_sb_size_lg2 ) << ")"
234  << " superblock_size(" << ( 1LU << m_sb_size_lg2 ) << ")" << std::endl ;
235 
236  for ( int32_t i = 0 ; i < m_sb_count
237  ; ++i , sb_state_ptr += m_sb_state_size ) {
238 
239  if ( *sb_state_ptr ) {
240 
241  const uint32_t block_count_lg2 = (*sb_state_ptr) >> state_shift ;
242  const uint32_t block_size_lg2 = m_sb_size_lg2 - block_count_lg2 ;
243  const uint32_t block_count = 1u << block_count_lg2 ;
244  const uint32_t block_used = (*sb_state_ptr) & state_used_mask ;
245 
246  s << "Superblock[ " << i << " / " << m_sb_count << " ] {"
247  << " block_size(" << ( 1 << block_size_lg2 ) << ")"
248  << " block_count( " << block_used
249  << " / " << block_count << " )"
250  << std::endl ;
251  }
252  }
253 
254  if ( ! accessible ) {
255  host.deallocate( sb_state_array, alloc_size );
256  }
257  }
258 
259  //--------------------------------------------------------------------------
260 
261 #ifdef KOKKOS_CUDA_9_DEFAULTED_BUG_WORKAROUND
262  KOKKOS_INLINE_FUNCTION MemoryPool( MemoryPool && rhs )
263  : m_tracker(std::move(rhs.m_tracker))
264  , m_sb_state_array(std::move(rhs.m_sb_state_array))
265  , m_sb_state_size(std::move(rhs.m_sb_state_size))
266  , m_sb_size_lg2(std::move(rhs.m_sb_size_lg2))
267  , m_max_block_size_lg2(std::move(rhs.m_max_block_size_lg2))
268  , m_min_block_size_lg2(std::move(rhs.m_min_block_size_lg2))
269  , m_sb_count(std::move(rhs.m_sb_count))
270  , m_hint_offset(std::move(rhs.m_hint_offset))
271  , m_data_offset(std::move(rhs.m_data_offset))
272  {
273  }
274  KOKKOS_INLINE_FUNCTION MemoryPool( const MemoryPool & rhs )
275  : m_tracker(rhs.m_tracker)
276  , m_sb_state_array(rhs.m_sb_state_array)
277  , m_sb_state_size(rhs.m_sb_state_size)
278  , m_sb_size_lg2(rhs.m_sb_size_lg2)
279  , m_max_block_size_lg2(rhs.m_max_block_size_lg2)
280  , m_min_block_size_lg2(rhs.m_min_block_size_lg2)
281  , m_sb_count(rhs.m_sb_count)
282  , m_hint_offset(rhs.m_hint_offset)
283  , m_data_offset(rhs.m_data_offset)
284  {
285  }
286  KOKKOS_INLINE_FUNCTION MemoryPool & operator = ( MemoryPool && rhs )
287  {
288  m_tracker = std::move(rhs.m_tracker);
289  m_sb_state_array = std::move(rhs.m_sb_state_array);
290  m_sb_state_size = std::move(rhs.m_sb_state_size);
291  m_sb_size_lg2 = std::move(rhs.m_sb_size_lg2);
292  m_max_block_size_lg2 = std::move(rhs.m_max_block_size_lg2);
293  m_min_block_size_lg2 = std::move(rhs.m_min_block_size_lg2);
294  m_sb_count = std::move(rhs.m_sb_count);
295  m_hint_offset = std::move(rhs.m_hint_offset);
296  m_data_offset = std::move(rhs.m_data_offset);
297  return *this;
298  }
299  KOKKOS_INLINE_FUNCTION MemoryPool & operator = ( const MemoryPool & rhs )
300  {
301  m_tracker = rhs.m_tracker;
302  m_sb_state_array = rhs.m_sb_state_array;
303  m_sb_state_size = rhs.m_sb_state_size;
304  m_sb_size_lg2 = rhs.m_sb_size_lg2;
305  m_max_block_size_lg2 = rhs.m_max_block_size_lg2;
306  m_min_block_size_lg2 = rhs.m_min_block_size_lg2;
307  m_sb_count = rhs.m_sb_count;
308  m_hint_offset = rhs.m_hint_offset;
309  m_data_offset = rhs.m_data_offset;
310  return *this;
311  }
312 #else
313  KOKKOS_INLINE_FUNCTION MemoryPool( MemoryPool && ) = default ;
314  KOKKOS_INLINE_FUNCTION MemoryPool( const MemoryPool & ) = default ;
315  KOKKOS_INLINE_FUNCTION MemoryPool & operator = ( MemoryPool && ) = default ;
316  KOKKOS_INLINE_FUNCTION MemoryPool & operator = ( const MemoryPool & ) = default ;
317 #endif
318 
319  KOKKOS_INLINE_FUNCTION MemoryPool()
320  : m_tracker()
321  , m_sb_state_array(0)
322  , m_sb_state_size(0)
323  , m_sb_size_lg2(0)
324  , m_max_block_size_lg2(0)
325  , m_min_block_size_lg2(0)
326  , m_sb_count(0)
327  , m_hint_offset(0)
328  , m_data_offset(0)
329  , m_unused_padding(0)
330  {}
331 
346  MemoryPool( const base_memory_space & memspace
347  , const size_t min_total_alloc_size
348  , size_t min_block_alloc_size = 0
349  , size_t max_block_alloc_size = 0
350  , size_t min_superblock_size = 0
351  )
352  : m_tracker()
353  , m_sb_state_array(0)
354  , m_sb_state_size(0)
355  , m_sb_size_lg2(0)
356  , m_max_block_size_lg2(0)
357  , m_min_block_size_lg2(0)
358  , m_sb_count(0)
359  , m_hint_offset(0)
360  , m_data_offset(0)
361  , m_unused_padding(0)
362  {
363  const uint32_t int_align_lg2 = 3 ; /* align as int[8] */
364  const uint32_t int_align_mask = ( 1u << int_align_lg2 ) - 1 ;
365  const uint32_t default_min_block_size = 1u << 6 ; /* 64 bytes */
366  const uint32_t default_max_block_size = 1u << 12 ;/* 4k bytes */
367  const uint32_t default_min_superblock_size = 1u << 20 ;/* 1M bytes */
368 
369  //--------------------------------------------------
370  // Default block and superblock sizes:
371 
372  if ( 0 == min_block_alloc_size ) {
373  // Default all sizes:
374 
375  min_superblock_size =
376  std::min( size_t(default_min_superblock_size)
377  , min_total_alloc_size );
378 
379  min_block_alloc_size =
380  std::min( size_t(default_min_block_size)
381  , min_superblock_size );
382 
383  max_block_alloc_size =
384  std::min( size_t(default_max_block_size)
385  , min_superblock_size );
386  }
387  else if ( 0 == min_superblock_size ) {
388 
389  // Choose superblock size as minimum of:
390  // max_block_per_superblock * min_block_size
391  // max_superblock_size
392  // min_total_alloc_size
393 
394  const size_t max_superblock =
395  min_block_alloc_size * max_block_per_superblock ;
396 
397  min_superblock_size =
398  std::min( max_superblock ,
399  std::min( size_t(max_superblock_size)
400  , min_total_alloc_size ) );
401  }
402 
403  if ( 0 == max_block_alloc_size ) {
404  max_block_alloc_size = min_superblock_size ;
405  }
406 
407  //--------------------------------------------------
408 
409  /* Enforce size constraints:
410  * min_block_alloc_size <= max_block_alloc_size
411  * max_block_alloc_size <= min_superblock_size
412  * min_superblock_size <= max_superblock_size
413  * min_superblock_size <= min_total_alloc_size
414  * min_superblock_size <= min_block_alloc_size *
415  * max_block_per_superblock
416  */
417 
418  Kokkos::Impl::memory_pool_bounds_verification
419  ( min_block_alloc_size
420  , max_block_alloc_size
421  , min_superblock_size
422  , max_superblock_size
423  , max_block_per_superblock
424  , min_total_alloc_size
425  );
426 
427  //--------------------------------------------------
428  // Block and superblock size is power of two:
429  // Maximum value is 'max_superblock_size'
430 
431  m_min_block_size_lg2 =
432  Kokkos::Impl::integral_power_of_two_that_contains(min_block_alloc_size);
433 
434  m_max_block_size_lg2 =
435  Kokkos::Impl::integral_power_of_two_that_contains(max_block_alloc_size);
436 
437  m_sb_size_lg2 =
438  Kokkos::Impl::integral_power_of_two_that_contains(min_superblock_size);
439 
440  {
441  // number of superblocks is multiple of superblock size that
442  // can hold min_total_alloc_size.
443 
444  const uint64_t sb_size_mask = ( 1LU << m_sb_size_lg2 ) - 1 ;
445 
446  m_sb_count = ( min_total_alloc_size + sb_size_mask ) >> m_sb_size_lg2 ;
447  }
448 
449  {
450  // Any superblock can be assigned to the smallest size block
451  // Size the block bitset to maximum number of blocks
452 
453  const uint32_t max_block_count_lg2 =
454  m_sb_size_lg2 - m_min_block_size_lg2 ;
455 
456  m_sb_state_size =
457  ( CB::buffer_bound_lg2( max_block_count_lg2 ) + int_align_mask ) & ~int_align_mask ;
458  }
459 
460  // Array of all superblock states
461 
462  const size_t all_sb_state_size =
463  ( m_sb_count * m_sb_state_size + int_align_mask ) & ~int_align_mask ;
464 
465  // Number of block sizes
466 
467  const int32_t number_block_sizes =
468  1 + m_max_block_size_lg2 - m_min_block_size_lg2 ;
469 
470  // Array length for possible block sizes
471  // Hint array is one uint32_t per block size
472 
473  const int32_t block_size_array_size =
474  ( number_block_sizes + int_align_mask ) & ~int_align_mask ;
475 
476  m_hint_offset = all_sb_state_size ;
477  m_data_offset = m_hint_offset +
478  block_size_array_size * HINT_PER_BLOCK_SIZE ;
479 
480  // Allocation:
481 
482  const size_t header_size = m_data_offset * sizeof(uint32_t);
483  const size_t alloc_size = header_size +
484  ( size_t(m_sb_count) << m_sb_size_lg2 );
485 
486  Record * rec = Record::allocate( memspace , "MemoryPool" , alloc_size );
487 
488  m_tracker.assign_allocated_record_to_uninitialized( rec );
489 
490  m_sb_state_array = (uint32_t *) rec->data();
491 
492  Kokkos::HostSpace host ;
493 
494  uint32_t * const sb_state_array =
495  accessible ? m_sb_state_array
496  : (uint32_t *) host.allocate(header_size);
497 
498  for ( int32_t i = 0 ; i < m_data_offset ; ++i ) sb_state_array[i] = 0 ;
499 
500  // Initial assignment of empty superblocks to block sizes:
501 
502  for ( int32_t i = 0 ; i < number_block_sizes ; ++i ) {
503  const uint32_t block_size_lg2 = i + m_min_block_size_lg2 ;
504  const uint32_t block_count_lg2 = m_sb_size_lg2 - block_size_lg2 ;
505  const uint32_t block_state = block_count_lg2 << state_shift ;
506  const uint32_t hint_begin = m_hint_offset + i * HINT_PER_BLOCK_SIZE ;
507 
508  // for block size index 'i':
509  // sb_id_hint = sb_state_array[ hint_begin ];
510  // sb_id_begin = sb_state_array[ hint_begin + 1 ];
511 
512  const int32_t jbeg = ( i * m_sb_count ) / number_block_sizes ;
513  const int32_t jend = ( ( i + 1 ) * m_sb_count ) / number_block_sizes ;
514 
515  sb_state_array[ hint_begin ] = uint32_t(jbeg);
516  sb_state_array[ hint_begin + 1 ] = uint32_t(jbeg);
517 
518  for ( int32_t j = jbeg ; j < jend ; ++j ) {
519  sb_state_array[ j * m_sb_state_size ] = block_state ;
520  }
521  }
522 
523  // Write out initialized state:
524 
525  if ( ! accessible ) {
526  Kokkos::Impl::DeepCopy< base_memory_space , Kokkos::HostSpace >
527  ( m_sb_state_array , sb_state_array , header_size );
528 
529  host.deallocate( sb_state_array, header_size );
530  }
531  else {
532  Kokkos::memory_fence();
533  }
534  }
535 
536  //--------------------------------------------------------------------------
537 
538 private:
539 
540  /* Given a size 'n' get the block size in which it can be allocated.
541  * Restrict lower bound to minimum block size.
542  */
543  KOKKOS_FORCEINLINE_FUNCTION
544  uint32_t get_block_size_lg2( uint32_t n ) const noexcept
545  {
546  const unsigned i = Kokkos::Impl::integral_power_of_two_that_contains( n );
547 
548  return i < m_min_block_size_lg2 ? m_min_block_size_lg2 : i ;
549  }
550 
551 public:
552 
553  /* Return 0 for invalid block size */
554  KOKKOS_INLINE_FUNCTION
555  uint32_t allocate_block_size( uint64_t alloc_size ) const noexcept
556  {
557  return alloc_size <= (1UL << m_max_block_size_lg2)
558  ? ( 1UL << get_block_size_lg2( uint32_t(alloc_size) ) )
559  : 0 ;
560  }
561 
562  //--------------------------------------------------------------------------
572  KOKKOS_FUNCTION
573  void * allocate( size_t alloc_size
574  , int32_t attempt_limit = 1 ) const noexcept
575  {
576  if ( size_t(1LU << m_max_block_size_lg2) < alloc_size ) {
577  Kokkos::abort("Kokkos MemoryPool allocation request exceeded specified maximum allocation size");
578  }
579 
580  if ( 0 == alloc_size ) return (void*) 0 ;
581 
582  void * p = 0 ;
583 
584  const uint32_t block_size_lg2 = get_block_size_lg2( alloc_size );
585 
586  // Allocation will fit within a superblock
587  // that has block sizes ( 1 << block_size_lg2 )
588 
589  const uint32_t block_count_lg2 = m_sb_size_lg2 - block_size_lg2 ;
590  const uint32_t block_state = block_count_lg2 << state_shift ;
591  const uint32_t block_count = 1u << block_count_lg2 ;
592 
593  // Superblock hints for this block size:
594  // hint_sb_id_ptr[0] is the dynamically changing hint
595  // hint_sb_id_ptr[1] is the static start point
596 
597  volatile uint32_t * const hint_sb_id_ptr
598  = m_sb_state_array /* memory pool state array */
599  + m_hint_offset /* offset to hint portion of array */
600  + HINT_PER_BLOCK_SIZE /* number of hints per block size */
601  * ( block_size_lg2 - m_min_block_size_lg2 ); /* block size id */
602 
603  const int32_t sb_id_begin = int32_t( hint_sb_id_ptr[1] );
604 
605  // Fast query clock register 'tic' to pseudo-randomize
606  // the guess for which block within a superblock should
607  // be claimed. If not available then a search occurs.
608 
609  const uint32_t block_id_hint =
610  (uint32_t)( Kokkos::Impl::clock_tic()
611 #if defined( KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_CUDA )
612  // Spread out potentially concurrent access
613  // by threads within a warp or thread block.
614  + ( threadIdx.x + blockDim.x * threadIdx.y )
615 #endif
616  );
617 
618  // expected state of superblock for allocation
619  uint32_t sb_state = block_state ;
620 
621  int32_t sb_id = -1 ;
622 
623  volatile uint32_t * sb_state_array = 0 ;
624 
625  while ( attempt_limit ) {
626 
627  int32_t hint_sb_id = -1 ;
628 
629  if ( sb_id < 0 ) {
630 
631  // No superblock specified, try the hint for this block size
632 
633  sb_id = hint_sb_id = int32_t( *hint_sb_id_ptr );
634 
635  sb_state_array = m_sb_state_array + ( sb_id * m_sb_state_size );
636  }
637 
638  // Require:
639  // 0 <= sb_id
640  // sb_state_array == m_sb_state_array + m_sb_state_size * sb_id
641 
642  if ( sb_state == ( state_header_mask & *sb_state_array ) ) {
643 
644  // This superblock state is as expected, for the moment.
645  // Attempt to claim a bit. The attempt updates the state
646  // so have already made sure the state header is as expected.
647 
648  const uint32_t count_lg2 = sb_state >> state_shift ;
649  const uint32_t mask = ( 1u << count_lg2 ) - 1 ;
650 
651  const Kokkos::pair<int,int> result =
652  CB::acquire_bounded_lg2( sb_state_array
653  , count_lg2
654  , block_id_hint & mask
655  , sb_state
656  );
657 
658  // If result.first < 0 then failed to acquire
659  // due to either full or buffer was wrong state.
660  // Could be wrong state if a deallocation raced the
661  // superblock to empty before the acquire could succeed.
662 
663  if ( 0 <= result.first ) { // acquired a bit
664 
665  const uint32_t size_lg2 = m_sb_size_lg2 - count_lg2 ;
666 
667  // Set the allocated block pointer
668 
669  p = ((char*)( m_sb_state_array + m_data_offset ))
670  + ( uint64_t(sb_id) << m_sb_size_lg2 ) // superblock memory
671  + ( uint64_t(result.first) << size_lg2 ); // block memory
672 
673 #if 0
674  printf( " MemoryPool(0x%lx) pointer(0x%lx) allocate(%lu) sb_id(%d) sb_state(0x%x) block_size(%d) block_capacity(%d) block_id(%d) block_claimed(%d)\n"
675  , (uintptr_t)m_sb_state_array
676  , (uintptr_t)p
677  , alloc_size
678  , sb_id
679  , sb_state
680  , (1u << size_lg2)
681  , (1u << count_lg2)
682  , result.first
683  , result.second );
684 #endif
685 
686  break ; // Success
687  }
688  }
689  //------------------------------------------------------------------
690  // Arrive here if failed to acquire a block.
691  // Must find a new superblock.
692 
693  // Start searching at designated index for this block size.
694  // Look for superblock that, in preferential order,
695  // 1) part-full superblock of this block size
696  // 2) empty superblock to claim for this block size
697  // 3) part-full superblock of the next larger block size
698 
699  sb_state = block_state ; // Expect to find the desired state
700  sb_id = -1 ;
701 
702  bool update_hint = false ;
703  int32_t sb_id_empty = -1 ;
704  int32_t sb_id_large = -1 ;
705  uint32_t sb_state_large = 0 ;
706 
707  sb_state_array = m_sb_state_array + sb_id_begin * m_sb_state_size ;
708 
709  for ( int32_t i = 0 , id = sb_id_begin ; i < m_sb_count ; ++i ) {
710 
711  // Query state of the candidate superblock.
712  // Note that the state may change at any moment
713  // as concurrent allocations and deallocations occur.
714 
715  const uint32_t full_state = *sb_state_array ;
716  const uint32_t used = full_state & state_used_mask ;
717  const uint32_t state = full_state & state_header_mask ;
718 
719  if ( state == block_state ) {
720 
721  // Superblock is assigned to this block size
722 
723  if ( used < block_count ) {
724 
725  // There is room to allocate one block
726 
727  sb_id = id ;
728 
729  // Is there room to allocate more than one block?
730 
731  update_hint = used + 1 < block_count ;
732 
733  break ;
734  }
735  }
736  else if ( 0 == used ) {
737 
738  // Superblock is empty
739 
740  if ( -1 == sb_id_empty ) {
741 
742  // Superblock is not assigned to this block size
743  // and is the first empty superblock encountered.
744  // Save this id to use if a partfull superblock is not found.
745 
746  sb_id_empty = id ;
747  }
748  }
749  else if ( ( -1 == sb_id_empty /* have not found an empty */ ) &&
750  ( -1 == sb_id_large /* have not found a larger */ ) &&
751  ( state < block_state /* a larger block */ ) &&
752  // is not full:
753  ( used < ( 1u << ( state >> state_shift ) ) ) ) {
754  // First superblock encountered that is
755  // larger than this block size and
756  // has room for an allocation.
757  // Save this id to use of partfull or empty superblock not found
758  sb_id_large = id ;
759  sb_state_large = state ;
760  }
761 
762  // Iterate around the superblock array:
763 
764  if ( ++id < m_sb_count ) {
765  sb_state_array += m_sb_state_size ;
766  }
767  else {
768  id = 0 ;
769  sb_state_array = m_sb_state_array ;
770  }
771  }
772 
773  // printf(" search m_sb_count(%d) sb_id(%d) sb_id_empty(%d) sb_id_large(%d)\n" , m_sb_count , sb_id , sb_id_empty , sb_id_large);
774 
775  if ( sb_id < 0 ) {
776 
777  // Did not find a partfull superblock for this block size.
778 
779  if ( 0 <= sb_id_empty ) {
780 
781  // Found first empty superblock following designated superblock
782  // Attempt to claim it for this block size.
783  // If the claim fails assume that another thread claimed it
784  // for this block size and try to use it anyway,
785  // but do not update hint.
786 
787  sb_id = sb_id_empty ;
788 
789  sb_state_array = m_sb_state_array + ( sb_id * m_sb_state_size );
790 
791  // If successfully changed assignment of empty superblock 'sb_id'
792  // to this block_size then update the hint.
793 
794  const uint32_t state_empty = state_header_mask & *sb_state_array ;
795 
796  // If this thread claims the empty block then update the hint
797  update_hint =
798  state_empty ==
799  Kokkos::atomic_compare_exchange
800  (sb_state_array,state_empty,block_state);
801  }
802  else if ( 0 <= sb_id_large ) {
803 
804  // Found a larger superblock with space available
805 
806  sb_id = sb_id_large ;
807  sb_state = sb_state_large ;
808 
809  sb_state_array = m_sb_state_array + ( sb_id * m_sb_state_size );
810  }
811  else {
812  // Did not find a potentially usable superblock
813  --attempt_limit ;
814  }
815  }
816 
817  if ( update_hint ) {
818  Kokkos::atomic_compare_exchange
819  ( hint_sb_id_ptr , uint32_t(hint_sb_id) , uint32_t(sb_id) );
820  }
821  } // end allocation attempt loop
822  //--------------------------------------------------------------------
823 
824  return p ;
825  }
826  // end allocate
827  //--------------------------------------------------------------------------
828 
835  KOKKOS_INLINE_FUNCTION
836  void deallocate( void * p , size_t /* alloc_size */ ) const noexcept
837  {
838  if ( 0 == p ) return ;
839 
840  // Determine which superblock and block
841  const ptrdiff_t d =
842  ((char*)p) - ((char*)( m_sb_state_array + m_data_offset ));
843 
844  // Verify contained within the memory pool's superblocks:
845  const int ok_contains =
846  ( 0 <= d ) && ( size_t(d) < ( size_t(m_sb_count) << m_sb_size_lg2 ) );
847 
848  int ok_block_aligned = 0 ;
849  int ok_dealloc_once = 0 ;
850 
851  if ( ok_contains ) {
852 
853  const int sb_id = d >> m_sb_size_lg2 ;
854 
855  // State array for the superblock.
856  volatile uint32_t * const sb_state_array =
857  m_sb_state_array + ( sb_id * m_sb_state_size );
858 
859  const uint32_t block_state = (*sb_state_array) & state_header_mask ;
860  const uint32_t block_size_lg2 =
861  m_sb_size_lg2 - ( block_state >> state_shift );
862 
863  ok_block_aligned = 0 == ( d & ( ( 1UL << block_size_lg2 ) - 1 ) );
864 
865  if ( ok_block_aligned ) {
866 
867  // Map address to block's bit
868  // mask into superblock and then shift down for block index
869 
870  const uint32_t bit =
871  ( d & ( ptrdiff_t( 1LU << m_sb_size_lg2 ) - 1 ) ) >> block_size_lg2 ;
872 
873  const int result =
874  CB::release( sb_state_array , bit , block_state );
875 
876  ok_dealloc_once = 0 <= result ;
877 
878 #if 0
879  printf( " MemoryPool(0x%lx) pointer(0x%lx) deallocate sb_id(%d) block_size(%d) block_capacity(%d) block_id(%d) block_claimed(%d)\n"
880  , (uintptr_t)m_sb_state_array
881  , (uintptr_t)p
882  , sb_id
883  , (1u << block_size_lg2)
884  , (1u << (m_sb_size_lg2 - block_size_lg2))
885  , bit
886  , result );
887 #endif
888  }
889  }
890 
891  if ( ! ok_contains || ! ok_block_aligned || ! ok_dealloc_once ) {
892 #if 0
893  printf( " MemoryPool(0x%lx) pointer(0x%lx) deallocate ok_contains(%d) ok_block_aligned(%d) ok_dealloc_once(%d)\n"
894  , (uintptr_t)m_sb_state_array
895  , (uintptr_t)p
896  , int(ok_contains)
897  , int(ok_block_aligned)
898  , int(ok_dealloc_once) );
899 #endif
900  Kokkos::abort("Kokkos MemoryPool::deallocate given erroneous pointer");
901  }
902  }
903  // end deallocate
904  //--------------------------------------------------------------------------
905 
906  KOKKOS_INLINE_FUNCTION
907  int number_of_superblocks() const noexcept { return m_sb_count ; }
908 
909  KOKKOS_INLINE_FUNCTION
910  void superblock_state( int sb_id
911  , int & block_size
912  , int & block_count_capacity
913  , int & block_count_used ) const noexcept
914  {
915  block_size = 0 ;
916  block_count_capacity = 0 ;
917  block_count_used = 0 ;
918 
920  < Kokkos::Impl::ActiveExecutionMemorySpace
921  , base_memory_space >::accessible ) {
922  // Can access the state array
923 
924  const uint32_t state =
925  ((uint32_t volatile *)m_sb_state_array)[sb_id*m_sb_state_size];
926 
927  const uint32_t block_count_lg2 = state >> state_shift ;
928  const uint32_t block_used = state & state_used_mask ;
929 
930  block_size = 1LU << ( m_sb_size_lg2 - block_count_lg2 );
931  block_count_capacity = 1LU << block_count_lg2 ;
932  block_count_used = block_used ;
933  }
934  }
935 };
936 
937 } // namespace Kokkos
938 
939 #endif /* #ifndef KOKKOS_MEMORYPOOL_HPP */
940 
void * allocate(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:64
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:72
Memory management for host memory.
Declaration of parallel operators.
Atomic functions.
second_type second
The second element of the pair.
Definition: Kokkos_Pair.hpp:74
Access relationship between DstMemorySpace and SrcMemorySpace.