Page Menu
Home
c4science
Search
Configure Global Search
Log In
Files
F120431782
Kokkos_MemoryPool.hpp
No One
Temporary
Actions
Download File
Edit File
Delete File
View Transforms
Subscribe
Mute Notifications
Award Token
Subscribers
None
File Metadata
Details
File Info
Storage
Attached
Created
Fri, Jul 4, 08:19
Size
58 KB
Mime Type
text/x-c++
Expires
Sun, Jul 6, 08:19 (1 d, 23 h)
Engine
blob
Format
Raw Data
Handle
27187992
Attached To
rLAMMPS lammps
Kokkos_MemoryPool.hpp
View Options
/*
//@HEADER
// ************************************************************************
//
// Kokkos v. 2.0
// Copyright (2014) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact H. Carter Edwards (hcedwar@sandia.gov)
//
// ************************************************************************
//@HEADER
*/
#ifndef KOKKOS_MEMORYPOOL_HPP
#define KOKKOS_MEMORYPOOL_HPP
#include <Kokkos_Core_fwd.hpp>
#include <Kokkos_Parallel.hpp>
#include <Kokkos_Atomic.hpp>
#include <impl/Kokkos_BitOps.hpp>
#include <impl/Kokkos_Error.hpp>
#include <impl/Kokkos_SharedAlloc.hpp>
#include <limits>
#include <algorithm>
#include <chrono>
// How should errors be handled? In general, production code should return a
// value indicating failure so the user can decide how the error is handled.
// While experimental, code can abort instead. If KOKKOS_ENABLE_MEMPOOL_PRINTERR is
// defined, the code will abort with an error message. Otherwise, the code will
// return with a value indicating failure when possible, or do nothing instead.
//#define KOKKOS_ENABLE_MEMPOOL_PRINTERR
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_INFO
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_CONSTRUCTOR_INFO
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_BLOCKSIZE_INFO
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_SUPERBLOCK_INFO
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_ACTIVE_SUPERBLOCKS
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_PAGE_INFO
//#define KOKKOS_ENABLE_MEMPOOL_PRINT_INDIVIDUAL_PAGE_INFO
//----------------------------------------------------------------------------
namespace Kokkos {
namespace Experimental {
namespace MempoolImpl {
template < typename T, typename ExecutionSpace >
struct initialize_array {
typedef ExecutionSpace execution_space;
typedef typename ExecutionSpace::size_type size_type;
T * m_data;
T m_value;
initialize_array( T * d, size_t size, T v ) : m_data( d ), m_value( v )
{
Kokkos::parallel_for( size, *this );
execution_space::fence();
}
KOKKOS_INLINE_FUNCTION
void operator()( size_type i ) const { m_data[i] = m_value; }
};
template <typename Bitset>
struct bitset_count
{
typedef typename Bitset::execution_space execution_space;
typedef typename execution_space::size_type size_type;
typedef typename Bitset::size_type value_type;
typedef typename Bitset::word_type word_type;
word_type * m_words;
value_type & m_result;
bitset_count( word_type * w, value_type num_words, value_type & r )
: m_words( w ), m_result( r )
{
parallel_reduce( num_words, *this, m_result );
}
KOKKOS_INLINE_FUNCTION
void init( value_type & v ) const
{ v = 0; }
KOKKOS_INLINE_FUNCTION
void join( volatile value_type & dst, volatile value_type const & src ) const
{ dst += src; }
KOKKOS_INLINE_FUNCTION
void operator()( size_type i, value_type & count ) const
{
count += Kokkos::Impl::bit_count( m_words[i] );
}
};
template < typename Device >
class Bitset {
public:
typedef typename Device::execution_space execution_space;
typedef typename Device::memory_space memory_space;
typedef unsigned word_type;
typedef unsigned size_type;
typedef Kokkos::Impl::DeepCopy< memory_space, Kokkos::HostSpace > raw_deep_copy;
// Define some constants.
enum {
// Size of bitset word. Should be 32.
WORD_SIZE = sizeof(word_type) * CHAR_BIT,
LG_WORD_SIZE = Kokkos::Impl::integral_power_of_two( WORD_SIZE ),
WORD_MASK = WORD_SIZE - 1
};
private:
word_type * m_words;
size_type m_size;
size_type m_num_words;
word_type m_last_word_mask;
public:
~Bitset() = default;
Bitset() = default;
Bitset( Bitset && ) = default;
Bitset( const Bitset & ) = default;
Bitset & operator = ( Bitset && ) = default;
Bitset & operator = ( const Bitset & ) = default;
void init( void * w, size_type s )
{
// Assumption: The size of the memory pointed to by w is a multiple of
// sizeof(word_type).
m_words = reinterpret_cast<word_type*>( w );
m_size = s;
m_num_words = ( s + WORD_SIZE - 1 ) >> LG_WORD_SIZE;
m_last_word_mask = m_size & WORD_MASK ? ( word_type(1) << ( m_size & WORD_MASK ) ) - 1 : 0;
reset();
}
size_type size() const { return m_size; }
size_type count() const
{
size_type val = 0;
bitset_count< Bitset > bc( m_words, m_num_words, val );
return val;
}
void set()
{
// Set all the bits.
initialize_array< word_type, execution_space > ia( m_words, m_num_words, ~word_type(0) );
if ( m_last_word_mask ) {
// Clear the unused bits in the last block.
raw_deep_copy( m_words + ( m_num_words - 1 ), &m_last_word_mask, sizeof(word_type) );
}
}
void reset()
{
initialize_array< word_type, execution_space > ia( m_words, m_num_words, word_type(0) );
}
KOKKOS_FORCEINLINE_FUNCTION
bool test( size_type i ) const
{
size_type word_pos = i >> LG_WORD_SIZE;
word_type word = volatile_load( &m_words[ word_pos ] );
word_type mask = word_type(1) << ( i & WORD_MASK );
return word & mask;
}
KOKKOS_FORCEINLINE_FUNCTION
bool set( size_type i ) const
{
size_type word_pos = i >> LG_WORD_SIZE;
word_type mask = word_type(1) << ( i & WORD_MASK );
return !( atomic_fetch_or( &m_words[ word_pos ], mask ) & mask );
}
KOKKOS_FORCEINLINE_FUNCTION
bool reset( size_type i ) const
{
size_type word_pos = i >> LG_WORD_SIZE;
word_type mask = word_type(1) << ( i & WORD_MASK );
return atomic_fetch_and( &m_words[ word_pos ], ~mask ) & mask;
}
KOKKOS_FORCEINLINE_FUNCTION
Kokkos::pair< bool, word_type >
fetch_word_set( size_type i ) const
{
size_type word_pos = i >> LG_WORD_SIZE;
word_type mask = word_type(1) << ( i & WORD_MASK );
Kokkos::pair<bool, word_type> result;
result.second = atomic_fetch_or( &m_words[ word_pos ], mask );
result.first = !( result.second & mask );
return result;
}
KOKKOS_FORCEINLINE_FUNCTION
Kokkos::pair< bool, word_type >
fetch_word_reset( size_type i ) const
{
size_type word_pos = i >> LG_WORD_SIZE;
word_type mask = word_type(1) << ( i & WORD_MASK );
Kokkos::pair<bool, word_type> result;
result.second = atomic_fetch_and( &m_words[ word_pos ], ~mask );
result.first = result.second & mask;
return result;
}
KOKKOS_FORCEINLINE_FUNCTION
Kokkos::pair< bool, word_type >
set_any_in_word( size_type & pos ) const
{
size_type word_pos = pos >> LG_WORD_SIZE;
word_type word = volatile_load( &m_words[ word_pos ] );
// Loop until there are no more unset bits in the word.
while ( ~word ) {
// Find the first unset bit in the word.
size_type bit = Kokkos::Impl::bit_scan_forward( ~word );
// Try to set the bit.
word_type mask = word_type(1) << bit;
word = atomic_fetch_or( &m_words[ word_pos ], mask );
if ( !( word & mask ) ) {
// Successfully set the bit.
pos = ( word_pos << LG_WORD_SIZE ) + bit;
return Kokkos::pair<bool, word_type>( true, word );
}
}
// Didn't find a free bit in this word.
return Kokkos::pair<bool, word_type>( false, word_type(0) );
}
KOKKOS_FORCEINLINE_FUNCTION
Kokkos::pair< bool, word_type >
set_any_in_word( size_type & pos, word_type word_mask ) const
{
size_type word_pos = pos >> LG_WORD_SIZE;
word_type word = volatile_load( &m_words[ word_pos ] );
word = ( ~word ) & word_mask;
// Loop until there are no more unset bits in the word.
while ( word ) {
// Find the first unset bit in the word.
size_type bit = Kokkos::Impl::bit_scan_forward( word );
// Try to set the bit.
word_type mask = word_type(1) << bit;
word = atomic_fetch_or( &m_words[ word_pos ], mask );
if ( !( word & mask ) ) {
// Successfully set the bit.
pos = ( word_pos << LG_WORD_SIZE ) + bit;
return Kokkos::pair<bool, word_type>( true, word );
}
word = ( ~word ) & word_mask;
}
// Didn't find a free bit in this word.
return Kokkos::pair<bool, word_type>( false, word_type(0) );
}
KOKKOS_FORCEINLINE_FUNCTION
Kokkos::pair< bool, word_type >
reset_any_in_word( size_type & pos ) const
{
size_type word_pos = pos >> LG_WORD_SIZE;
word_type word = volatile_load( &m_words[ word_pos ] );
// Loop until there are no more set bits in the word.
while ( word ) {
// Find the first unset bit in the word.
size_type bit = Kokkos::Impl::bit_scan_forward( word );
// Try to reset the bit.
word_type mask = word_type(1) << bit;
word = atomic_fetch_and( &m_words[ word_pos ], ~mask );
if ( word & mask ) {
// Successfully reset the bit.
pos = ( word_pos << LG_WORD_SIZE ) + bit;
return Kokkos::pair<bool, word_type>( true, word );
}
}
// Didn't find a free bit in this word.
return Kokkos::pair<bool, word_type>( false, word_type(0) );
}
KOKKOS_FORCEINLINE_FUNCTION
Kokkos::pair< bool, word_type >
reset_any_in_word( size_type & pos, word_type word_mask ) const
{
size_type word_pos = pos >> LG_WORD_SIZE;
word_type word = volatile_load( &m_words[ word_pos ] );
word = word & word_mask;
// Loop until there are no more set bits in the word.
while ( word ) {
// Find the first unset bit in the word.
size_type bit = Kokkos::Impl::bit_scan_forward( word );
// Try to reset the bit.
word_type mask = word_type(1) << bit;
word = atomic_fetch_and( &m_words[ word_pos ], ~mask );
if ( word & mask ) {
// Successfully reset the bit.
pos = ( word_pos << LG_WORD_SIZE ) + bit;
return Kokkos::pair<bool, word_type>( true, word );
}
word = word & word_mask;
}
// Didn't find a free bit in this word.
return Kokkos::pair<bool, word_type>( false, word_type(0) );
}
};
template < typename UInt32View, typename BSHeaderView, typename SBHeaderView,
typename MempoolBitset >
struct create_histogram {
typedef typename UInt32View::execution_space execution_space;
typedef typename execution_space::size_type size_type;
typedef Kokkos::pair< double, uint32_t > value_type;
size_t m_start;
UInt32View m_page_histogram;
BSHeaderView m_blocksize_info;
SBHeaderView m_sb_header;
MempoolBitset m_sb_blocks;
size_t m_lg_max_sb_blocks;
uint32_t m_lg_min_block_size;
uint32_t m_blocks_per_page;
value_type & m_result;
create_histogram( size_t start, size_t end, UInt32View ph, BSHeaderView bsi,
SBHeaderView sbh, MempoolBitset sbb, size_t lmsb,
uint32_t lmbs, uint32_t bpp, value_type & r )
: m_start( start ), m_page_histogram( ph ), m_blocksize_info( bsi ),
m_sb_header( sbh ), m_sb_blocks( sbb ), m_lg_max_sb_blocks( lmsb ),
m_lg_min_block_size( lmbs ), m_blocks_per_page( bpp ), m_result( r )
{
Kokkos::parallel_reduce( end - start, *this, m_result );
execution_space::fence();
}
KOKKOS_INLINE_FUNCTION
void init( value_type & v ) const
{
v.first = 0.0;
v.second = 0;
}
KOKKOS_INLINE_FUNCTION
void join( volatile value_type & dst, volatile value_type const & src ) const
{
dst.first += src.first;
dst.second += src.second;
}
KOKKOS_INLINE_FUNCTION
void operator()( size_type i, value_type & r ) const
{
size_type i2 = i + m_start;
uint32_t lg_block_size = m_sb_header(i2).m_lg_block_size;
// A superblock only has a block size of 0 when it is empty.
if ( lg_block_size != 0 ) {
uint32_t block_size_id = lg_block_size - m_lg_min_block_size;
uint32_t blocks_per_sb = m_blocksize_info[block_size_id].m_blocks_per_sb;
uint32_t pages_per_sb = m_blocksize_info[block_size_id].m_pages_per_sb;
uint32_t total_allocated_blocks = 0;
for ( uint32_t j = 0; j < pages_per_sb; ++j ) {
unsigned start_pos = ( i2 << m_lg_max_sb_blocks ) + j * m_blocks_per_page;
unsigned end_pos = start_pos + m_blocks_per_page;
uint32_t page_allocated_blocks = 0;
for ( unsigned k = start_pos; k < end_pos; ++k ) {
page_allocated_blocks += m_sb_blocks.test( k );
}
total_allocated_blocks += page_allocated_blocks;
atomic_increment( &m_page_histogram(page_allocated_blocks) );
}
r.first += double(total_allocated_blocks) / blocks_per_sb;
r.second += blocks_per_sb;
}
}
};
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_SUPERBLOCK_INFO
template < typename UInt32View, typename SBHeaderView, typename MempoolBitset >
struct count_allocated_blocks {
typedef typename UInt32View::execution_space execution_space;
typedef typename execution_space::size_type size_type;
UInt32View m_num_allocated_blocks;
SBHeaderView m_sb_header;
MempoolBitset m_sb_blocks;
size_t m_sb_size;
size_t m_lg_max_sb_blocks;
count_allocated_blocks( size_t num_sb, UInt32View nab, SBHeaderView sbh,
MempoolBitset sbb, size_t sbs, size_t lmsb )
: m_num_allocated_blocks( nab ), m_sb_header( sbh ),
m_sb_blocks( sbb ), m_sb_size( sbs ), m_lg_max_sb_blocks( lmsb )
{
Kokkos::parallel_for( num_sb, *this );
execution_space::fence();
}
KOKKOS_INLINE_FUNCTION
void operator()( size_type i ) const
{
uint32_t lg_block_size = m_sb_header(i).m_lg_block_size;
// A superblock only has a block size of 0 when it is empty.
if ( lg_block_size != 0 ) {
// Count the allocated blocks in the superblock.
uint32_t blocks_per_sb = lg_block_size > 0 ? m_sb_size >> lg_block_size : 0;
unsigned start_pos = i << m_lg_max_sb_blocks;
unsigned end_pos = start_pos + blocks_per_sb;
uint32_t count = 0;
for ( unsigned j = start_pos; j < end_pos; ++j ) {
count += m_sb_blocks.test( j );
}
m_num_allocated_blocks(i) = count;
}
}
};
#endif
}
/// \class MemoryPool
/// \brief Bitset based memory manager for pools of same-sized chunks of memory.
/// \tparam Device Kokkos device that gives the execution and memory space the
/// allocator will be used in.
///
/// MemoryPool is a memory space that can be on host or device. It provides a
/// pool memory allocator for fast allocation of same-sized chunks of memory.
/// The memory is only accessible on the host / device this allocator is
/// associated with.
///
/// This allocator is based on ideas from the following GPU allocators:
/// Halloc (https://github.com/canonizer/halloc).
/// ScatterAlloc (https://github.com/ComputationalRadiationPhysics/scatteralloc)
template < typename Device >
class MemoryPool {
private:
// The allocator uses superblocks. A superblock is divided into pages, and a
// page is divided into blocks. A block is the chunk of memory that is given
// out by the allocator. A page always has a number of blocks equal to the
// size of the word used by the bitset. Thus, the pagesize can vary between
// superblocks as it is based on the block size of the superblock. The
// allocator supports all powers of 2 from MIN_BLOCK_SIZE to the size of a
// superblock as block sizes.
// Superblocks are divided into 4 categories:
// 1. empty - is completely empty; there are no active allocations
// 2. partfull - partially full; there are some active allocations
// 3. full - full enough with active allocations that new allocations
// will likely fail
// 4. active - is currently the active superblock for a block size
//
// An inactive superblock is one that is empty, partfull, or full.
//
// New allocations occur only from an active superblock. If a superblock is
// made inactive after an allocation request is made to it but before the
// allocation request is fulfilled, the allocation will still be attempted
// from that superblock. Deallocations can occur to partfull, full, or
// active superblocks. Superblocks move between categories as allocations
// and deallocations happen. Superblocks all start empty.
//
// Here are the possible moves between categories:
// empty -> active During allocation, there is no active superblock
// or the active superblock is full.
// active -> full During allocation, the full threshold of the
// superblock is reached when increasing the fill
// level.
// full -> partfull During deallocation, the full threshold of the
// superblock is crossed when decreasing the fill
// level.
// partfull -> empty Deallocation of the last allocated block of an
// inactive superblock.
// partfull -> active During allocation, the active superblock is full.
//
// When a new active superblock is needed, partfull superblocks of the same
// block size are chosen over empty superblocks.
//
// The empty and partfull superblocks are tracked using bitsets that represent
// the superblocks in those repsective categories. Empty superblocks use a
// single bitset, while partfull superblocks use a bitset per block size
// (contained sequentially in a single bitset). Active superblocks are
// tracked by the active superblocks array. Full superblocks aren't tracked
// at all.
typedef typename Device::execution_space execution_space;
typedef typename Device::memory_space backend_memory_space;
typedef Device device_type;
typedef MempoolImpl::Bitset< device_type > MempoolBitset;
// Define some constants.
enum {
MIN_BLOCK_SIZE = 64,
LG_MIN_BLOCK_SIZE = Kokkos::Impl::integral_power_of_two( MIN_BLOCK_SIZE ),
MAX_BLOCK_SIZES = 31 - LG_MIN_BLOCK_SIZE + 1,
// Size of bitset word.
BLOCKS_PER_PAGE = MempoolBitset::WORD_SIZE,
LG_BLOCKS_PER_PAGE = MempoolBitset::LG_WORD_SIZE,
INVALID_SUPERBLOCK = ~uint32_t(0),
SUPERBLOCK_LOCK = ~uint32_t(0) - 1,
MAX_TRIES = 32 // Cap on the number of pages searched
// before an allocation returns empty.
};
public:
// Stores information about each superblock.
struct SuperblockHeader {
uint32_t m_full_pages;
uint32_t m_empty_pages;
uint32_t m_lg_block_size;
uint32_t m_is_active;
KOKKOS_FUNCTION
SuperblockHeader() :
m_full_pages(0), m_empty_pages(0), m_lg_block_size(0), m_is_active(false) {}
};
// Stores information about each block size.
struct BlockSizeHeader {
uint32_t m_blocks_per_sb;
uint32_t m_pages_per_sb;
uint32_t m_sb_full_level;
uint32_t m_page_full_level;
KOKKOS_FUNCTION
BlockSizeHeader() :
m_blocks_per_sb(0), m_pages_per_sb(0), m_sb_full_level(0), m_page_full_level(0) {}
};
private:
typedef Kokkos::Impl::SharedAllocationTracker Tracker;
typedef View< uint32_t *, device_type > UInt32View;
typedef View< SuperblockHeader *, device_type > SBHeaderView;
// The letters 'sb' used in any variable name mean superblock.
size_t m_lg_sb_size; // Log2 of superblock size.
size_t m_sb_size; // Superblock size.
size_t m_lg_max_sb_blocks; // Log2 of the number of blocks of the
// minimum block size in a superblock.
size_t m_num_sb; // Number of superblocks.
size_t m_ceil_num_sb; // Number of superblocks rounded up to the smallest
// multiple of the bitset word size. Used by
// bitsets representing superblock categories to
// ensure different block sizes never share a word
// in the bitset.
size_t m_num_block_size; // Number of block sizes supported.
size_t m_data_size; // Amount of memory available to the allocator.
size_t m_sb_blocks_size; // Amount of memory for free / empty blocks bitset.
size_t m_empty_sb_size; // Amount of memory for empty superblocks bitset.
size_t m_partfull_sb_size; // Amount of memory for partfull superblocks bitset.
size_t m_total_size; // Total amount of memory allocated.
char * m_data; // Beginning device memory location used for
// superblocks.
UInt32View m_active; // Active superblocks IDs.
SBHeaderView m_sb_header; // Header info for superblocks.
MempoolBitset m_sb_blocks; // Bitsets representing free / allocated status
// of blocks in superblocks.
MempoolBitset m_empty_sb; // Bitset representing empty superblocks.
MempoolBitset m_partfull_sb; // Bitsets representing partially full superblocks.
Tracker m_track; // Tracker for superblock memory.
BlockSizeHeader m_blocksize_info[MAX_BLOCK_SIZES]; // Header info for block sizes.
// There were several methods tried for storing the block size header info: in a View,
// in a View of const data, and in a RandomAccess View. All of these were slower than
// storing it in a static array that is a member variable to the class. In the latter
// case, the block size info gets copied into the constant memory on the GPU along with
// the class when it is copied there for exeucting a parallel loop. Instead of storing
// the values, computing the values every time they were needed was also tried. This
// method was slightly slower than storing them in the static array.
public:
//! Tag this class as a kokkos memory space
typedef MemoryPool memory_space;
~MemoryPool() = default;
MemoryPool() = default;
MemoryPool( MemoryPool && ) = default;
MemoryPool( const MemoryPool & ) = default;
MemoryPool & operator = ( MemoryPool && ) = default;
MemoryPool & operator = ( const MemoryPool & ) = default;
/// \brief Initializes the memory pool.
/// \param memspace The memory space from which the memory pool will allocate memory.
/// \param total_size The requested memory amount controlled by the allocator. The
/// actual amount is rounded up to the smallest multiple of the
/// superblock size >= the requested size.
/// \param log2_superblock_size Log2 of the size of superblocks used by the allocator.
/// In most use cases, the default value should work.
inline
MemoryPool( const backend_memory_space & memspace,
size_t total_size, size_t log2_superblock_size = 20 )
: m_lg_sb_size( log2_superblock_size ),
m_sb_size( size_t(1) << m_lg_sb_size ),
m_lg_max_sb_blocks( m_lg_sb_size - LG_MIN_BLOCK_SIZE ),
m_num_sb( ( total_size + m_sb_size - 1 ) >> m_lg_sb_size ),
m_ceil_num_sb( ( ( m_num_sb + BLOCKS_PER_PAGE - 1 ) >> LG_BLOCKS_PER_PAGE ) <<
LG_BLOCKS_PER_PAGE ),
m_num_block_size( m_lg_sb_size - LG_MIN_BLOCK_SIZE + 1 ),
m_data_size( m_num_sb * m_sb_size ),
m_sb_blocks_size( ( m_num_sb << m_lg_max_sb_blocks ) / CHAR_BIT ),
m_empty_sb_size( m_ceil_num_sb / CHAR_BIT ),
m_partfull_sb_size( m_ceil_num_sb * m_num_block_size / CHAR_BIT ),
m_total_size( m_data_size + m_sb_blocks_size + m_empty_sb_size + m_partfull_sb_size ),
m_data(0),
m_active( "Active superblocks" ),
m_sb_header( "Superblock headers" ),
m_track()
{
// Assumption. The minimum block size must be a power of 2.
static_assert( Kokkos::Impl::is_integral_power_of_two( MIN_BLOCK_SIZE ), "" );
// Assumption. Require a superblock be large enough so it takes at least 1
// whole bitset word to represent it using the minimum blocksize.
if ( m_sb_size < MIN_BLOCK_SIZE * BLOCKS_PER_PAGE ) {
printf( "\n** MemoryPool::MemoryPool() Superblock size must be >= %u **\n",
MIN_BLOCK_SIZE * BLOCKS_PER_PAGE );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
Kokkos::abort( "" );
}
// Assumption. A superblock's size can be at most 2^31. Verify this.
if ( m_lg_sb_size > 31 ) {
printf( "\n** MemoryPool::MemoryPool() Superblock size must be < %u **\n",
( uint32_t(1) << 31 ) );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
Kokkos::abort( "" );
}
// Assumption. The Bitset only uses unsigned for size types which limits
// the amount of memory the allocator can manage. Verify the memory size
// is below this limit.
if ( m_data_size > size_t(MIN_BLOCK_SIZE) * std::numeric_limits<unsigned>::max() ) {
printf( "\n** MemoryPool::MemoryPool() Allocator can only manage %lu bytes of memory; requested %lu **\n",
size_t(MIN_BLOCK_SIZE) * std::numeric_limits<unsigned>::max(), total_size );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
Kokkos::abort( "" );
}
// Allocate memory for Views. This is done here instead of at construction
// so that the runtime checks can be performed before allocating memory.
resize( m_active, m_num_block_size );
resize( m_sb_header, m_num_sb );
// Allocate superblock memory.
typedef Kokkos::Impl::SharedAllocationRecord< backend_memory_space, void > SharedRecord;
SharedRecord * rec =
SharedRecord::allocate( memspace, "mempool", m_total_size );
m_track.assign_allocated_record_to_uninitialized( rec );
m_data = reinterpret_cast<char *>( rec->data() );
// Set and initialize the free / empty block bitset memory.
m_sb_blocks.init( m_data + m_data_size, m_num_sb << m_lg_max_sb_blocks );
// Set and initialize the empty superblock block bitset memory.
m_empty_sb.init( m_data + m_data_size + m_sb_blocks_size, m_num_sb );
// Start with all superblocks in the empty category.
m_empty_sb.set();
// Set and initialize the partfull superblock block bitset memory.
m_partfull_sb.init( m_data + m_data_size + m_sb_blocks_size + m_empty_sb_size,
m_ceil_num_sb * m_num_block_size );
// Initialize all active superblocks to be invalid.
typename UInt32View::HostMirror host_active = create_mirror_view( m_active );
for ( size_t i = 0; i < m_num_block_size; ++i ) host_active(i) = INVALID_SUPERBLOCK;
deep_copy( m_active, host_active );
// A superblock is considered full when this percentage of its pages are full.
const double superblock_full_fraction = .8;
// A page is considered full when this percentage of its blocks are full.
const double page_full_fraction = .875;
// Initialize the blocksize info.
for ( size_t i = 0; i < m_num_block_size; ++i ) {
uint32_t lg_block_size = i + LG_MIN_BLOCK_SIZE;
uint32_t blocks_per_sb = m_sb_size >> lg_block_size;
uint32_t pages_per_sb = ( blocks_per_sb + BLOCKS_PER_PAGE - 1 ) >> LG_BLOCKS_PER_PAGE;
m_blocksize_info[i].m_blocks_per_sb = blocks_per_sb;
m_blocksize_info[i].m_pages_per_sb = pages_per_sb;
// Set the full level for the superblock.
m_blocksize_info[i].m_sb_full_level =
static_cast<uint32_t>( pages_per_sb * superblock_full_fraction );
if ( m_blocksize_info[i].m_sb_full_level == 0 ) {
m_blocksize_info[i].m_sb_full_level = 1;
}
// Set the full level for the page.
uint32_t blocks_per_page =
blocks_per_sb < BLOCKS_PER_PAGE ? blocks_per_sb : BLOCKS_PER_PAGE;
m_blocksize_info[i].m_page_full_level =
static_cast<uint32_t>( blocks_per_page * page_full_fraction );
if ( m_blocksize_info[i].m_page_full_level == 0 ) {
m_blocksize_info[i].m_page_full_level = 1;
}
}
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_CONSTRUCTOR_INFO
printf( "\n" );
printf( " m_lg_sb_size: %12lu\n", m_lg_sb_size );
printf( " m_sb_size: %12lu\n", m_sb_size );
printf( " m_max_sb_blocks: %12lu\n", size_t(1) << m_lg_max_sb_blocks );
printf( "m_lg_max_sb_blocks: %12lu\n", m_lg_max_sb_blocks );
printf( " m_num_sb: %12lu\n", m_num_sb );
printf( " m_ceil_num_sb: %12lu\n", m_ceil_num_sb );
printf( " m_num_block_size: %12lu\n", m_num_block_size );
printf( " data bytes: %12lu\n", m_data_size );
printf( " sb_blocks bytes: %12lu\n", m_sb_blocks_size );
printf( " empty_sb bytes: %12lu\n", m_empty_sb_size );
printf( " partfull_sb bytes: %12lu\n", m_partfull_sb_size );
printf( " total bytes: %12lu\n", m_total_size );
printf( " m_empty_sb size: %12u\n", m_empty_sb.size() );
printf( "m_partfull_sb size: %12u\n", m_partfull_sb.size() );
printf( "\n" );
fflush( stdout );
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_BLOCKSIZE_INFO
// Print the blocksize info for all the block sizes.
printf( "SIZE BLOCKS_PER_SB PAGES_PER_SB SB_FULL_LEVEL PAGE_FULL_LEVEL\n" );
for ( size_t i = 0; i < m_num_block_size; ++i ) {
printf( "%4zu %13u %12u %13u %15u\n", i + LG_MIN_BLOCK_SIZE,
m_blocksize_info[i].m_blocks_per_sb, m_blocksize_info[i].m_pages_per_sb,
m_blocksize_info[i].m_sb_full_level, m_blocksize_info[i].m_page_full_level );
}
printf( "\n" );
#endif
}
/// \brief The actual block size allocated given alloc_size.
KOKKOS_INLINE_FUNCTION
size_t allocate_block_size( const size_t alloc_size ) const
{ return size_t(1) << ( get_block_size_index( alloc_size ) + LG_MIN_BLOCK_SIZE ); }
/// \brief Allocate a chunk of memory.
/// \param alloc_size Size of the requested allocated in number of bytes.
///
/// The function returns a void pointer to a memory location on success and
/// NULL on failure.
KOKKOS_FUNCTION
void * allocate( size_t alloc_size ) const
{
void * p = 0;
// Only support allocations up to the superblock size. Just return 0
// (failed allocation) for any size above this.
if ( alloc_size <= m_sb_size )
{
int block_size_id = get_block_size_index( alloc_size );
uint32_t blocks_per_sb = m_blocksize_info[block_size_id].m_blocks_per_sb;
uint32_t pages_per_sb = m_blocksize_info[block_size_id].m_pages_per_sb;
#ifdef KOKKOS_IMPL_CUDA_CLANG_WORKAROUND
// Without this test it looks like pages_per_sb might come back wrong.
if ( pages_per_sb == 0 ) return NULL;
#endif
unsigned word_size = blocks_per_sb > 32 ? 32 : blocks_per_sb;
unsigned word_mask = ( uint64_t(1) << word_size ) - 1;
// Instead of forcing an atomic read to guarantee the updated value,
// reading the old value is actually beneficial because more threads will
// attempt allocations on the old active superblock instead of waiting on
// the new active superblock. This will help hide the latency of
// switching the active superblock.
uint32_t sb_id = volatile_load( &m_active(block_size_id) );
// If the active is locked, keep reading it atomically until the lock is
// released.
while ( sb_id == SUPERBLOCK_LOCK ) {
sb_id = atomic_fetch_or( &m_active(block_size_id), uint32_t(0) );
}
load_fence();
bool allocation_done = false;
while ( !allocation_done ) {
bool need_new_sb = false;
if ( sb_id != INVALID_SUPERBLOCK ) {
// Use the value from the clock register as the hash value.
uint64_t hash_val = get_clock_register();
// Get the starting position for this superblock's bits in the bitset.
uint32_t pos_base = sb_id << m_lg_max_sb_blocks;
// Mod the hash value to choose a page in the superblock. The
// initial block searched is the first block of that page.
uint32_t pos_rel = uint32_t( hash_val & ( pages_per_sb - 1 ) ) << LG_BLOCKS_PER_PAGE;
// Get the absolute starting position for this superblock's bits in the bitset.
uint32_t pos = pos_base + pos_rel;
// Keep track of the number of pages searched. Pages in the superblock are
// searched linearly from the starting page. All pages in the superblock are
// searched until either a location is found, or it is proven empty.
uint32_t pages_searched = 0;
bool search_done = false;
while ( !search_done ) {
bool success = false;
unsigned prev_val = 0;
Kokkos::tie( success, prev_val ) = m_sb_blocks.set_any_in_word( pos, word_mask );
if ( !success ) {
if ( ++pages_searched >= pages_per_sb ) {
// Searched all the pages in this superblock. Look for a new superblock.
//
// The previous method tried limiting the number of pages searched, but
// that caused a huge performance issue in CUDA where the outer loop
// executed massive numbers of times. Threads weren't able to find a
// free location when the superblock wasn't full and were able to execute
// the outer loop many times before the superblock was switched for a new
// one. Switching to an exhaustive search eliminated this possiblity and
// didn't slow anything down for the tests.
need_new_sb = true;
search_done = true;
}
else {
// Move to the next page making sure the new search position
// doesn't go past this superblock's bits.
pos += BLOCKS_PER_PAGE;
pos = ( pos < pos_base + blocks_per_sb ) ? pos : pos_base;
}
}
else {
// Reserved a memory location to allocate.
memory_fence();
search_done = true;
allocation_done = true;
uint32_t lg_block_size = block_size_id + LG_MIN_BLOCK_SIZE;
p = m_data + ( size_t(sb_id) << m_lg_sb_size ) +
( ( pos - pos_base ) << lg_block_size );
uint32_t used_bits = Kokkos::Impl::bit_count( prev_val );
if ( used_bits == 0 ) {
// This page was empty. Decrement the number of empty pages for
// the superblock.
atomic_decrement( &m_sb_header(sb_id).m_empty_pages );
}
else if ( used_bits == m_blocksize_info[block_size_id].m_page_full_level - 1 )
{
// This page is full. Increment the number of full pages for
// the superblock.
uint32_t full_pages = atomic_fetch_add( &m_sb_header(sb_id).m_full_pages, 1 );
// This allocation made the superblock full, so a new one needs to be found.
if ( full_pages == m_blocksize_info[block_size_id].m_sb_full_level - 1 ) {
need_new_sb = true;
}
}
}
}
}
else {
// This is the first allocation for this block size. A superblock needs
// to be set as the active one. If this point is reached any other time,
// it is an error.
need_new_sb = true;
}
if ( need_new_sb ) {
uint32_t new_sb_id = find_superblock( block_size_id, sb_id );
if ( new_sb_id == sb_id ) {
allocation_done = true;
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_INFO
printf( "** No superblocks available. **\n" );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
#endif
}
else {
sb_id = new_sb_id;
}
}
}
}
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_INFO
else {
printf( "** Requested allocation size (%zu) larger than superblock size (%lu). **\n",
alloc_size, m_sb_size );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
}
#endif
return p;
}
/// \brief Release allocated memory back to the pool.
/// \param alloc_ptr Pointer to chunk of memory previously allocated by
/// the allocator.
/// \param alloc_size Size of the allocated memory in number of bytes.
KOKKOS_FUNCTION
void deallocate( void * alloc_ptr, size_t alloc_size ) const
{
char * ap = static_cast<char *>( alloc_ptr );
// Only deallocate memory controlled by this pool.
if ( ap >= m_data && ap + alloc_size <= m_data + m_data_size ) {
// Get the superblock for the address. This can be calculated by math on
// the address since the superblocks are stored contiguously in one memory
// chunk.
uint32_t sb_id = ( ap - m_data ) >> m_lg_sb_size;
// Get the starting position for this superblock's bits in the bitset.
uint32_t pos_base = sb_id << m_lg_max_sb_blocks;
// Get the relative position for this memory location's bit in the bitset.
uint32_t offset = ( ap - m_data ) - ( size_t(sb_id) << m_lg_sb_size );
uint32_t lg_block_size = m_sb_header(sb_id).m_lg_block_size;
uint32_t block_size_id = lg_block_size - LG_MIN_BLOCK_SIZE;
uint32_t pos_rel = offset >> lg_block_size;
bool success = false;
unsigned prev_val = 0;
memory_fence();
Kokkos::tie( success, prev_val ) = m_sb_blocks.fetch_word_reset( pos_base + pos_rel );
// If the memory location was previously deallocated, do nothing.
if ( success ) {
uint32_t page_fill_level = Kokkos::Impl::bit_count( prev_val );
if ( page_fill_level == 1 ) {
// This page is now empty. Increment the number of empty pages for the
// superblock.
uint32_t empty_pages = atomic_fetch_add( &m_sb_header(sb_id).m_empty_pages, 1 );
if ( !volatile_load( &m_sb_header(sb_id).m_is_active ) &&
empty_pages == m_blocksize_info[block_size_id].m_pages_per_sb - 1 )
{
// This deallocation caused the superblock to be empty. Change the
// superblock category from partially full to empty.
unsigned pos = block_size_id * m_ceil_num_sb + sb_id;
if ( m_partfull_sb.reset( pos ) ) {
// Reset the empty pages and block size for the superblock.
volatile_store( &m_sb_header(sb_id).m_empty_pages, uint32_t(0) );
volatile_store( &m_sb_header(sb_id).m_lg_block_size, uint32_t(0) );
store_fence();
m_empty_sb.set( sb_id );
}
}
}
else if ( page_fill_level == m_blocksize_info[block_size_id].m_page_full_level ) {
// This page is no longer full. Decrement the number of full pages for
// the superblock.
uint32_t full_pages = atomic_fetch_sub( &m_sb_header(sb_id).m_full_pages, 1 );
if ( !volatile_load( &m_sb_header(sb_id).m_is_active ) &&
full_pages == m_blocksize_info[block_size_id].m_sb_full_level )
{
// This deallocation caused the number of full pages to decrease below
// the full threshold. Change the superblock category from full to
// partially full.
unsigned pos = block_size_id * m_ceil_num_sb + sb_id;
m_partfull_sb.set( pos );
}
}
}
}
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINTERR
else {
printf( "\n** MemoryPool::deallocate() ADDRESS_OUT_OF_RANGE(0x%llx) **\n",
reinterpret_cast<uint64_t>( alloc_ptr ) );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
}
#endif
}
/// \brief Tests if the memory pool has no more memory available to allocate.
KOKKOS_INLINE_FUNCTION
bool is_empty() const
{
// The allocator is empty if all superblocks are full. A superblock is
// full if it has >= 80% of its pages allocated.
// Look at all the superblocks. If one is not full, then the allocator
// isn't empty.
for ( size_t i = 0; i < m_num_sb; ++i ) {
uint32_t lg_block_size = m_sb_header(i).m_lg_block_size;
// A superblock only has a block size of 0 when it is empty.
if ( lg_block_size == 0 ) return false;
uint32_t block_size_id = lg_block_size - LG_MIN_BLOCK_SIZE;
uint32_t full_pages = volatile_load( &m_sb_header(i).m_full_pages );
if ( full_pages < m_blocksize_info[block_size_id].m_sb_full_level ) return false;
}
// All the superblocks were full. The allocator is empty.
return true;
}
// The following functions are used for debugging.
void print_status() const
{
printf( "\n" );
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_SUPERBLOCK_INFO
typename SBHeaderView::HostMirror host_sb_header = create_mirror_view( m_sb_header );
deep_copy( host_sb_header, m_sb_header );
UInt32View num_allocated_blocks( "Allocated Blocks", m_num_sb );
// Count the number of allocated blocks per superblock.
{
MempoolImpl::count_allocated_blocks< UInt32View, SBHeaderView, MempoolBitset >
mch( m_num_sb, num_allocated_blocks, m_sb_header,
m_sb_blocks, m_sb_size, m_lg_max_sb_blocks );
}
typename UInt32View::HostMirror host_num_allocated_blocks =
create_mirror_view( num_allocated_blocks );
deep_copy( host_num_allocated_blocks, num_allocated_blocks );
// Print header info of all superblocks.
printf( "SB_ID SIZE ACTIVE EMPTY_PAGES FULL_PAGES USED_BLOCKS\n" );
for ( size_t i = 0; i < m_num_sb; ++i ) {
printf( "%5zu %4u %6d %11u %10u %10u\n", i,
host_sb_header(i).m_lg_block_size, host_sb_header(i).m_is_active,
host_sb_header(i).m_empty_pages, host_sb_header(i).m_full_pages,
host_num_allocated_blocks(i) );
}
printf( "\n" );
#endif
UInt32View page_histogram( "Page Histogram", 33 );
// Get a View version of the blocksize info.
typedef View< BlockSizeHeader *, device_type > BSHeaderView;
BSHeaderView blocksize_info( "BlockSize Headers", MAX_BLOCK_SIZES );
Kokkos::Impl::DeepCopy< backend_memory_space, Kokkos::HostSpace >
dc( blocksize_info.ptr_on_device(), m_blocksize_info,
sizeof(BlockSizeHeader) * m_num_block_size );
Kokkos::pair< double, uint32_t > result = Kokkos::pair< double, uint32_t >( 0.0, 0 );
// Create the page histogram.
{
MempoolImpl::create_histogram< UInt32View, BSHeaderView, SBHeaderView, MempoolBitset >
mch( 0, m_num_sb, page_histogram, blocksize_info, m_sb_header, m_sb_blocks,
m_lg_max_sb_blocks, LG_MIN_BLOCK_SIZE, BLOCKS_PER_PAGE, result );
}
typename UInt32View::HostMirror host_page_histogram = create_mirror_view( page_histogram );
deep_copy( host_page_histogram, page_histogram );
// Find the used and total pages and blocks.
uint32_t used_pages = 0;
uint32_t used_blocks = 0;
for ( uint32_t i = 1; i < 33; ++i ) {
used_pages += host_page_histogram(i);
used_blocks += i * host_page_histogram(i);
}
uint32_t total_pages = used_pages + host_page_histogram(0);
unsigned num_empty_sb = m_empty_sb.count();
unsigned num_non_empty_sb = m_num_sb - num_empty_sb;
unsigned num_partfull_sb = m_partfull_sb.count();
uint32_t total_blocks = result.second;
double ave_sb_full = num_non_empty_sb == 0 ? 0.0 : result.first / num_non_empty_sb;
double percent_used_sb = double( m_num_sb - num_empty_sb ) / m_num_sb;
double percent_used_pages = total_pages == 0 ? 0.0 : double(used_pages) / total_pages;
double percent_used_blocks = total_blocks == 0 ? 0.0 : double(used_blocks) / total_blocks;
// Count active superblocks.
typename UInt32View::HostMirror host_active = create_mirror_view( m_active );
deep_copy( host_active, m_active );
unsigned num_active_sb = 0;
for ( size_t i = 0; i < m_num_block_size; ++i ) {
num_active_sb += host_active(i) != INVALID_SUPERBLOCK;
}
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_ACTIVE_SUPERBLOCKS
// Print active superblocks.
printf( "BS_ID SB_ID\n" );
for ( size_t i = 0; i < m_num_block_size; ++i ) {
uint32_t sb_id = host_active(i);
if ( sb_id == INVALID_SUPERBLOCK ) {
printf( "%5zu I\n", i );
}
else if ( sb_id == SUPERBLOCK_LOCK ) {
printf( "%5zu L\n", i );
}
else {
printf( "%5zu %7u\n", i, sb_id );
}
}
printf( "\n" );
fflush( stdout );
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_PAGE_INFO
// Print the summary page histogram.
printf( "USED_BLOCKS PAGE_COUNT\n" );
for ( uint32_t i = 0; i < 33; ++i ) {
printf( "%10u %10u\n", i, host_page_histogram[i] );
}
printf( "\n" );
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_INDIVIDUAL_PAGE_INFO
// Print the page histogram for a few individual superblocks.
// const uint32_t num_sb_id = 2;
// uint32_t sb_id[num_sb_id] = { 0, 10 };
const uint32_t num_sb_id = 1;
uint32_t sb_id[num_sb_id] = { 0 };
for ( uint32_t i = 0; i < num_sb_id; ++i ) {
deep_copy( page_histogram, 0 );
{
MempoolImpl::create_histogram< UInt32View, BSHeaderView, SBHeaderView, MempoolBitset >
mch( sb_id[i], sb_id[i] + 1, page_histogram, blocksize_info, m_sb_header,
m_sb_blocks, m_lg_max_sb_blocks, LG_MIN_BLOCK_SIZE, BLOCKS_PER_PAGE, result );
}
deep_copy( host_page_histogram, page_histogram );
printf( "SB_ID USED_BLOCKS PAGE_COUNT\n" );
for ( uint32_t j = 0; j < 33; ++j ) {
printf( "%5u %10u %10u\n", sb_id[i], j, host_page_histogram[j] );
}
printf( "\n" );
}
/*
// Print the blocks used for each page of a few individual superblocks.
for ( uint32_t i = 0; i < num_sb_id; ++i ) {
uint32_t lg_block_size = host_sb_header(sb_id[i]).m_lg_block_size;
if ( lg_block_size != 0 ) {
printf( "SB_ID BLOCK ID USED_BLOCKS\n" );
uint32_t block_size_id = lg_block_size - LG_MIN_BLOCK_SIZE;
uint32_t pages_per_sb = m_blocksize_info[block_size_id].m_pages_per_sb;
for ( uint32_t j = 0; j < pages_per_sb; ++j ) {
unsigned start_pos = ( sb_id[i] << m_lg_max_sb_blocks ) + j * BLOCKS_PER_PAGE;
unsigned end_pos = start_pos + BLOCKS_PER_PAGE;
uint32_t num_allocated_blocks = 0;
for ( unsigned k = start_pos; k < end_pos; ++k ) {
num_allocated_blocks += m_sb_blocks.test( k );
}
printf( "%5u %8u %11u\n", sb_id[i], j, num_allocated_blocks );
}
printf( "\n" );
}
}
*/
#endif
printf( " Used blocks: %10u / %10u = %10.6lf\n", used_blocks, total_blocks,
percent_used_blocks );
printf( " Used pages: %10u / %10u = %10.6lf\n", used_pages, total_pages,
percent_used_pages );
printf( " Used SB: %10zu / %10zu = %10.6lf\n", m_num_sb - num_empty_sb, m_num_sb,
percent_used_sb );
printf( " Active SB: %10u\n", num_active_sb );
printf( " Empty SB: %10u\n", num_empty_sb );
printf( " Partfull SB: %10u\n", num_partfull_sb );
printf( " Full SB: %10lu\n",
m_num_sb - num_active_sb - num_empty_sb - num_partfull_sb );
printf( "Ave. SB Full %%: %10.6lf\n", ave_sb_full );
printf( "\n" );
fflush( stdout );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
}
KOKKOS_INLINE_FUNCTION
size_t get_min_block_size() const { return MIN_BLOCK_SIZE; }
size_t get_mem_size() const { return m_data_size; }
private:
/// \brief Returns the index into the active array for the given size.
///
/// Computes log2 of the largest power of two >= the given size
/// ( ie ceil( log2(size) ) ) shifted by LG_MIN_BLOCK_SIZE.
KOKKOS_FORCEINLINE_FUNCTION
int get_block_size_index( const size_t size ) const
{
// We know the size fits in a 32 bit unsigned because the size of a
// superblock is limited to 2^31, so casting to an unsigned is safe.
// Find the most significant nonzero bit.
uint32_t first_nonzero_bit =
Kokkos::Impl::bit_scan_reverse( static_cast<unsigned>( size ) );
// If size is an integral power of 2, ceil( log2(size) ) is equal to the
// most significant nonzero bit. Otherwise, you need to add 1. Since the
// minimum block size is MIN_BLOCK_SIZE, make sure ceil( log2(size) ) is at
// least LG_MIN_BLOCK_SIZE.
uint32_t lg2_size = first_nonzero_bit + !Kokkos::Impl::is_integral_power_of_two( size );
lg2_size = lg2_size > LG_MIN_BLOCK_SIZE ? lg2_size : LG_MIN_BLOCK_SIZE;
// Return ceil( log2(size) ) shifted so that the value for MIN_BLOCK_SIZE
// is 0.
return lg2_size - LG_MIN_BLOCK_SIZE;
}
/// \brief Finds a superblock with free space to become a new active superblock.
///
/// If this function is called, the current active superblock needs to be replaced
/// because it is full. Initially, only the thread that sets the active superblock
/// to full calls this function. Other threads can still allocate from the "full"
/// active superblock because a full superblock still has locations available. If
/// a thread tries to allocate from the active superblock when it has no free
/// locations, then that thread will call this function, too, and spin on a lock
/// waiting until the active superblock has been replaced.
KOKKOS_FUNCTION
uint32_t find_superblock( int block_size_id, uint32_t old_sb ) const
{
// Try to grab the lock on the head.
uint32_t lock_sb =
Kokkos::atomic_compare_exchange( &m_active(block_size_id), old_sb, SUPERBLOCK_LOCK );
load_fence();
// Initialize the new superblock to be the previous one so the previous
// superblock is returned if a new superblock can't be found.
uint32_t new_sb = lock_sb;
if ( lock_sb == old_sb ) {
// This thread has the lock.
// 1. Look for a partially filled superblock that is of the right block
// size.
size_t max_tries = m_ceil_num_sb >> LG_BLOCKS_PER_PAGE;
size_t tries = 0;
bool search_done = false;
// Set the starting search position to the beginning of this block
// size's bitset.
unsigned pos = block_size_id * m_ceil_num_sb;
while ( !search_done ) {
bool success = false;
unsigned prev_val = 0;
Kokkos::tie( success, prev_val ) = m_partfull_sb.reset_any_in_word( pos );
if ( !success ) {
if ( ++tries >= max_tries ) {
// Exceeded number of words for this block size's bitset.
search_done = true;
}
else {
pos += BLOCKS_PER_PAGE;
}
}
else {
// Found a superblock.
// It is possible that the newly found superblock is the same as the
// old superblock. In this case putting the old value back in yields
// correct behavior. This could happen as follows. This thread
// grabs the lock and transitions the superblock to the full state.
// Before it searches for a new superblock, other threads perform
// enough deallocations to transition the superblock to the partially
// full state. This thread then searches for a partially full
// superblock and finds the one it removed. There's potential for
// this to cause a performance issue if the same superblock keeps
// being removed and added due to the right mix and ordering of
// allocations and deallocations.
search_done = true;
new_sb = pos - block_size_id * m_ceil_num_sb;
// Set the head status for the superblock.
volatile_store( &m_sb_header(new_sb).m_is_active, uint32_t(true) );
// If there was a previous active superblock, mark it as not active.
// It is now in the full category and as such isn't tracked.
if ( lock_sb != INVALID_SUPERBLOCK ) {
volatile_store( &m_sb_header(lock_sb).m_is_active, uint32_t(false) );
}
store_fence();
}
}
// 2. Look for an empty superblock.
if ( new_sb == lock_sb ) {
tries = 0;
search_done = false;
// Set the starting search position to the beginning of this block
// size's bitset.
pos = 0;
while ( !search_done ) {
bool success = false;
unsigned prev_val = 0;
Kokkos::tie( success, prev_val ) = m_empty_sb.reset_any_in_word( pos );
if ( !success ) {
if ( ++tries >= max_tries ) {
// Exceeded number of words for this block size's bitset.
search_done = true;
}
else {
pos += BLOCKS_PER_PAGE;
}
}
else {
// Found a superblock.
// It is possible that the newly found superblock is the same as
// the old superblock. In this case putting the old value back in
// yields correct behavior. This could happen as follows. This
// thread grabs the lock and transitions the superblock to the full
// state. Before it searches for a new superblock, other threads
// perform enough deallocations to transition the superblock to the
// partially full state and then the empty state. This thread then
// searches for a partially full superblock and none exist. This
// thread then searches for an empty superblock and finds the one
// it removed. The likelihood of this happening is so remote that
// the potential for this to cause a performance issue is
// infinitesimal.
search_done = true;
new_sb = pos;
// Set the empty pages, block size, and head status for the
// superblock.
volatile_store( &m_sb_header(new_sb).m_empty_pages,
m_blocksize_info[block_size_id].m_pages_per_sb );
volatile_store( &m_sb_header(new_sb).m_lg_block_size,
block_size_id + LG_MIN_BLOCK_SIZE );
volatile_store( &m_sb_header(new_sb).m_is_active, uint32_t(true) );
// If there was a previous active superblock, mark it as not active.
// It is now in the full category and as such isn't tracked.
if ( lock_sb != INVALID_SUPERBLOCK ) {
volatile_store( &m_sb_header(lock_sb).m_is_active, uint32_t(false) );
}
store_fence();
}
}
}
// Write the new active superblock to release the lock.
atomic_exchange( &m_active(block_size_id), new_sb );
}
else {
// Either another thread has the lock and is switching the active
// superblock for this block size or another thread has already changed
// the active superblock since this thread read its value. Keep
// atomically reading the active superblock until it isn't locked to get
// the new active superblock.
do {
new_sb = atomic_fetch_or( &m_active(block_size_id), uint32_t(0) );
} while ( new_sb == SUPERBLOCK_LOCK );
load_fence();
// Assertions:
// 1. An invalid superblock should never be found here.
// 2. If the new superblock is the same as the previous superblock, the
// allocator is empty.
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINTERR
if ( new_sb == INVALID_SUPERBLOCK ) {
printf( "\n** MemoryPool::find_superblock() FOUND_INACTIVE_SUPERBLOCK **\n" );
#ifdef KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST
fflush( stdout );
#endif
Kokkos::abort( "" );
}
#endif
}
return new_sb;
}
/// Returns 64 bits from a clock register.
KOKKOS_FORCEINLINE_FUNCTION
uint64_t get_clock_register(void) const
{
#if defined( __CUDA_ARCH__ )
// Return value of 64-bit hi-res clock register.
return clock64();
#elif defined( __i386__ ) || defined( __x86_64 )
// Return value of 64-bit hi-res clock register.
unsigned a = 0, d = 0;
__asm__ volatile( "rdtsc" : "=a" (a), "=d" (d) );
return ( (uint64_t) a ) | ( ( (uint64_t) d ) << 32 );
#elif defined( __powerpc ) || defined( __powerpc__ ) || defined( __powerpc64__ ) || \
defined( __POWERPC__ ) || defined( __ppc__ ) || defined( __ppc64__ )
unsigned int cycles = 0;
asm volatile( "mftb %0" : "=r" (cycles) );
return (uint64_t) cycles;
#else
const uint64_t ticks =
std::chrono::high_resolution_clock::now().time_since_epoch().count();
return ticks;
#endif
}
};
} // namespace Experimental
} // namespace Kokkos
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINTERR
#undef KOKKOS_ENABLE_MEMPOOL_PRINTERR
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_INFO
#undef KOKKOS_ENABLE_MEMPOOL_PRINT_INFO
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_BLOCKSIZE_INFO
#undef KOKKOS_ENABLE_MEMPOOL_PRINT_BLOCKSIZE_INFO
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_SUPERBLOCK_INFO
#undef KOKKOS_ENABLE_MEMPOOL_PRINT_SUPERBLOCK_INFO
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_PAGE_INFO
#undef KOKKOS_ENABLE_MEMPOOL_PRINT_PAGE_INFO
#endif
#ifdef KOKKOS_ENABLE_MEMPOOL_PRINT_INDIVIDUAL_PAGE_INFO
#undef KOKKOS_ENABLE_MEMPOOL_PRINT_INDIVIDUAL_PAGE_INFO
#endif
#endif // KOKKOS_MEMORYPOOL_HPP
Event Timeline
Log In to Comment