Cabana_Sort.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_SORT_HPP
#define CABANA_SORT_HPP
 
#include <Cabana_AoSoA.hpp>
#include <Cabana_DeepCopy.hpp>
#include <Cabana_Slice.hpp>
#include <Cabana_Utils.hpp>
 
#include <Kokkos_Core.hpp>
#include <Kokkos_Profiling_ScopedRegion.hpp>
#include <Kokkos_Sort.hpp>
 
#include <type_traits>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
template <class MemorySpace>
class BinningData
{
  public:
    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;
    using size_type = typename memory_space::size_type;
    using CountView = Kokkos::View<const int*, memory_space>;
    using OffsetView = Kokkos::View<size_type*, memory_space>;

    BinningData()
        : _nbin( 0 )
    {
    }

    BinningData( const std::size_t begin, const std::size_t end,
                 CountView counts, OffsetView offsets,
                 OffsetView permute_vector )
        : _begin( begin )
        , _end( end )
        , _nbin( counts.extent( 0 ) )
        , _counts( counts )
        , _offsets( offsets )
        , _permute_vector( permute_vector )
    {
    }

    KOKKOS_INLINE_FUNCTION
    int numBin() const { return _nbin; }

    KOKKOS_INLINE_FUNCTION
    int binSize( const size_type bin_id ) const { return _counts( bin_id ); }

    KOKKOS_INLINE_FUNCTION
    size_type binOffset( const size_type bin_id ) const
    {
        return _offsets( bin_id );
    }

    KOKKOS_INLINE_FUNCTION
    size_type permutation( const size_type tuple_id ) const
    {
        return _permute_vector( tuple_id );
    }

    KOKKOS_INLINE_FUNCTION
    std::size_t rangeBegin() const { return _begin; }

    KOKKOS_INLINE_FUNCTION
    std::size_t rangeEnd() const { return _end; }
 
  private:
    std::size_t _begin;
    std::size_t _end;
    int _nbin;
    CountView _counts;
    OffsetView _offsets;
    OffsetView _permute_vector;
};
 
//---------------------------------------------------------------------------//
template <typename>
struct is_binning_data : public std::false_type
{
};
 
template <typename MemorySpace>
struct is_binning_data<BinningData<MemorySpace>> : public std::true_type
{
};

template <typename MemorySpace>
struct is_binning_data<const BinningData<MemorySpace>> : public std::true_type
{
};
 
