Cabana_SoA.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_SOA_HPP
#define CABANA_SOA_HPP
 
#include <Cabana_MemberTypes.hpp>
 
#include <Kokkos_Core.hpp>
 
#include <cstdlib>
#include <type_traits>
#include <utility>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
// SoA forward declaration.
template <typename Types, int VectorLength>
struct SoA;
 
//---------------------------------------------------------------------------//
template <class>
struct is_soa_impl : public std::false_type
{
};
 
template <class DataTypes, int VectorLength>
struct is_soa_impl<SoA<DataTypes, VectorLength>> : public std::true_type
{
};

template <class T>
struct is_soa : public is_soa_impl<typename std::remove_cv<T>::type>::type
{
};
 
//---------------------------------------------------------------------------//
 
namespace Impl
{
 
//---------------------------------------------------------------------------//
// Given an array layout and a potentially multi dimensional type T along with
// its rank, compose the inner array type.
template <typename T, std::size_t Rank, int VectorLength>
struct InnerArrayTypeImpl;
 
// rank-0 specialization.
template <typename T, int VectorLength>
struct InnerArrayTypeImpl<T, 0, VectorLength>
{
    using value_type = typename std::remove_all_extents<T>::type;
    using type = value_type[VectorLength];
};
 
// rank-1 specialization.
template <typename T, int VectorLength>
struct InnerArrayTypeImpl<T, 1, VectorLength>
{
    using value_type = typename std::remove_all_extents<T>::type;
    static constexpr std::size_t D0 = std::extent<T, 0>::value;
    using type = value_type[D0][VectorLength];
};
 
// rank-2 specialization.
template <typename T, int VectorLength>
struct InnerArrayTypeImpl<T, 2, VectorLength>
{
    using value_type = typename std::remove_all_extents<T>::type;
    static constexpr std::size_t D0 = std::extent<T, 0>::value;
    static constexpr std::size_t D1 = std::extent<T, 1>::value;
    using type = value_type[D0][D1][VectorLength];
};
 
// rank-3 specialization.
template <typename T, int VectorLength>
struct InnerArrayTypeImpl<T, 3, VectorLength>
{
    using value_type = typename std::remove_all_extents<T>::type;
    static constexpr std::size_t D0 = std::extent<T, 0>::value;
    static constexpr std::size_t D1 = std::extent<T, 1>::value;
    static constexpr std::size_t D2 = std::extent<T, 2>::value;
    using type = value_type[D0][D1][D2][VectorLength];
};
 
//---------------------------------------------------------------------------//
// Inner array type.
template <typename T, int VectorLength>
struct InnerArrayType
{
    using type =
        typename InnerArrayTypeImpl<T, std::rank<T>::value, VectorLength>::type;
};
 
//---------------------------------------------------------------------------//
template <std::size_t M, int VectorLength, typename T>
struct StructMember
{
    using array_type = typename InnerArrayType<T, VectorLength>::type;
    array_type _data;
};
 
//---------------------------------------------------------------------------//
// SoA implementation detail to hide the index sequence.
template <int VectorLength, typename Sequence, typename... Types>
struct SoAImpl;
 
template <int VectorLength, std::size_t... Indices, typename... Types>
struct SoAImpl<VectorLength, std::index_sequence<Indices...>, Types...>
    : StructMember<Indices, VectorLength, Types>...
{
};
 
//---------------------------------------------------------------------------//
// Given an SoA cast it to to one of its member types.
template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION const typename SoA_t::template base<M>&
soaMemberCast( const SoA_t& soa )
{
    static_assert( is_soa<SoA_t>::value, "soaMemberCast only for SoAs" );
    return static_cast<const typename SoA_t::template base<M>&>( soa );
}
 
template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename SoA_t::template base<M>&
soaMemberCast( SoA_t& soa )
{
    static_assert( is_soa<SoA_t>::value, "soaMemberCast only for SoAs" );
    return static_cast<typename SoA_t::template base<M>&>( soa );
}
 
//---------------------------------------------------------------------------//
// Get a pointer to the first element of a member in a given SoA.
template <std::size_t M, class SoA_t>
typename SoA_t::template member_pointer_type<M> soaMemberPtr( SoA_t* p )
{
    static_assert( is_soa<SoA_t>::value, "soaMemberPtr only for SoAs" );
    void* member = static_cast<typename SoA_t::template base<M>*>( p );
    return static_cast<typename SoA_t::template member_pointer_type<M>>(
        member );
}
 
//---------------------------------------------------------------------------//

} // end namespace Impl
 
