ChaiBuffer.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2021, Lawrence Livermore National Security, LLC and LvArray contributors.
 * All rights reserved.
 * See the LICENSE file for details.
 * SPDX-License-Identifier: (BSD-3-Clause)
 */

#pragma once

// Source includes
#include "LvArrayConfig.hpp"
#include "Macros.hpp"
#include "typeManipulation.hpp"
#include "arrayManipulation.hpp"
#include "system.hpp"
#include "bufferManipulation.hpp"

// TPL includes
#include <chai/ArrayManager.hpp>

// System includes
#include <mutex>


namespace LvArray
{

namespace internal
{

inline chai::ArrayManager & getArrayManager()
{
  static chai::ArrayManager & arrayManager = *chai::ArrayManager::getInstance();
  return arrayManager;
}

static std::mutex chaiLock;

inline chai::ExecutionSpace toChaiExecutionSpace( MemorySpace const space )
{
  if( space == MemorySpace::undefined )
    return chai::NONE;
  if( space == MemorySpace::host )
    return chai::CPU;
#if defined(LVARRAY_USE_CUDA)
  if( space == MemorySpace::cuda )
    return chai::GPU;
#endif
#if defined(LVARRAY_USE_HIP)
  if( space == MemorySpace::hip )
    return chai::GPU;
#endif

  LVARRAY_ERROR( "Unrecognized memory space " << static_cast< int >( space ) );

  return chai::NONE;
}

inline MemorySpace toMemorySpace( chai::ExecutionSpace const space )
{
  if( space == chai::NONE )
    return MemorySpace::undefined;
  if( space == chai::CPU )
    return MemorySpace::host;
#if defined(LVARRAY_USE_CUDA)
  if( space == chai::GPU )
    return MemorySpace::cuda;
#endif
#if defined(LVARRAY_USE_HIP)
  if( space == chai::GPU )
    return MemorySpace::hip;
#endif

  LVARRAY_ERROR( "Unrecognized execution space " << static_cast< int >( space ) );

  return MemorySpace::undefined;
}

} // namespace internal

template< typename T >
class ChaiBuffer
{
public:

  using value_type = T;

  constexpr static bool hasShallowCopy = true;

  using T_non_const = std::remove_const_t< T >;

  LVARRAY_HOST_DEVICE inline constexpr
  ChaiBuffer():
    m_pointer( nullptr ),
    m_capacity( 0 ),
    m_pointerRecord( nullptr )
  {}

  LVARRAY_HOST_DEVICE
  ChaiBuffer( bool ):
    m_pointer( nullptr ),
    m_capacity( 0 )
  #if !defined(LVARRAY_DEVICE_COMPILE)
    , m_pointerRecord( new chai::PointerRecord{} )
  #else
    , m_pointerRecord( nullptr )
  #endif
  {
  #if defined(LVARRAY_DEVICE_COMPILE)
    LVARRAY_ERROR( "Creating a new ChaiBuffer on device is not supported. This is often the result of capturing an array on device instead of a view." );
  #else
    m_pointerRecord->m_size = 0;
    setName( "" );

    for( int space = chai::CPU; space < chai::NUM_EXECUTION_SPACES; ++space )
    {
      m_pointerRecord->m_allocators[ space ] = internal::getArrayManager().getAllocatorId( chai::ExecutionSpace( space ) );
    }
  #endif
  }

  LVARRAY_HOST_DEVICE
  ChaiBuffer( std::initializer_list< MemorySpace > const & spaces,
              std::initializer_list< umpire::Allocator > const & allocators ):
    m_pointer( nullptr ),
    m_capacity( 0 )
  #if !defined(LVARRAY_DEVICE_COMPILE)
    , m_pointerRecord( new chai::PointerRecord{} )
  #else
    , m_pointerRecord( nullptr )
  #endif
  {
  #if defined(LVARRAY_DEVICE_COMPILE)
    LVARRAY_ERROR( "Creating a new ChaiBuffer on device is not supported." );
    LVARRAY_UNUSED_VARIABLE( spaces );
    LVARRAY_UNUSED_VARIABLE( allocators );
  #else
    m_pointerRecord->m_size = 0;
    setName( "" );

    LVARRAY_ERROR_IF_NE( spaces.size(), allocators.size() );

    for( int space = chai::CPU; space < chai::NUM_EXECUTION_SPACES; ++space )
    {
      m_pointerRecord->m_allocators[ space ] = internal::getArrayManager().getAllocatorId( chai::ExecutionSpace( space ) );
    }

    for( std::size_t i = 0; i < spaces.size(); ++i )
    {
      m_pointerRecord->m_allocators[ internal::toChaiExecutionSpace( spaces.begin()[ i ] ) ] = allocators.begin()[ i ].getId();
    }
  #endif
  }

