Cabana_DeepCopy.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_DEEPCOPY_HPP
#define CABANA_DEEPCOPY_HPP
 
#include <Cabana_AoSoA.hpp>
#include <Cabana_ParticleList.hpp>
#include <Cabana_Slice.hpp>
#include <impl/Cabana_TypeTraits.hpp>
 
#include <Kokkos_Core.hpp>
 
#include <exception>
#include <type_traits>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
template <class Space, class SrcAoSoA>
inline AoSoA<typename SrcAoSoA::member_types, Space, SrcAoSoA::vector_length>
create_mirror(
    const Space&, const SrcAoSoA& src,
    std::enable_if_t<( !std::is_same_v<typename SrcAoSoA::memory_space,
                                       typename Space::memory_space> )>* =
        nullptr )
{
    static_assert( is_aosoa<SrcAoSoA>::value,
                   "create_mirror() requires an AoSoA" );
    return AoSoA<typename SrcAoSoA::member_types, Space,
                 SrcAoSoA::vector_length>(
        std::string( src.label() ).append( "_mirror" ), src.size() );
}
 
//---------------------------------------------------------------------------//
template <class Space, class SrcAoSoA>
inline SrcAoSoA create_mirror_view(
    const Space&, const SrcAoSoA& src,
    std::enable_if_t<( std::is_same_v<typename SrcAoSoA::memory_space,
                                      typename Space::memory_space> )>* =
        nullptr )
{
    static_assert( is_aosoa<SrcAoSoA>::value,
                   "create_mirror_view() requires an AoSoA" );
    return src;
}
 
//---------------------------------------------------------------------------//
template <class Space, class SrcAoSoA>
inline AoSoA<typename SrcAoSoA::member_types, Space, SrcAoSoA::vector_length>
create_mirror_view(
    const Space& space, const SrcAoSoA& src,
    std::enable_if_t<( !std::is_same_v<typename SrcAoSoA::memory_space,
                                       typename Space::memory_space> )>* =
        nullptr )
{
    static_assert( is_aosoa<SrcAoSoA>::value,
                   "create_mirror_view() requires an AoSoA" );
    return create_mirror( space, src );
}
 
//---------------------------------------------------------------------------//
template <class Space, class SrcAoSoA>
inline SrcAoSoA create_mirror_view_and_copy(
    const Space&, const SrcAoSoA& src,
    std::enable_if_t<( std::is_same_v<typename SrcAoSoA::memory_space,
                                      typename Space::memory_space> &&
                       is_aosoa<SrcAoSoA>::value )>* = nullptr )
{
    return src;
}
 
//---------------------------------------------------------------------------//
template <class Space, class SrcAoSoA>
inline AoSoA<typename SrcAoSoA::member_types, Space, SrcAoSoA::vector_length>
create_mirror_view_and_copy(
    const Space& space, const SrcAoSoA& src,
    std::enable_if_t<( !std::is_same_v<typename SrcAoSoA::memory_space,
                                       typename Space::memory_space> &&
                       is_aosoa<SrcAoSoA>::value )>* = nullptr )
{
    auto dst = create_mirror( space, src );
 
    Kokkos::deep_copy(
        typename decltype( dst )::soa_view( dst.data(), dst.numSoA() ),
        typename SrcAoSoA::soa_view( src.data(), src.numSoA() ) );
 
    return dst;
}
 
