Cabana_AoSoA.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef CABANA_AOSOA_HPP
#define CABANA_AOSOA_HPP
 
#include <Cabana_MemberTypes.hpp>
#include <Cabana_Slice.hpp>
#include <Cabana_SoA.hpp>
#include <Cabana_Tuple.hpp>
#include <Cabana_Types.hpp>
#include <Cabana_Utils.hpp>
#include <impl/Cabana_Index.hpp>
#include <impl/Cabana_PerformanceTraits.hpp>
 
#include <Kokkos_Core.hpp>
 
#include <cmath>
#include <cstdlib>
#include <string>
#include <type_traits>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
// AoSoA forward declaration.
template <class DataTypes, class MemorySpace, int VectorLength,
          class MemoryTraits>
class AoSoA;
 
//---------------------------------------------------------------------------//
template <class>
struct is_aosoa_impl : public std::false_type
{
};
 
template <class DataTypes, class MemorySpace, int VectorLength,
          class MemoryTraits>
struct is_aosoa_impl<AoSoA<DataTypes, MemorySpace, VectorLength, MemoryTraits>>
    : public std::true_type
{
};

template <class T>
struct is_aosoa : public is_aosoa_impl<typename std::remove_cv<T>::type>::type
{
};
 
//---------------------------------------------------------------------------//
template <std::size_t M, class AoSoA_t>
typename AoSoA_t::template member_slice_type<M>
slice( const AoSoA_t& aosoa, const std::string& slice_label = "" )
{
    static_assert(
        0 == sizeof( typename AoSoA_t::soa_type ) %
                 sizeof( typename AoSoA_t::template member_value_type<M> ),
        "Slice stride cannot be calculated for misaligned memory!" );
 
    return typename AoSoA_t::template member_slice_type<M>(
        Impl::soaMemberPtr<M>( aosoa.data() ), aosoa.size(), aosoa.numSoA(),
        slice_label );
}
 
//---------------------------------------------------------------------------//
template <class DataTypes, class MemorySpace,
          int VectorLength = Impl::PerformanceTraits<
              typename MemorySpace::execution_space>::vector_length,
          class MemoryTraits = Kokkos::MemoryManaged>
class AoSoA
{
  public:
    using aosoa_type =
        AoSoA<DataTypes, MemorySpace, VectorLength, MemoryTraits>;

    using host_mirror_type = AoSoA<DataTypes, Kokkos::HostSpace, VectorLength>;
 
    static_assert( is_member_types<DataTypes>::value,
                   "AoSoA data types must be member types" );
    static_assert( CheckMemberTypes<DataTypes>::value,
                   "AoSoA data type failure" );
    using member_types = DataTypes;
 
    // FIXME: extracting the self type for backwards compatibility with previous
    // template on DeviceType. Should simply be MemorySpace after next release.
    using memory_space = typename MemorySpace::memory_space;
    // FIXME: replace warning with memory space assert after next release.
    static_assert( Impl::deprecated( Kokkos::is_device<MemorySpace>() ) );

    using device_type [[deprecated]] = typename memory_space::device_type;
    using execution_space = typename memory_space::execution_space;
 
    static_assert( Impl::IsVectorLengthValid<VectorLength>::value,
                   "Vector length must be valid" );
    static constexpr int vector_length = VectorLength;

    using memory_traits = MemoryTraits;

    using size_type = typename memory_space::size_type;

    using soa_type = SoA<member_types, vector_length>;

    using soa_view = Kokkos::View<soa_type*, memory_space, memory_traits>;

    static constexpr std::size_t number_of_members = member_types::size;

    static constexpr std::size_t max_supported_rank = 3;

    using index_type = Impl::Index<vector_length>;

    using tuple_type = Tuple<member_types>;

    template <std::size_t M>
    using member_data_type = typename MemberTypeAtIndex<M, member_types>::type;

    template <std::size_t M>
    using member_value_type =
        typename std::remove_all_extents<member_data_type<M>>::type;

    template <std::size_t M>
    using member_pointer_type =
        typename std::add_pointer<member_value_type<M>>::type;

    template <std::size_t M>
    using member_slice_type =
        Slice<member_data_type<M>, memory_space, DefaultAccessMemory,
              vector_length,
              sizeof( soa_type ) / sizeof( member_value_type<M> )>;
 
  public:
    AoSoA( const std::string& label = "" )
        : _size( 0 )
        , _capacity( 0 )
        , _num_soa( 0 )
        , _data( Kokkos::ViewAllocateWithoutInitializing( label ), 0 )
    {
        static_assert(
            !memory_traits::is_unmanaged,
            "Construction by allocation cannot use unmanaged memory" );
    }

    AoSoA( const std::string label, const size_type n )
        : _size( n )
        , _capacity( 0 )
        , _num_soa( 0 )
        , _data( Kokkos::ViewAllocateWithoutInitializing( label ), 0 )
    {
        static_assert(
            !memory_traits::is_unmanaged,
            "Construction by allocation cannot use unmanaged memory" );
        resize( _size );
    }