  LVARRAY_HOST_DEVICE inline
  ChaiBuffer( ChaiBuffer const & src ):
    m_pointer( src.m_pointer ),
    m_capacity( src.m_capacity ),
    m_pointerRecord( src.m_pointerRecord )
  {
  #if defined(LVARRAY_USE_DEVICE) && !defined(LVARRAY_DEVICE_COMPILE)
    move( internal::toMemorySpace( internal::getArrayManager().getExecutionSpace() ), true );
  #endif
  }

  LVARRAY_HOST_DEVICE inline
  ChaiBuffer( ChaiBuffer const & src, std::ptrdiff_t const size ):
    m_pointer( src.m_pointer ),
    m_capacity( src.m_capacity ),
    m_pointerRecord( src.m_pointerRecord )
  {
  #if defined(LVARRAY_USE_DEVICE) && !defined(LVARRAY_DEVICE_COMPILE)
    moveNested( internal::toMemorySpace( internal::getArrayManager().getExecutionSpace() ), size, true );
  #else
    LVARRAY_UNUSED_VARIABLE( size );
  #endif
  }

  LVARRAY_HOST_DEVICE inline constexpr
  ChaiBuffer( ChaiBuffer && src ):
    m_pointer( src.m_pointer ),
    m_capacity( src.m_capacity ),
    m_pointerRecord( src.m_pointerRecord )
  {
    src.m_capacity = 0;
    src.m_pointer = nullptr;
    src.m_pointerRecord = nullptr;
  }

  template< typename U >
  LVARRAY_HOST_DEVICE inline constexpr
  ChaiBuffer( ChaiBuffer< U > const & src ):
    m_pointer( reinterpret_cast< T * >( src.data() ) ),
    m_capacity( typeManipulation::convertSize< T, U >( src.capacity() ) ),
    m_pointerRecord( &src.pointerRecord() )
  {}

  LVARRAY_HOST_DEVICE inline LVARRAY_INTEL_CONSTEXPR
  ChaiBuffer & operator=( ChaiBuffer const & src )
  {
    m_capacity = src.m_capacity;
    m_pointer = src.m_pointer;
    m_pointerRecord = src.m_pointerRecord;
    return *this;
  }

  LVARRAY_HOST_DEVICE inline LVARRAY_INTEL_CONSTEXPR
  ChaiBuffer & operator=( ChaiBuffer && src )
  {
    m_capacity = src.m_capacity;
    m_pointer = src.m_pointer;
    m_pointerRecord = src.m_pointerRecord;

    src.m_capacity = 0;
    src.m_pointer = nullptr;
    src.m_pointerRecord = nullptr;

    return *this;
  }

  LVARRAY_HOST_DEVICE
  void reallocate( std::ptrdiff_t const size,
                   MemorySpace const space,
                   std::ptrdiff_t const newCapacity )
  {
  #if defined(LVARRAY_DEVICE_COMPILE)
    LVARRAY_ERROR( "Allocation from device is not supported." );
    LVARRAY_UNUSED_VARIABLE( size );
    LVARRAY_UNUSED_VARIABLE( space );
    LVARRAY_UNUSED_VARIABLE( newCapacity );
  #else
    move( space, true );
    chai::PointerRecord * const newRecord = new chai::PointerRecord{};
    newRecord->m_size = newCapacity * sizeof( T );
    newRecord->m_user_callback = m_pointerRecord->m_user_callback;

    for( int s = chai::CPU; s < chai::NUM_EXECUTION_SPACES; ++s )
    {
      newRecord->m_allocators[ s ] = m_pointerRecord->m_allocators[ s ];
    }

    chai::ExecutionSpace const chaiSpace = internal::toChaiExecutionSpace( space );

    internal::chaiLock.lock();
    internal::getArrayManager().allocate( newRecord, chaiSpace );
    internal::chaiLock.unlock();

    T * const newPointer = static_cast< T * >( newRecord->m_pointers[ chaiSpace ] );

    if( size > 0 )
    {
      LVARRAY_ERROR_IF_NE_MSG( space, MemorySpace::host, "Calling reallocate with a non-zero current size is not yet supporeted for the GPU." );
      std::ptrdiff_t const overlapAmount = std::min( newCapacity, size );
      arrayManipulation::uninitializedMove( newPointer, overlapAmount, m_pointer );
      arrayManipulation::destroy( m_pointer, size );
    }

    free();
    m_capacity = newCapacity;
    m_pointer = newPointer;
    m_pointerRecord = newRecord;
    registerTouch( space );
  #endif
  }

  LVARRAY_HOST_DEVICE inline
  void free()
  {
  #if defined(LVARRAY_DEVICE_COMPILE)
    LVARRAY_ERROR( "Deallocation from device is not supported." );
  #else
    std::lock_guard< std::mutex > lock( internal::chaiLock );
    internal::getArrayManager().free( m_pointerRecord );
    m_capacity = 0;
    m_pointer = nullptr;
    m_pointerRecord = nullptr;
  #endif
  }

