Cabana_Slice.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_SLICE_HPP
#define CABANA_SLICE_HPP
 
#include <Cabana_Types.hpp>
#include <Cabana_Utils.hpp>
#include <impl/Cabana_Index.hpp>
#include <impl/Cabana_TypeTraits.hpp>
 
#include <Kokkos_Core.hpp>
 
#include <cassert>
#include <cstdlib>
#include <cstring>
#include <string>
#include <type_traits>
 
//---------------------------------------------------------------------------//
namespace Cabana
{
namespace Impl
{
 
//---------------------------------------------------------------------------//
// Given a tuple member type T of the given rank get the Kokkos view
// data layout parameters. The tuple index effectively introduces 2 new
// dimensions to the problem on top of the member dimensions - one for the
// struct index and one for the vector index.
template <typename T, std::size_t Rank, int VectorLength, int Stride>
struct KokkosDataTypeImpl;
 
// Rank-0
template <typename T, int VectorLength, int Stride>
struct KokkosDataTypeImpl<T, 0, VectorLength, Stride>
{
    using value_type = typename std::remove_all_extents<T>::type;
    using data_type = value_type* [VectorLength];
    using array_layout = Kokkos::LayoutStride;
 
    inline static array_layout createLayout( const std::size_t num_soa )
    {
        return array_layout( num_soa, Stride, VectorLength, 1 );
    }
};
 
// Rank-1
template <typename T, int VectorLength, int Stride>
struct KokkosDataTypeImpl<T, 1, VectorLength, Stride>
{
    using value_type = typename std::remove_all_extents<T>::type;
    static constexpr std::size_t D0 = std::extent<T, 0>::value;
    using data_type = value_type* [VectorLength][D0];
    using array_layout = Kokkos::LayoutStride;
 
    inline static array_layout createLayout( const std::size_t num_soa )
    {
        return array_layout( num_soa, Stride, VectorLength, 1, D0,
                             VectorLength );
    }
};
 
// Rank-2
template <typename T, int VectorLength, int Stride>
struct KokkosDataTypeImpl<T, 2, VectorLength, Stride>
{
    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 data_type = value_type* [VectorLength][D0][D1];
    using array_layout = Kokkos::LayoutStride;
 
    inline static array_layout createLayout( const std::size_t num_soa )
    {
        return array_layout( num_soa, Stride, VectorLength, 1, D0,
                             VectorLength * D1, D1, VectorLength );
    }
};
 
// Rank-3
template <typename T, int VectorLength, int Stride>
struct KokkosDataTypeImpl<T, 3, VectorLength, Stride>
{
    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 data_type = value_type* [VectorLength][D0][D1][D2];
    using array_layout = Kokkos::LayoutStride;
 
    inline static array_layout createLayout( const std::size_t num_soa )
    {
        return array_layout( num_soa, Stride, VectorLength, 1, D0,
                             VectorLength * D1 * D2, D1, VectorLength * D2, D2,
                             VectorLength );
    }
};
 
// Data type specialization.
template <typename T, int VectorLength, int Stride>
struct KokkosDataType
{
    using kokkos_data_type =
        KokkosDataTypeImpl<T, std::rank<T>::value, VectorLength, Stride>;
    using data_type = typename kokkos_data_type::data_type;
    using array_layout = typename kokkos_data_type::array_layout;
 
    inline static array_layout createLayout( const std::size_t num_soa )
    {
        return kokkos_data_type::createLayout( num_soa );
    }
};
 
//---------------------------------------------------------------------------//
// Kokkos view wrapper for tuple members
template <typename T, int VectorLength, int Stride,
          typename std::enable_if<
              Impl::IsVectorLengthValid<VectorLength>::value, int>::type = 0>
struct KokkosViewWrapper
{
    using data_type =
        typename KokkosDataType<T, VectorLength, Stride>::data_type;
 
    using array_layout =
        typename KokkosDataType<T, VectorLength, Stride>::array_layout;
 
    inline static array_layout createLayout( const std::size_t num_soa )
    {
        return KokkosDataType<T, VectorLength, Stride>::createLayout( num_soa );
    }
};
 
//---------------------------------------------------------------------------//

} // end namespace Impl
 
//---------------------------------------------------------------------------//
//---------------------------------------------------------------------------//
template <typename DataType, typename MemorySpace, typename MemoryAccessType,
          int VectorLength, int Stride>
class Slice
{
  public:
    // Ensure the vector length is valid.
    static_assert( Impl::IsVectorLengthValid<VectorLength>::value,
                   "Vector length must be valid" );