namespace Impl
{
//---------------------------------------------------------------------------//
template <class KeyViewType, class Comparator,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto kokkosBinSort( KeyViewType keys, Comparator comp,
                    const bool sort_within_bins, const std::size_t begin,
                    const std::size_t end )
{
    Kokkos::Profiling::ScopedRegion region( "Cabana::BinSort" );
    using memory_space = typename KeyViewType::memory_space;
    static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
 
    Kokkos::BinSort<KeyViewType, Comparator> bin_sort( keys, begin, end, comp,
                                                       sort_within_bins );
    bin_sort.create_permute_vector();
 
    return BinningData<memory_space>( begin, end, bin_sort.get_bin_count(),
                                      bin_sort.get_bin_offsets(),
                                      bin_sort.get_permute_vector() );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType,
          class ExecutionSpace = typename KeyViewType::execution_space>
Kokkos::MinMaxScalar<typename KeyViewType::non_const_value_type>
keyMinMax( KeyViewType keys, const std::size_t begin, const std::size_t end )
{
    Kokkos::Profiling::ScopedRegion region( "Cabana::keyMinMax" );
 
    using memory_space = typename KeyViewType::memory_space;
    static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
 
    using KeyValueType = typename KeyViewType::non_const_value_type;
    Kokkos::MinMaxScalar<KeyValueType> result;
    Kokkos::MinMax<KeyValueType> reducer( result );
    Kokkos::parallel_reduce(
        "Cabana::keyMinMax", Kokkos::RangePolicy<ExecutionSpace>( begin, end ),
        KOKKOS_LAMBDA( std::size_t i, decltype( result )& local_minmax ) {
            auto const val = keys( i );
            if ( val < local_minmax.min_val )
            {
                local_minmax.min_val = val;
            }
            if ( val > local_minmax.max_val )
            {
                local_minmax.max_val = val;
            }
        },
        reducer );
    Kokkos::fence();
 
    return result;
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto kokkosBinSort1d( KeyViewType keys, const int nbin,
                      const bool sort_within_bins, const std::size_t begin,
                      const std::size_t end )
{
    // Find the minimum and maximum key values.
    auto key_bounds =
        Impl::keyMinMax<KeyViewType, ExecutionSpace>( keys, begin, end );
 
    // Create a sorting comparator.
    Kokkos::BinOp1D<KeyViewType> comp( nbin, key_bounds.min_val,
                                       key_bounds.max_val );
 
    // BinSort
    return kokkosBinSort<KeyViewType, decltype( comp ), ExecutionSpace>(
        keys, comp, sort_within_bins, begin, end );
}
 
//---------------------------------------------------------------------------//
 
} // end namespace Impl
 
//---------------------------------------------------------------------------//
template <class KeyViewType, class Comparator,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto sortByKeyWithComparator(
    KeyViewType keys, Comparator comp, const std::size_t begin,
    const std::size_t end,
    typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                            int>::type* = 0 )
{
    auto bin_data = Impl::kokkosBinSort<KeyViewType, ExecutionSpace>(
        keys, comp, true, begin, end );
    return bin_data.permuteVector();
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType, class Comparator,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto sortByKeyWithComparator(
    KeyViewType keys, Comparator comp,
    typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                            int>::type* = 0 )
{
    Impl::kokkosBinSort<KeyViewType, ExecutionSpace>( keys, comp, true, 0,
                                                      keys.extent( 0 ) );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType, class Comparator,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto binByKeyWithComparator(
    KeyViewType keys, Comparator comp, const std::size_t begin,
    const std::size_t end,
    typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                            int>::type* = 0 )
{
    return Impl::kokkosBinSort<KeyViewType, ExecutionSpace>( keys, comp, false,
                                                             begin, end );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType, class Comparator,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto binByKeyWithComparator(
    KeyViewType keys, Comparator comp,
    typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                            int>::type* = 0 )
{
    return Impl::kokkosBinSort<KeyViewType, ExecutionSpace>(
        keys, comp, false, 0, keys.extent( 0 ) );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto sortByKey( KeyViewType keys, const std::size_t begin,
                const std::size_t end,
                typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                                        int>::type* = 0 )
{
    int nbin = ( end - begin ) / 2;
    return Impl::kokkosBinSort1d<KeyViewType, ExecutionSpace>( keys, nbin, true,
                                                               begin, end );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto sortByKey( KeyViewType keys,
                typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                                        int>::type* = 0 )
{
    return sortByKey<KeyViewType, ExecutionSpace>( keys, 0, keys.extent( 0 ) );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto binByKey( KeyViewType keys, const int nbin, const std::size_t begin,
               const std::size_t end,
               typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                                       int>::type* = 0 )
{
    return Impl::kokkosBinSort1d<KeyViewType, ExecutionSpace>(
        keys, nbin, false, begin, end );
}
 
//---------------------------------------------------------------------------//
template <class KeyViewType,
          class ExecutionSpace = typename KeyViewType::execution_space>
auto binByKey( KeyViewType keys, const int nbin,
               typename std::enable_if<( Kokkos::is_view<KeyViewType>::value ),
                                       int>::type* = 0 )
{
    return Impl::kokkosBinSort1d<KeyViewType, ExecutionSpace>(
        keys, nbin, false, 0, keys.extent( 0 ) );
}
 
//---------------------------------------------------------------------------//
template <class SliceType,
          class ExecutionSpace = typename SliceType::execution_space>
auto sortByKey(
    SliceType slice, const std::size_t begin, const std::size_t end,
    typename std::enable_if<( is_slice<SliceType>::value ), int>::type* = 0 )
{
    Kokkos::View<typename SliceType::value_type*,
                 typename SliceType::memory_space>
        keys( Kokkos::ViewAllocateWithoutInitializing( "slice_keys" ),
              slice.size() );
 
    copySliceToView( keys, slice, 0, slice.size() );
    return sortByKey<decltype( keys ), ExecutionSpace>( keys, begin, end );
}
 
//---------------------------------------------------------------------------//
template <class SliceType,
          class ExecutionSpace = typename SliceType::execution_space>
auto sortByKey(
    SliceType slice,
    typename std::enable_if<( is_slice<SliceType>::value ), int>::type* = 0 )
{
    return sortByKey<SliceType, ExecutionSpace>( slice, 0, slice.size() );
}
 
//---------------------------------------------------------------------------//
template <class SliceType,
          class ExecutionSpace = typename SliceType::execution_space>
auto binByKey(
    SliceType slice, const int nbin, const std::size_t begin,
    const std::size_t end,
    typename std::enable_if<( is_slice<SliceType>::value ), int>::type* = 0 )
{
    Kokkos::View<typename SliceType::value_type*,
                 typename SliceType::memory_space>
        keys( Kokkos::ViewAllocateWithoutInitializing( "slice_keys" ),
              slice.size() );
 
    copySliceToView( keys, slice, 0, slice.size() );
    return binByKey<decltype( keys ), ExecutionSpace>( keys, nbin, begin, end );
}
 
//---------------------------------------------------------------------------//
template <class SliceType,
          class ExecutionSpace = typename SliceType::execution_space>
auto binByKey(
    SliceType slice, const int nbin,
    typename std::enable_if<( is_slice<SliceType>::value ), int>::type* = 0 )
{
    return binByKey<SliceType, ExecutionSpace>( slice, nbin, 0, slice.size() );
}
 
//---------------------------------------------------------------------------//
template <class BinningDataType, class AoSoA_t,
          class ExecutionSpace = typename BinningDataType::execution_space>
void permute(
    const BinningDataType& binning_data, AoSoA_t& aosoa,
    typename std::enable_if<( is_binning_data<BinningDataType>::value &&
                              is_aosoa<AoSoA_t>::value ),
                            int>::type* = 0 )
{
    Kokkos::Profiling::ScopedRegion region( "Cabana::permute" );
 
    using memory_space = typename BinningDataType::memory_space;
    static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
 
    auto begin = binning_data.rangeBegin();
    auto end = binning_data.rangeEnd();
 
    using memory_space = typename BinningDataType::memory_space;
    Kokkos::View<typename AoSoA_t::tuple_type*, memory_space> scratch_tuples(
        Kokkos::ViewAllocateWithoutInitializing( "scratch_tuples" ),
        end - begin );
 
    auto permute_to_scratch = KOKKOS_LAMBDA( const std::size_t i )
    {
        scratch_tuples( i - begin ) =
            aosoa.getTuple( binning_data.permutation( i - begin ) );
    };
    Kokkos::parallel_for( "Cabana::kokkosBinSort::permute_to_scratch",
                          Kokkos::RangePolicy<ExecutionSpace>( begin, end ),
                          permute_to_scratch );
    Kokkos::fence();
 
    auto copy_back = KOKKOS_LAMBDA( const std::size_t i )
    {
        aosoa.setTuple( i, scratch_tuples( i - begin ) );
    };
    Kokkos::parallel_for( "Cabana::kokkosBinSort::copy_back",
                          Kokkos::RangePolicy<ExecutionSpace>( begin, end ),
                          copy_back );
    Kokkos::fence();
}
 
//---------------------------------------------------------------------------//
template <class BinningDataType, class SliceType,
          class ExecutionSpace = typename BinningDataType::execution_space>
void permute(
    const BinningDataType& binning_data, SliceType& slice,
    typename std::enable_if<( is_binning_data<BinningDataType>::value &&
                              is_slice<SliceType>::value ),
                            int>::type* = 0 )
{
    Kokkos::Profiling::ScopedRegion region( "Cabana::permute" );
 
    using memory_space = typename BinningDataType::memory_space;
    static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
 
    auto begin = binning_data.rangeBegin();
    auto end = binning_data.rangeEnd();
 
    // Get the number of components in the slice.
    std::size_t num_comp = 1;
    for ( std::size_t d = 2; d < slice.viewRank(); ++d )
        num_comp *= slice.extent( d );
 
    // Get the raw slice data.
    auto slice_data = slice.data();
 
    Kokkos::View<typename SliceType::value_type**, memory_space> scratch_array(
        Kokkos::ViewAllocateWithoutInitializing( "scratch_array" ), end - begin,
        num_comp );
 
    auto permute_to_scratch = KOKKOS_LAMBDA( const std::size_t i )
    {
        auto permute_i = binning_data.permutation( i - begin );
        auto s = SliceType::index_type::s( permute_i );
        auto a = SliceType::index_type::a( permute_i );
        std::size_t slice_offset = s * slice.stride( 0 ) + a;
        for ( std::size_t n = 0; n < num_comp; ++n )
            scratch_array( i - begin, n ) =
                slice_data[slice_offset + SliceType::vector_length * n];
    };
    Kokkos::parallel_for( "Cabana::kokkosBinSort::permute_to_scratch",
                          Kokkos::RangePolicy<ExecutionSpace>( begin, end ),
                          permute_to_scratch );
    Kokkos::fence();
 
    auto copy_back = KOKKOS_LAMBDA( const std::size_t i )
    {
        auto s = SliceType::index_type::s( i );
        auto a = SliceType::index_type::a( i );
        std::size_t slice_offset = s * slice.stride( 0 ) + a;
        for ( std::size_t n = 0; n < num_comp; ++n )
            slice_data[slice_offset + SliceType::vector_length * n] =
                scratch_array( i - begin, n );
    };
    Kokkos::parallel_for( "Cabana::kokkosBinSort::copy_back",
                          Kokkos::RangePolicy<ExecutionSpace>( begin, end ),
                          copy_back );
    Kokkos::fence();
}
 
//---------------------------------------------------------------------------//
template <class BinningDataType, class ViewType,
          class MemorySpace = typename BinningDataType::memory_space>
void permute(
    const BinningDataType& binning_data, ViewType& view,
    typename std::enable_if<( is_binning_data<BinningDataType>::value &&
                              Kokkos::is_view<ViewType>::value ),
                            int>::type* = 0 )
{
    Kokkos::Profiling::pushRegion( "Cabana::permute" );
 
    auto begin = binning_data.rangeBegin();
    auto end = binning_data.rangeEnd();
 
    // Get the number of components (1 for View).
    std::size_t num_comp = 1;
    for ( std::size_t d = 1; d < view.rank; ++d )
        num_comp *= view.extent( d );
 
    Kokkos::View<typename ViewType::value_type**, MemorySpace> scratch(
        Kokkos::ViewAllocateWithoutInitializing( "scratch" ), end - begin,
        num_comp );
 
    auto permute_to_scratch = KOKKOS_LAMBDA( const std::size_t i )
    {
        auto permute_i = binning_data.permutation( i - begin );
        for ( std::size_t n = 0; n < num_comp; ++n )
            scratch( i - begin, n ) = view( permute_i, n );
    };
    auto policy =
        Kokkos::RangePolicy<typename BinningDataType::execution_space>( begin,
                                                                        end );
    Kokkos::parallel_for( "Cabana::kokkosBinSort::permute_to_scratch", policy,
                          permute_to_scratch );
    Kokkos::fence();
 
    auto copy_back = KOKKOS_LAMBDA( const std::size_t i )
    {
        for ( std::size_t n = 0; n < num_comp; ++n )
            view( i, n ) = scratch( i - begin, n );
    };
    Kokkos::parallel_for( "Cabana::kokkosBinSort::copy_back", policy,
                          copy_back );
    Kokkos::fence();
 
    Kokkos::Profiling::popRegion();
}
 
//---------------------------------------------------------------------------//
 
} // end namespace Cabana
 
#endif // end CABANA_SORT_HPP