  LVARRAY_HOST_DEVICE inline constexpr
  std::ptrdiff_t capacity() const
  { return m_capacity; }

  LVARRAY_HOST_DEVICE inline constexpr
  T * data() const
  { return m_pointer; }

  LVARRAY_HOST_DEVICE inline constexpr
  chai::PointerRecord & pointerRecord() const
  { return *m_pointerRecord; }

  template< typename INDEX_TYPE >
  LVARRAY_HOST_DEVICE inline constexpr
  T & operator[]( INDEX_TYPE const i ) const
  { return m_pointer[ i ]; }

  inline
  void moveNested( MemorySpace const space, std::ptrdiff_t const size, bool const touch ) const
  {
  #if defined(LVARRAY_USE_CUDA) || defined(LVARRAY_USE_HIP )
    chai::ExecutionSpace const chaiSpace = internal::toChaiExecutionSpace( space );
    if( m_pointerRecord == nullptr ||
        m_capacity == 0 ||
        chaiSpace == chai::NONE ) return;

    chai::ExecutionSpace const prevSpace = m_pointerRecord->m_last_space;

    if( prevSpace == chai::CPU && prevSpace != chaiSpace ) moveInnerData( space, size, touch );

    move( space, touch );

    if( prevSpace == chai::GPU && prevSpace != chaiSpace ) moveInnerData( space, size, touch );
  #else
    LVARRAY_ERROR_IF_NE( space, MemorySpace::host );
    LVARRAY_UNUSED_VARIABLE( size );
    LVARRAY_UNUSED_VARIABLE( touch );
  #endif
  }

  void move( MemorySpace const space, bool const touch ) const
  {
  #if defined(LVARRAY_USE_CUDA) || defined(LVARRAY_USE_HIP)
    chai::ExecutionSpace const chaiSpace = internal::toChaiExecutionSpace( space );
    if( m_pointerRecord == nullptr ||
        m_capacity == 0 ||
        chaiSpace == chai::NONE ) return;

    auto & am = internal::getArrayManager();
    const_cast< T * & >( m_pointer ) =
      static_cast< T * >( am.move( const_cast< T_non_const * >( m_pointer ), m_pointerRecord, chaiSpace ) );

    if( !std::is_const< T >::value && touch ) m_pointerRecord->m_touched[ chaiSpace ] = true;
    m_pointerRecord->m_last_space = chaiSpace;
  #else
    LVARRAY_ERROR_IF_NE( space, MemorySpace::host );
    LVARRAY_UNUSED_VARIABLE( touch );
  #endif
  }

  MemorySpace getPreviousSpace() const
  { return internal::toMemorySpace( m_pointerRecord->m_last_space ); }

  inline constexpr
  void registerTouch( MemorySpace const space ) const
  {
    chai::ExecutionSpace const chaiSpace = internal::toChaiExecutionSpace( space );
    m_pointerRecord->m_touched[ chaiSpace ] = true;
    m_pointerRecord->m_last_space = chaiSpace;
  }

  template< typename U=ChaiBuffer< T > >
  void setName( std::string const & name )
  {
    std::string const typeString = system::demangleType< U >();
    m_pointerRecord->m_user_callback =
      [name, typeString]( chai::PointerRecord const * const record, chai::Action const act, chai::ExecutionSpace const s )
    {
      if( act == chai::ACTION_MOVE )
      {
        std::string const size = system::calculateSize( record->m_size );
        std::string const paddedSize = std::string( 9 - size.size(), ' ' ) + size;
        char const * const spaceStr = ( s == chai::CPU ) ? "HOST  " : "DEVICE";
        LVARRAY_LOG( "Moved " << paddedSize << " to the " << spaceStr << ": " << typeString << " " << name );
      }
    };
  }

private:

  template< typename U=T_non_const >
  std::enable_if_t< bufferManipulation::HasMemberFunction_move< U > >
  moveInnerData( MemorySpace const space, std::ptrdiff_t const size, bool const touch ) const
  {
    if( space == MemorySpace::undefined ) return;

    for( std::ptrdiff_t i = 0; i < size; ++i )
    {
      const_cast< T_non_const * >( m_pointer )[ i ].move( space, touch );
    }
  }

  template< typename U=T_non_const >
  std::enable_if_t< !bufferManipulation::HasMemberFunction_move< U > >
  moveInnerData( MemorySpace const, std::ptrdiff_t const, bool const ) const
  {}

  T * LVARRAY_RESTRICT m_pointer = nullptr;

  std::ptrdiff_t m_capacity = 0;

  chai::PointerRecord * m_pointerRecord = nullptr;
};

} /* namespace LvArray */