    using slice_type =
        Slice<DataType, MemorySpace, MemoryAccessType, VectorLength, Stride>;

    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;
 
    static_assert( is_memory_access_tag<MemoryAccessType>::value,
                   "Slice memory access type must be a Cabana access type" );
    using memory_access_type = MemoryAccessType;

    static constexpr int vector_length = VectorLength;

    static constexpr int soa_stride = Stride;

    using index_type = Impl::Index<vector_length>;

    static constexpr std::size_t max_supported_rank = 3;

    static constexpr std::size_t max_label_length = 128;

    using view_wrapper =
        Impl::KokkosViewWrapper<DataType, vector_length, soa_stride>;

    using kokkos_view =
        Kokkos::View<typename view_wrapper::data_type,
                     typename view_wrapper::array_layout, MemorySpace,
                     typename MemoryAccessType::kokkos_memory_traits>;

    using size_type = typename kokkos_view::size_type;

    using reference_type = typename kokkos_view::reference_type;
    using value_type = typename kokkos_view::value_type;
    using pointer_type = typename kokkos_view::pointer_type;
    using view_layout = typename kokkos_view::array_layout;

    static constexpr std::size_t data_rank = std::rank<DataType>::value;
    static constexpr std::size_t view_rank = data_rank + 2;

    using default_access_slice =
        Slice<DataType, MemorySpace, DefaultAccessMemory, VectorLength, Stride>;
    using atomic_access_slice =
        Slice<DataType, MemorySpace, AtomicAccessMemory, VectorLength, Stride>;
    using random_access_slice =
        Slice<DataType, MemorySpace, RandomAccessMemory, VectorLength, Stride>;
 
    // Declare slices of different memory access types to be friends.
    friend class Slice<DataType, MemorySpace, DefaultAccessMemory, VectorLength,
                       Stride>;
    friend class Slice<DataType, MemorySpace, AtomicAccessMemory, VectorLength,
                       Stride>;
    friend class Slice<DataType, MemorySpace, RandomAccessMemory, VectorLength,
                       Stride>;
 
    // Equivalent Kokkos view rank. This rank assumes that the struct and
    // array dimension are merged. For the true rank of the raw AoSoA data use
    // the rank() function below which will account for the extra dimension
    // from separate struct and array indices. This enumeration is for
    // compatibility with Kokkos views.
    enum
    {
        rank = data_rank + 1
    };
 
  public:
    Slice()
        : _size( 0 )
    {
    }

    Slice( const pointer_type data, const size_type size,
           const size_type num_soa, const std::string& label = "" )
        : _view( data, view_wrapper::createLayout( num_soa ) )
        , _size( size )
    {
        std::strcpy( _label, label.c_str() );
    }

    template <class MAT>
    Slice( const Slice<DataType, MemorySpace, MAT, VectorLength, Stride>& rhs )
        : _view( rhs._view )
        , _size( rhs._size )
    {
        std::strcpy( _label, rhs._label );
    }

    template <class MAT>
    Slice& operator=(
        const Slice<DataType, MemorySpace, MAT, VectorLength, Stride>& rhs )
    {
        _view = rhs._view;
        _size = rhs._size;
        std::strcpy( _label, rhs._label );
        return *this;
    }

    std::string label() const { return std::string( _label ); }

    KOKKOS_INLINE_FUNCTION
    size_type size() const { return _size; }

    KOKKOS_INLINE_FUNCTION
    size_type numSoA() const { return _view.extent( 0 ); }

    KOKKOS_INLINE_FUNCTION
    size_type arraySize( const size_type s ) const
    {
        // Check if this is not the last struct index.
        // If it isn't, the data array is guaranteed to be full.
        if ( static_cast<size_type>( s ) < _view.extent( 0 ) - 1 )
        {
            return vector_length;
        }
        else
        {
            // This is the last struct index, which may be partially full.
            // We calculate the remainder to see how many elements it holds.
            const size_type rem = _size % vector_length;
            // If rem is 0 and size is positive, the last chunk is full.
            // Otherwise, the remainder is the correct size (e.g., for _size=0).
            return ( rem == 0 && _size > 0 ) ? vector_length : rem;
        }
    }
 