    AoSoA( soa_type* ptr, const size_type num_soa, const size_type n )
        : _size( n )
        , _capacity( num_soa * vector_length )
        , _num_soa( num_soa )
        , _data( ptr, num_soa )
    {
        static_assert( memory_traits::is_unmanaged,
                       "Pointer construction requires unmanaged memory" );
    }

    std::string label() const { return _data.label(); }

    KOKKOS_FUNCTION
    size_type size() const { return _size; }

    KOKKOS_FUNCTION
    bool empty() const { return ( size() == 0 ); }

    KOKKOS_FUNCTION
    size_type capacity() const { return _capacity; }

    void resize( const size_type n )
    {
        static_assert( !memory_traits::is_unmanaged,
                       "Cannot resize unmanaged memory" );
 
        // Reserve memory if needed.
        reserve( n );
 
        // Update the sizes of the data. This is potentially different than
        // the amount of allocated data.
        _size = n;
        _num_soa = std::ceil( static_cast<double>( n ) / vector_length );
    }

    void reserve( const size_type n )
    {
        static_assert( !memory_traits::is_unmanaged,
                       "Cannot reserve unmanaged memory" );
 
        // If we aren't asking for more memory then we have nothing to do.
        if ( n <= _capacity )
            return;
 
        // Figure out the new capacity.
        size_type num_soa_alloc = std::floor( n / vector_length );
        if ( 0 < n % vector_length )
            ++num_soa_alloc;
 
        // If we aren't asking for any more SoA objects then we still have
        // nothing to do.
        if ( num_soa_alloc <= _num_soa )
            return;
 
        // Assign the new capacity.
        _capacity = num_soa_alloc * vector_length;
 
        // We need more SoA objects so allocate a new view and copy the
        // existing data.
        soa_view resized_data(
            Kokkos::ViewAllocateWithoutInitializing( _data.label() ),
            num_soa_alloc );
        if ( _num_soa > 0 )
            Kokkos::deep_copy(
                Kokkos::subview(
                    resized_data,
                    Kokkos::pair<size_type, size_type>( 0, _num_soa ) ),
                Kokkos::subview( _data, Kokkos::pair<size_type, size_type>(
                                            0, _num_soa ) ) );
        _data = resized_data;
    }

    void shrinkToFit()
    {
        static_assert( !memory_traits::is_unmanaged,
                       "Cannot shrink unmanaged memory" );
 
        // If we aren't asking for any fewer SoA objects then we have nothing
        // to do. The amount of allocated data has to be at least as big as
        // _num_soa so we just need to check here that they are equivalent. If
        // they are equivalent, the container is already as small as it can be.
        if ( _data.size() == _num_soa )
            return;
 
        // Assign the new capacity.
        _capacity = _num_soa * vector_length;
 
        // We need fewer SoA objects so allocate a new view and copy the
        // existing data.
        soa_view resized_data(
            Kokkos::ViewAllocateWithoutInitializing( _data.label() ),
            _num_soa );
        if ( _num_soa > 0 )
            Kokkos::deep_copy(
                resized_data,
                Kokkos::subview( _data, Kokkos::pair<size_type, size_type>(
                                            0, _num_soa ) ) );
        _data = resized_data;
    }

    KOKKOS_INLINE_FUNCTION
    size_type numSoA() const { return _num_soa; }

    KOKKOS_INLINE_FUNCTION
    size_type arraySize( const size_type s ) const
    {
        // the SoA struct size should be full size, i.e. vector_length if:
        // 1) s is not the last SoA struct
        // or 2) if _size = _num_soa * vector_length
        return ( ( (size_type)s < _num_soa - 1 ) ||
                 ( _size % vector_length == 0 ) )
                   ? vector_length              // if s is a full SoA struct
                   : ( _size % vector_length ); // if s is the last SoA struct
    }

    KOKKOS_FORCEINLINE_FUNCTION
    soa_type& access( const size_type s ) const { return _data( s ); }

    KOKKOS_INLINE_FUNCTION
    tuple_type getTuple( const size_type i ) const
    {
        tuple_type tpl;
        Impl::tupleCopy( tpl, 0, _data( index_type::s( i ) ),
                         index_type::a( i ) );
        return tpl;
    }

    KOKKOS_INLINE_FUNCTION
    void setTuple( const size_type i, const tuple_type& tpl ) const
    {
        Impl::tupleCopy( _data( index_type::s( i ) ), index_type::a( i ), tpl,
                         0 );
    }

    soa_type* data() const { return _data.data(); }
 
  private:
    // Total number of tuples in the container.
    size_type _size;
 
    // Allocated number of tuples in all arrays in all structs.
    size_type _capacity;
 
    // Number of structs-of-arrays in the array.
    size_type _num_soa;
 
    // Structs-of-Arrays managed data. This Kokkos View manages the block of
    // memory owned by this class such that the copy constructor and
    // assignment operator for this class perform a shallow and reference
    // counted copy of the data.
    soa_view _data;
};
 
//---------------------------------------------------------------------------//
 
} // end namespace Cabana
 
#endif // CABANA_AOSOA_HPP