//---------------------------------------------------------------------------//
// Get template helper.

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_reference_type<M>>::type
get( SoA_t& soa, const std::size_t a )
{
    return Impl::soaMemberCast<M>( soa )._data[a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_const_reference_type<M>>::type
get( const SoA_t& soa, const std::size_t a )
{
    return Impl::soaMemberCast<M>( soa )._data[a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_reference_type<M>>::type
get( SoA_t& soa, const std::size_t a, const std::size_t d0 )
{
    return Impl::soaMemberCast<M>( soa )._data[d0][a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_const_reference_type<M>>::type
get( const SoA_t& soa, const std::size_t a, const std::size_t d0 )
{
    return Impl::soaMemberCast<M>( soa )._data[d0][a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_reference_type<M>>::type
get( SoA_t& soa, const std::size_t a, const std::size_t d0,
     const std::size_t d1 )
{
    return Impl::soaMemberCast<M>( soa )._data[d0][d1][a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_const_reference_type<M>>::type
get( const SoA_t& soa, const std::size_t a, const std::size_t d0,
     const std::size_t d1 )
{
    return Impl::soaMemberCast<M>( soa )._data[d0][d1][a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_reference_type<M>>::type
get( SoA_t& soa, const std::size_t a, const std::size_t d0,
     const std::size_t d1, const std::size_t d2 )
{
    return Impl::soaMemberCast<M>( soa )._data[d0][d1][d2][a];
}

template <std::size_t M, class SoA_t>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    is_soa<SoA_t>::value,
    typename SoA_t::template member_const_reference_type<M>>::type
get( const SoA_t& soa, const std::size_t a, const std::size_t d0,
     const std::size_t d1, const std::size_t d2 )
{
    return Impl::soaMemberCast<M>( soa )._data[d0][d1][d2][a];
}
 
//---------------------------------------------------------------------------//
template <typename... Types, int VectorLength>
struct SoA<MemberTypes<Types...>, VectorLength>
    : Impl::SoAImpl<VectorLength, std::make_index_sequence<sizeof...( Types )>,
                    Types...>
{
    static constexpr int vector_length = VectorLength;

    using member_types = MemberTypes<Types...>;

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

    static constexpr std::size_t max_supported_rank = 3;

    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_reference_type = member_value_type<M>&;

    template <std::size_t M>
    using member_const_reference_type = member_value_type<M> const&;

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

    template <std::size_t M>
    using base = Impl::StructMember<M, vector_length, member_data_type<M>>;
 
    // -------------------------------
    // Member data type properties.

    template <std::size_t M>
    KOKKOS_FORCEINLINE_FUNCTION constexpr std::size_t rank() const
    {
        return std::rank<member_data_type<M>>::value;
    }

    template <std::size_t M, std::size_t D>
    KOKKOS_FORCEINLINE_FUNCTION constexpr std::size_t extent() const
    {
        return std::extent<member_data_type<M>, D>::value;
    }
};
 
//---------------------------------------------------------------------------//
 
namespace Impl
{
 
//---------------------------------------------------------------------------//
// Member element copy operators.
//---------------------------------------------------------------------------//
 
// Copy a single member from one SoA to another.
 
// Rank 0
template <std::size_t M, int DstVectorLength, int SrcVectorLength,
          typename... Types>
KOKKOS_INLINE_FUNCTION typename std::enable_if<
    ( 0 == std::rank<typename MemberTypeAtIndex<
               M, MemberTypes<Types...>>::type>::value ),
    void>::type
soaElementMemberCopy( SoA<MemberTypes<Types...>, DstVectorLength>& dst,
                      const std::size_t dst_idx,
                      const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
                      const std::size_t src_idx )
{
    get<M>( dst, dst_idx ) = get<M>( src, src_idx );
}
 
// Rank 1
template <std::size_t M, int DstVectorLength, int SrcVectorLength,
          typename... Types>
KOKKOS_INLINE_FUNCTION typename std::enable_if<
    ( 1 == std::rank<typename MemberTypeAtIndex<
               M, MemberTypes<Types...>>::type>::value ),
    void>::type
soaElementMemberCopy( SoA<MemberTypes<Types...>, DstVectorLength>& dst,
                      const std::size_t dst_idx,
                      const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
                      const std::size_t src_idx )
{
    for ( std::size_t i0 = 0; i0 < dst.template extent<M, 0>(); ++i0 )
        get<M>( dst, dst_idx, i0 ) = get<M>( src, src_idx, i0 );
}
 
// Rank 2
template <std::size_t M, int DstVectorLength, int SrcVectorLength,
          typename... Types>
KOKKOS_INLINE_FUNCTION typename std::enable_if<
    ( 2 == std::rank<typename MemberTypeAtIndex<
               M, MemberTypes<Types...>>::type>::value ),
    void>::type
soaElementMemberCopy( SoA<MemberTypes<Types...>, DstVectorLength>& dst,
                      const std::size_t dst_idx,
                      const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
                      const std::size_t src_idx )
{
    for ( std::size_t i0 = 0; i0 < dst.template extent<M, 0>(); ++i0 )
        for ( std::size_t i1 = 0; i1 < dst.template extent<M, 1>(); ++i1 )
            get<M>( dst, dst_idx, i0, i1 ) = get<M>( src, src_idx, i0, i1 );
}
 
// Rank 3
template <std::size_t M, int DstVectorLength, int SrcVectorLength,
          typename... Types>
KOKKOS_INLINE_FUNCTION typename std::enable_if<
    ( 3 == std::rank<typename MemberTypeAtIndex<
               M, MemberTypes<Types...>>::type>::value ),
    void>::type
soaElementMemberCopy( SoA<MemberTypes<Types...>, DstVectorLength>& dst,
                      const std::size_t dst_idx,
                      const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
                      const std::size_t src_idx )
{
    for ( std::size_t i0 = 0; i0 < dst.template extent<M, 0>(); ++i0 )
        for ( std::size_t i1 = 0; i1 < dst.template extent<M, 1>(); ++i1 )
            for ( std::size_t i2 = 0; i2 < dst.template extent<M, 2>(); ++i2 )
                get<M>( dst, dst_idx, i0, i1, i2 ) =
                    get<M>( src, src_idx, i0, i1, i2 );
}
 
// Copy the values of all members of an SoA from a source to a destination at
// the given indices.
template <std::size_t M, int DstVectorLength, int SrcVectorLength,
          typename... Types>
KOKKOS_INLINE_FUNCTION void soaElementCopy(
    SoA<MemberTypes<Types...>, DstVectorLength>& dst, const std::size_t dst_idx,
    const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
    const std::size_t src_idx, std::integral_constant<std::size_t, M> )
{
    soaElementMemberCopy<M>( dst, dst_idx, src, src_idx );
    soaElementCopy( dst, dst_idx, src, src_idx,
                    std::integral_constant<std::size_t, M - 1>() );
}
 
template <int DstVectorLength, int SrcVectorLength, typename... Types>
KOKKOS_INLINE_FUNCTION void soaElementCopy(
    SoA<MemberTypes<Types...>, DstVectorLength>& dst, const std::size_t dst_idx,
    const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
    const std::size_t src_idx, std::integral_constant<std::size_t, 0> )
{
    soaElementMemberCopy<0>( dst, dst_idx, src, src_idx );
}
 
// Copy the data from one struct at a given index to another.
template <int DstVectorLength, int SrcVectorLength, typename... Types>
KOKKOS_INLINE_FUNCTION void
tupleCopy( SoA<MemberTypes<Types...>, DstVectorLength>& dst,
           const std::size_t dst_idx,
           const SoA<MemberTypes<Types...>, SrcVectorLength>& src,
           const std::size_t src_idx )
{
    soaElementCopy(
        dst, dst_idx, src, src_idx,
        std::integral_constant<std::size_t, sizeof...( Types ) - 1>() );
}
 
//---------------------------------------------------------------------------//
 
} // end namespace Impl
} // end namespace Cabana
 
#endif // end CABANA_SOA_HPP