    // ------------
    // 2-D accessor

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 0 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        access( const size_type s, const size_type a ) const
    {
        return _view( s, a );
    }

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 1 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        access( const size_type s, const size_type a, const size_type d0 ) const
    {
        return _view( s, a, d0 );
    }

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 2 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        access( const size_type s, const size_type a, const size_type d0,
                const size_type d1 ) const
    {
        return _view( s, a, d0, d1 );
    }

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 3 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        access( const size_type s, const size_type a, const size_type d0,
                const size_type d1, const size_type d2 ) const
    {
        return _view( s, a, d0, d1, d2 );
    }
 
    // ------------
    // 1-D accessor

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 0 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        operator()( const size_type i ) const
    {
        return access( index_type::s( i ), index_type::a( i ) );
    }

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 1 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        operator()( const size_type i, const size_type d0 ) const
    {
        return access( index_type::s( i ), index_type::a( i ), d0 );
    }

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 2 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        operator()( const size_type i, const size_type d0,
                    const size_type d1 ) const
    {
        return access( index_type::s( i ), index_type::a( i ), d0, d1 );
    }

    template <typename U = DataType>
    KOKKOS_FORCEINLINE_FUNCTION
        typename std::enable_if<( 3 == std::rank<U>::value &&
                                  std::is_same<U, DataType>::value ),
                                reference_type>::type
        operator()( const size_type i, const size_type d0, const size_type d1,
                    const size_type d2 ) const
    {
        return access( index_type::s( i ), index_type::a( i ), d0, d1, d2 );
    }
 
    // -------------------------------
    // Raw data access.

    KOKKOS_INLINE_FUNCTION
    pointer_type data() const { return _view.data(); }

    KOKKOS_INLINE_FUNCTION
    constexpr size_type viewRank() const { return _view.rank; }

    KOKKOS_INLINE_FUNCTION
    size_type extent( const size_type d ) const { return _view.extent( d ); }

    KOKKOS_INLINE_FUNCTION
    size_type stride( const size_type d ) const { return _view.stride( d ); }

    KOKKOS_INLINE_FUNCTION
    kokkos_view view() const { return _view; }
 
  private:
    // The data view. This view is unmanaged and has access traits specified
    // by the template parameters of this class.
    kokkos_view _view;
 
    // Number of tuples in the slice.
    size_type _size;
 