//---------------------------------------------------------------------------//
template <class DstAoSoA, class SrcAoSoA>
inline void
deep_copy( DstAoSoA& dst, const SrcAoSoA& src,
           std::enable_if_t<( is_aosoa<DstAoSoA>::value &&
                              is_aosoa<SrcAoSoA>::value )>* = nullptr )
{
    using dst_type = DstAoSoA;
    using src_type = SrcAoSoA;
    using dst_memory_space = typename dst_type::memory_space;
    using src_memory_space = typename src_type::memory_space;
    using dst_soa_type = typename dst_type::soa_type;
    using src_soa_type = typename src_type::soa_type;
 
    // Check that the data types are the same.
    static_assert( std::is_same_v<typename dst_type::member_types,
                                  typename src_type::member_types>,
                   "Cabana::deep_copy: Attempted to deep copy AoSoA objects of "
                   "different member types" );
 
    // Check for the same number of values.
    if ( dst.size() != src.size() )
    {
        throw std::runtime_error(
            "Cabana::deep_copy: Attempted to deep copy "
            "AoSoA objects of different sizes. (Labels: " +
            src.label() + ", " + dst.label() + ")" );
    }
 
    // Get the pointers to the beginning of the data blocks.
    void* dst_data = dst.data();
    const void* src_data = src.data();
 
    // Return if both pointers are null.
    if ( dst_data == nullptr && src_data == nullptr )
    {
        return;
    }
 
    // Get the number of SoA's in each object.
    auto dst_num_soa = dst.numSoA();
    auto src_num_soa = src.numSoA();
 
    // Return if the AoSoA memory occupies the same space.
    if ( ( dst_data == src_data ) && ( dst_num_soa * sizeof( dst_soa_type ) ==
                                       src_num_soa * sizeof( src_soa_type ) ) )
    {
        return;
    }
 
    // If the inner array size is the same and both AoSoAs have the same number
    // of values then we can do a byte-wise copy directly.
    if ( std::is_same_v<dst_soa_type, src_soa_type> )
    {
        Kokkos::deep_copy( Kokkos::View<char*, dst_memory_space>(
                               reinterpret_cast<char*>( dst.data() ),
                               dst.numSoA() * sizeof( dst_soa_type ) ),
                           Kokkos::View<char*, src_memory_space>(
                               reinterpret_cast<char*>( src.data() ),
                               src.numSoA() * sizeof( src_soa_type ) ) );
    }
 
    // Otherwise copy the data element-by-element because the data layout is
    // different.
    else
    {
        // Create an AoSoA in the destination space with the same data layout
        // as the source.
        auto src_copy_on_dst = create_mirror_view_and_copy(
            typename dst_type::memory_space(), src );
 
        // Copy via tuples.
        auto copy_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            dst.setTuple( i, src_copy_on_dst.getTuple( i ) );
        };
        Kokkos::RangePolicy<typename dst_type::execution_space> exec_policy(
            0, dst.size() );
        Kokkos::parallel_for( "Cabana::deep_copy", exec_policy, copy_func );
        Kokkos::fence();
    }
}
 
//---------------------------------------------------------------------------//
template <class DstMemorySpace, class SrcMemorySpace, int DstVectorLength,
          int SrcVectorLength, class... FieldTags>
inline void deep_copy(
    ParticleList<DstMemorySpace, DstVectorLength, FieldTags...>& dst,
    const ParticleList<SrcMemorySpace, SrcVectorLength, FieldTags...>& src )
{
    // Copy particle data to new memory space.
    auto aosoa_src = src.aosoa();
    auto& aosoa_dst = dst.aosoa();
 
    // Set the new data.
    Cabana::deep_copy( aosoa_dst, aosoa_src );
}
 
//---------------------------------------------------------------------------//
template <class AoSoA_t>
inline void deep_copy( AoSoA_t& aosoa,
                       const typename AoSoA_t::tuple_type& tuple )
{
    static_assert( is_aosoa<AoSoA_t>::value,
                   "Cabana::deep_copy: Only AoSoAs can be assigned tuples" );
    auto assign_func = KOKKOS_LAMBDA( const std::size_t i )
    {
        aosoa.setTuple( i, tuple );
    };
    Kokkos::RangePolicy<typename AoSoA_t::execution_space> exec_policy(
        0, aosoa.size() );
    Kokkos::parallel_for( "Cabana::deep_copy", exec_policy, assign_func );
    Kokkos::fence();
}
 
