lammps/lib/kokkos/core/unit_test/TestSharedAlloc.hpp

/*
//@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
*/

#include <gtest/gtest.h>

#include <stdexcept>
#include <sstream>
#include <iostream>

#include <Kokkos_Core.hpp>

/*--------------------------------------------------------------------------*/

namespace Test {

struct SharedAllocDestroy {

  volatile int * count ;

  SharedAllocDestroy() = default ;
  SharedAllocDestroy( int * arg ) : count( arg ) {}

  void destroy_shared_allocation()
    {
      Kokkos::atomic_fetch_add( count , 1 );
    }

};

template< class MemorySpace , class ExecutionSpace >
void test_shared_alloc()
{
#if defined( KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST )

  typedef const Kokkos::Experimental::Impl::SharedAllocationHeader   Header ;
  typedef Kokkos::Experimental::Impl::SharedAllocationTracker  Tracker ;
  typedef Kokkos::Experimental::Impl::SharedAllocationRecord< void , void >                       RecordBase ;
  typedef Kokkos::Experimental::Impl::SharedAllocationRecord< MemorySpace , void >                RecordMemS ;
  typedef Kokkos::Experimental::Impl::SharedAllocationRecord< MemorySpace , SharedAllocDestroy >  RecordFull ;

  static_assert( sizeof(Tracker) == sizeof(int*), "SharedAllocationTracker has wrong size!" );

  MemorySpace s ;

  const size_t N = 1200 ;
  const size_t size = 8 ;

  RecordMemS * rarray[ N ];
  Header     * harray[ N ];

  RecordMemS ** const r = rarray ;
  Header     ** const h = harray ;

  Kokkos::RangePolicy< ExecutionSpace > range(0,N);
  
  //----------------------------------------
  {
  // Since always executed on host space, leave [=]
    Kokkos::parallel_for( range , [=]( size_t i ){
      char name[64] ;
      sprintf(name,"test_%.2d",int(i));

      r[i] = RecordMemS::allocate( s , name , size * ( i + 1 ) );
      h[i] = Header::get_header( r[i]->data() );

      ASSERT_EQ( r[i]->use_count() , 0 );

      for ( size_t j = 0 ; j < ( i / 10 ) + 1 ; ++j ) RecordBase::increment( r[i] );

      ASSERT_EQ( r[i]->use_count() , ( i / 10 ) + 1 );
      ASSERT_EQ( r[i] , RecordMemS::get_record( r[i]->data() ) );
    });

    // Sanity check for the whole set of allocation records to which this record belongs.
    RecordBase::is_sane( r[0] );
    // RecordMemS::print_records( std::cout , s , true );

    Kokkos::parallel_for( range , [=]( size_t i ){
      while ( 0 != ( r[i] = static_cast< RecordMemS *>( RecordBase::decrement( r[i] ) ) ) ) {
        if ( r[i]->use_count() == 1 ) RecordBase::is_sane( r[i] );
      }
    });
  }
  //----------------------------------------
  {
    int destroy_count = 0 ;
    SharedAllocDestroy counter( & destroy_count );

    Kokkos::parallel_for( range , [=]( size_t i ){
      char name[64] ;
      sprintf(name,"test_%.2d",int(i));

      RecordFull * rec = RecordFull::allocate( s , name , size * ( i + 1 ) );

      rec->m_destroy = counter ;

      r[i] = rec ;
      h[i] = Header::get_header( r[i]->data() );

      ASSERT_EQ( r[i]->use_count() , 0 );

      for ( size_t j = 0 ; j < ( i / 10 ) + 1 ; ++j ) RecordBase::increment( r[i] );

      ASSERT_EQ( r[i]->use_count() , ( i / 10 ) + 1 );
      ASSERT_EQ( r[i] , RecordMemS::get_record( r[i]->data() ) );
    });

    RecordBase::is_sane( r[0] );

    Kokkos::parallel_for( range , [=]( size_t i ){
      while ( 0 != ( r[i] = static_cast< RecordMemS *>( RecordBase::decrement( r[i] ) ) ) ) {
        if ( r[i]->use_count() == 1 ) RecordBase::is_sane( r[i] );
      }
    });

    ASSERT_EQ( destroy_count , int(N) );
  }

  //----------------------------------------
  {
    int destroy_count = 0 ;

    {
      RecordFull * rec = RecordFull::allocate( s , "test" , size );

      // ... Construction of the allocated { rec->data() , rec->size() }

      // Copy destruction function object into the allocation record
      rec->m_destroy = SharedAllocDestroy( & destroy_count );

      ASSERT_EQ( rec->use_count() , 0 );

      // Start tracking, increments the use count from 0 to 1
      Tracker track ;

      track.assign_allocated_record_to_uninitialized( rec );

      ASSERT_EQ( rec->use_count() , 1 );
      ASSERT_EQ( track.use_count() , 1 );

      // Verify construction / destruction increment
      for ( size_t i = 0 ; i < N ; ++i ) {
        ASSERT_EQ( rec->use_count() , 1 );
        {
          Tracker local_tracker ;
          local_tracker.assign_allocated_record_to_uninitialized( rec );
          ASSERT_EQ( rec->use_count() , 2 );
          ASSERT_EQ( local_tracker.use_count() , 2 );
        }
        ASSERT_EQ( rec->use_count() , 1 );
        ASSERT_EQ( track.use_count() , 1 );
      }

      Kokkos::parallel_for( range , [=]( size_t i ){
        Tracker local_tracker ;
        local_tracker.assign_allocated_record_to_uninitialized( rec );
        ASSERT_GT( rec->use_count() , 1 );
      });

      ASSERT_EQ( rec->use_count() , 1 );
      ASSERT_EQ( track.use_count() , 1 );

      // Destruction of 'track' object deallocates the 'rec' and invokes the destroy function object.
    }

    ASSERT_EQ( destroy_count , 1 );
  }

#endif /* #if defined( KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST ) */

}


}