    // Slice label.
    char _label[max_label_length];
};
 
//---------------------------------------------------------------------------//
template <typename>
struct is_slice_impl : public std::false_type
{
};
 
// True only if the type is a member slice *AND* the member slice is templated
// on an AoSoA type.
template <typename DataType, typename MemorySpace, typename MemoryAccessType,
          int VectorLength, int Stride>
struct is_slice_impl<
    Slice<DataType, MemorySpace, MemoryAccessType, VectorLength, Stride>>
    : public std::true_type
{
};

template <class T>
struct is_slice : public is_slice_impl<typename std::remove_cv<T>::type>::type
{
};
 
//---------------------------------------------------------------------------//
// Copy to View.
//---------------------------------------------------------------------------//

template <class ExecutionSpace, class ViewType, class SliceType>
void copySliceToView(
    ExecutionSpace exec_space, ViewType& view, const SliceType& slice,
    const std::size_t begin, const std::size_t end,
    typename std::enable_if<( SliceType::view_rank == 2 ), int*>::type = 0 )
{
    Kokkos::parallel_for(
        "Cabana::copySliceToView::Rank0",
        Kokkos::RangePolicy<ExecutionSpace>( exec_space, begin, end ),
        KOKKOS_LAMBDA( const int i ) { view( i - begin ) = slice( i ); } );
}

template <class ExecutionSpace, class ViewType, class SliceType>
void copySliceToView(
    ExecutionSpace exec_space, ViewType& view, const SliceType& slice,
    const std::size_t begin, const std::size_t end,
    typename std::enable_if<( SliceType::view_rank == 3 ), int*>::type = 0 )
{
    Kokkos::parallel_for(
        "Cabana::copySliceToView::Rank1",
        Kokkos::RangePolicy<ExecutionSpace>( exec_space, begin, end ),
        KOKKOS_LAMBDA( const int i ) {
            for ( std::size_t d0 = 0; d0 < slice.extent( 2 ); ++d0 )
                view( i - begin, d0 ) = slice( i, d0 );
        } );
}

template <class ExecutionSpace, class ViewType, class SliceType>
void copySliceToView(
    ExecutionSpace exec_space, ViewType& view, const SliceType& slice,
    const std::size_t begin, const std::size_t end,
    typename std::enable_if<( SliceType::view_rank == 4 ), int*>::type = 0 )
{
    Kokkos::parallel_for(
        "Cabana::copySliceToView::Rank2",
        Kokkos::RangePolicy<ExecutionSpace>( exec_space, begin, end ),
        KOKKOS_LAMBDA( const int i ) {
            for ( std::size_t d0 = 0; d0 < slice.extent( 2 ); ++d0 )
                for ( std::size_t d1 = 0; d1 < slice.extent( 3 ); ++d1 )
                    view( i - begin, d0, d1 ) = slice( i, d0, d1 );
        } );
}

template <class ViewType, class SliceType>
void copySliceToView( ViewType& view, const SliceType& slice,
                      const std::size_t begin, const std::size_t end )
{
    using exec_space = typename SliceType::execution_space;
    copySliceToView( exec_space{}, view, slice, begin, end );
}
 
//---------------------------------------------------------------------------//
// Copy from View.
//---------------------------------------------------------------------------//

template <class ExecutionSpace, class SliceType, class ViewType>
void copyViewToSlice(
    ExecutionSpace exec_space, SliceType& slice, const ViewType& view,
    const std::size_t begin, const std::size_t end,
    typename std::enable_if<( SliceType::view_rank == 2 ), int*>::type = 0 )
{
    Kokkos::parallel_for(
        "Cabana::copyViewToSlice::Rank0",
        Kokkos::RangePolicy<ExecutionSpace>( exec_space, begin, end ),
        KOKKOS_LAMBDA( const int i ) { slice( i - begin ) = view( i ); } );
}

template <class ExecutionSpace, class SliceType, class ViewType>
void copyViewToSlice(
    ExecutionSpace exec_space, SliceType& slice, const ViewType& view,
    const std::size_t begin, const std::size_t end,
    typename std::enable_if<( SliceType::view_rank == 3 ), int*>::type = 0 )
{
    Kokkos::parallel_for(
        "Cabana::copyViewToSlice::Rank1",
        Kokkos::RangePolicy<ExecutionSpace>( exec_space, begin, end ),
        KOKKOS_LAMBDA( const int i ) {
            for ( std::size_t d0 = 0; d0 < slice.extent( 2 ); ++d0 )
                slice( i - begin, d0 ) = view( i, d0 );
        } );
}

template <class ExecutionSpace, class SliceType, class ViewType>
void copyViewToSlice(
    ExecutionSpace exec_space, SliceType& slice, const ViewType& view,
    const std::size_t begin, const std::size_t end,
    typename std::enable_if<( SliceType::view_rank == 4 ), int*>::type = 0 )
{
    Kokkos::parallel_for(
        "Cabana::copyViewToSlice::Rank2",
        Kokkos::RangePolicy<ExecutionSpace>( exec_space, begin, end ),
        KOKKOS_LAMBDA( const int i ) {
            for ( std::size_t d0 = 0; d0 < slice.extent( 2 ); ++d0 )
                for ( std::size_t d1 = 0; d1 < slice.extent( 3 ); ++d1 )
                    slice( i - begin, d0, d1 ) = view( i, d0, d1 );
        } );
}

template <class ViewType, class SliceType>
void copyViewToSlice( SliceType& slice, const ViewType& view,
                      const std::size_t begin, const std::size_t end )
{
    using exec_space = typename SliceType::execution_space;
    copyViewToSlice( exec_space{}, slice, view, begin, end );
}

template <class SliceType>
void checkSize(
    [[maybe_unused]] SliceType slice, [[maybe_unused]] const std::size_t size,
    typename std::enable_if<is_slice<SliceType>::value, int>::type* = 0 )
{
    // Allow unused for release builds.
    assert( slice.size() == size );
}

template <class ViewType>
void checkSize(
    [[maybe_unused]] ViewType view, [[maybe_unused]] const std::size_t size,
    typename std::enable_if<Kokkos::is_view<ViewType>::value, int>::type* = 0 )
{
    // Allow unused for release builds.
    assert( view.extent( 0 ) == size );
}
 
//---------------------------------------------------------------------------//

template <class SliceType>
KOKKOS_INLINE_FUNCTION auto
size( SliceType slice,
      typename std::enable_if<is_slice<SliceType>::value, int>::type* = 0 )
{
    return slice.size();
}

template <class ViewType>
KOKKOS_INLINE_FUNCTION auto size(
    ViewType view,
    typename std::enable_if<Kokkos::is_view<ViewType>::value, int>::type* = 0 )
{
    return view.extent( 0 );
}
 
} // end namespace Cabana
 
//---------------------------------------------------------------------------//
 
#endif // end CABANA_SLICE_HPP