//---------------------------------------------------------------------------//
template <class DstSlice, class SrcSlice>
inline void
deep_copy( DstSlice& dst, const SrcSlice& src,
           std::enable_if_t<( is_slice<DstSlice>::value &&
                              is_slice<SrcSlice>::value )>* = nullptr )
{
    using dst_type = DstSlice;
    using src_type = SrcSlice;
 
    // Check that the data types are the same.
    static_assert( std::is_same_v<typename dst_type::value_type,
                                  typename src_type::value_type>,
                   "Cabana::deep_copy: Attempted to deep copy Slice objects of "
                   "different value types" );
 
    // Check that the element dimensions are the same.
    static_assert( SrcSlice::view_layout::D0 == SrcSlice::view_layout::D0,
                   "Cabana::deep_copy: Slice dimension 0 is different" );
    static_assert( SrcSlice::view_layout::D1 == SrcSlice::view_layout::D1,
                   "Cabana::deep_copy: Slice dimension 1 is different" );
    static_assert( SrcSlice::view_layout::D2 == SrcSlice::view_layout::D2,
                   "Cabana::deep_copy: Slice dimension 2 is different" );
    static_assert( SrcSlice::view_layout::D3 == SrcSlice::view_layout::D3,
                   "Cabana::deep_copy: Slice dimension 3 is different" );
    static_assert( SrcSlice::view_layout::D4 == SrcSlice::view_layout::D4,
                   "Cabana::deep_copy: Slice dimension 4 is different" );
    static_assert( SrcSlice::view_layout::D5 == SrcSlice::view_layout::D5,
                   "Cabana::deep_copy: Slice dimension 5 is different" );
 
    // Check for the same number of elements.
    if ( dst.size() != src.size() )
    {
        throw std::runtime_error(
            "Cabana::deep_copy: Attempted to deep copy Slice objects of "
            "different sizes. (Labels: " +
            src.label() + ", " + dst.label() + ")" );
    }
 
    // Get the pointers to the beginning of the data blocks.
    auto dst_data = dst.data();
    const auto src_data = src.data();
 
    // Return if both pointers are null.
    if ( dst_data == nullptr && src_data == nullptr )
    {
        return;
    }
 
    // Get the number of SoA's in each object.
    auto dst_num_soa = dst.numSoA();
    auto src_num_soa = src.numSoA();
 
    // Return if the slice memory occupies the same space.
    if ( ( dst_data == src_data ) &&
         ( dst_num_soa * dst.stride( 0 ) == src_num_soa * src.stride( 0 ) ) )
    {
        return;
    }
 
    // Get the number of components in each slice element.
    std::size_t num_comp = 1;
    for ( std::size_t d = 2; d < dst.viewRank(); ++d )
        num_comp *= dst.extent( d );
 
    // Gather the slice data in a flat view in the source space and copy it to
    // the destination space.
    Kokkos::View<typename dst_type::value_type*,
                 typename dst_type::memory_space>
        gather_dst( "gather_dst", num_comp * dst.size() );
    {
        Kokkos::View<typename src_type::value_type*,
                     typename src_type::memory_space>
            gather_src( "gather_src", num_comp * src.size() );
        auto gather_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            auto src_offset = SrcSlice::index_type::s( i ) * src.stride( 0 ) +
                              SrcSlice::index_type::a( i );
            for ( std::size_t n = 0; n < num_comp; ++n )
                gather_src( i * num_comp + n ) =
                    src_data[src_offset + SrcSlice::vector_length * n];
        };
        Kokkos::RangePolicy<typename src_type::execution_space> gather_policy(
            0, src.size() );
        Kokkos::parallel_for( "Cabana::deep_copy::gather", gather_policy,
                              gather_func );
        Kokkos::fence();
        Kokkos::deep_copy( gather_dst, gather_src );
    }
 
    // Scatter back into the destination slice from the gathered slice.
    auto scatter_func = KOKKOS_LAMBDA( const std::size_t i )
    {
        auto dst_offset = DstSlice::index_type::s( i ) * dst.stride( 0 ) +
                          DstSlice::index_type::a( i );
        for ( std::size_t n = 0; n < num_comp; ++n )
            dst_data[dst_offset + DstSlice::vector_length * n] =
                gather_dst( i * num_comp + n );
    };
    Kokkos::RangePolicy<typename dst_type::execution_space> scatter_policy(
        0, dst.size() );
    Kokkos::parallel_for( "Cabana::deep_copy::scatter", scatter_policy,
                          scatter_func );
    Kokkos::fence();
}
 
//---------------------------------------------------------------------------//
template <class Slice_t>
inline void deep_copy( Slice_t& slice,
                       const typename Slice_t::value_type scalar )
{
    static_assert( is_slice<Slice_t>::value,
                   "Cabana::deep_copy: Only slices can be assigned scalars" );
    Kokkos::deep_copy( slice.view(), scalar );
}
 
//---------------------------------------------------------------------------//
template <class DstMemorySpace, class SrcMemorySpace, int VectorLength,
          class... FieldTags>
auto create_mirror_view_and_copy(
    DstMemorySpace,
    ParticleList<SrcMemorySpace, VectorLength, FieldTags...> plist_src,
    std::enable_if_t<std::is_same_v<SrcMemorySpace, DstMemorySpace>>* =
        nullptr )
{
    return plist_src;
}

template <class DstMemorySpace, class SrcMemorySpace, int VectorLength,
          class... FieldTags>
auto create_mirror_view_and_copy(
    DstMemorySpace,
    ParticleList<SrcMemorySpace, VectorLength, FieldTags...> plist_src,
    std::enable_if_t<!std::is_same_v<SrcMemorySpace, DstMemorySpace>>* =
        nullptr )
{
    // Extract the original AoSoA.
    auto aosoa_src = plist_src.aosoa();
 
    // Create an AoSoA in the new memory space.
    using src_plist_type =
        ParticleList<SrcMemorySpace, VectorLength, FieldTags...>;
    using member_types = typename src_plist_type::member_types;
    AoSoA<member_types, DstMemorySpace, VectorLength> aosoa_dst(
        aosoa_src.label(), aosoa_src.size() );
 
    // Copy data to new AoAoA.
    deep_copy( aosoa_dst, aosoa_src );
 
    // Create new list with the copied data.
    return ParticleList<DstMemorySpace, VectorLength, FieldTags...>(
        aosoa_dst );
}
 
} // end namespace Cabana
 
#endif // end CABANA_DEEPCOPY_HPP