block_radix_sort.hpp

// Copyright (c) 2017-2021 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#ifndef ROCPRIM_BLOCK_BLOCK_RADIX_SORT_HPP_
#define ROCPRIM_BLOCK_BLOCK_RADIX_SORT_HPP_

#include <type_traits>

#include "../config.hpp"
#include "../detail/various.hpp"
#include "../detail/radix_sort.hpp"
#include "../warp/detail/warp_scan_crosslane.hpp"

#include "../intrinsics.hpp"
#include "../functional.hpp"
#include "../types.hpp"

#include "block_exchange.hpp"
#include "block_radix_rank.hpp"


BEGIN_ROCPRIM_NAMESPACE

template<
    class Key,
    unsigned int BlockSizeX,
    unsigned int ItemsPerThread,
    class Value = empty_type,
    unsigned int BlockSizeY = 1,
    unsigned int BlockSizeZ = 1
>
class block_radix_sort
{
    static constexpr unsigned int BlockSize           = BlockSizeX * BlockSizeY * BlockSizeZ;
    static constexpr bool         with_values         = !std::is_same<Value, empty_type>::value;
    static constexpr unsigned int radix_bits_per_pass = 4;

    using bit_key_type    = typename ::rocprim::detail::radix_key_codec<Key>::bit_key_type;
    using block_rank_type = ::rocprim::block_radix_rank<BlockSizeX,
                                                        radix_bits_per_pass,
                                                        block_radix_rank_algorithm::basic_memoize,
                                                        BlockSizeY,
                                                        BlockSizeZ>;
    using bit_keys_exchange_type = ::rocprim::block_exchange<bit_key_type, BlockSizeX, ItemsPerThread, BlockSizeY, BlockSizeZ>;
    using values_exchange_type
        = ::rocprim::block_exchange<Value, BlockSizeX, ItemsPerThread, BlockSizeY, BlockSizeZ>;

    // Struct used for creating a raw_storage object for this primitive's temporary storage.
    union storage_type_
    {
        typename bit_keys_exchange_type::storage_type bit_keys_exchange;
        typename values_exchange_type::storage_type   values_exchange;
        typename block_rank_type::storage_type        rank;
    };

public:

    #ifndef DOXYGEN_SHOULD_SKIP_THIS // hides storage_type implementation for Doxygen
    using storage_type = detail::raw_storage<storage_type_>;
    #else
    using storage_type = storage_type_; // only for Doxygen
    #endif

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort(Key (&keys)[ItemsPerThread],
              storage_type& storage,
              unsigned int begin_bit = 0,
              unsigned int end_bit = 8 * sizeof(Key))
    {
        empty_type values[ItemsPerThread];
        sort_impl<false>(keys, values, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort(Key (&keys)[ItemsPerThread],
              unsigned int begin_bit = 0,
              unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort(keys, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_desc(Key (&keys)[ItemsPerThread],
                   storage_type& storage,
                   unsigned int begin_bit = 0,
                   unsigned int end_bit = 8 * sizeof(Key))
    {
        empty_type values[ItemsPerThread];
        sort_impl<true>(keys, values, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort_desc(Key (&keys)[ItemsPerThread],
                   unsigned int begin_bit = 0,
                   unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort_desc(keys, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort(Key (&keys)[ItemsPerThread],
              typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
              storage_type& storage,
              unsigned int begin_bit = 0,
              unsigned int end_bit = 8 * sizeof(Key))
    {
        sort_impl<false>(keys, values, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort(Key (&keys)[ItemsPerThread],
              typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
              unsigned int begin_bit = 0,
              unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort(keys, values, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_desc(Key (&keys)[ItemsPerThread],
                   typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
                   storage_type& storage,
                   unsigned int begin_bit = 0,
                   unsigned int end_bit = 8 * sizeof(Key))
    {
        sort_impl<true>(keys, values, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort_desc(Key (&keys)[ItemsPerThread],
                   typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
                   unsigned int begin_bit = 0,
                   unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort_desc(keys, values, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_to_striped(Key (&keys)[ItemsPerThread],
                         storage_type& storage,
                         unsigned int begin_bit = 0,
                         unsigned int end_bit = 8 * sizeof(Key))
    {
        empty_type values[ItemsPerThread];
        sort_impl<false, true>(keys, values, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort_to_striped(Key (&keys)[ItemsPerThread],
                         unsigned int begin_bit = 0,
                         unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort_to_striped(keys, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_desc_to_striped(Key (&keys)[ItemsPerThread],
                              storage_type& storage,
                              unsigned int begin_bit = 0,
                              unsigned int end_bit = 8 * sizeof(Key))
    {
        empty_type values[ItemsPerThread];
        sort_impl<true, true>(keys, values, storage, begin_bit, end_bit);
    }

    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort_desc_to_striped(Key (&keys)[ItemsPerThread],
                              unsigned int begin_bit = 0,
                              unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort_desc_to_striped(keys, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_to_striped(Key (&keys)[ItemsPerThread],
                         typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
                         storage_type& storage,
                         unsigned int begin_bit = 0,
                         unsigned int end_bit = 8 * sizeof(Key))
    {
        sort_impl<false, true>(keys, values, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort_to_striped(Key (&keys)[ItemsPerThread],
                         typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
                         unsigned int begin_bit = 0,
                         unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort_to_striped(keys, values, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_desc_to_striped(Key (&keys)[ItemsPerThread],
                              typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
                              storage_type& storage,
                              unsigned int begin_bit = 0,
                              unsigned int end_bit = 8 * sizeof(Key))
    {
        sort_impl<true, true>(keys, values, storage, begin_bit, end_bit);
    }

    template<bool WithValues = with_values>
    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE
    void sort_desc_to_striped(Key (&keys)[ItemsPerThread],
                              typename std::enable_if<WithValues, Value>::type (&values)[ItemsPerThread],
                              unsigned int begin_bit = 0,
                              unsigned int end_bit = 8 * sizeof(Key))
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        sort_desc_to_striped(keys, values, storage, begin_bit, end_bit);
    }

private:

    template<bool Descending, bool ToStriped = false, class SortedValue>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void sort_impl(Key (&keys)[ItemsPerThread],
                   SortedValue (&values)[ItemsPerThread],
                   storage_type& storage,
                   unsigned int begin_bit,
                   unsigned int end_bit)
    {
        using key_codec = ::rocprim::detail::radix_key_codec<Key, Descending>;

        bit_key_type bit_keys[ItemsPerThread];
        ROCPRIM_UNROLL
        for(unsigned int i = 0; i < ItemsPerThread; i++)
        {
            bit_keys[i] = key_codec::encode(keys[i]);
        }

        while(true)
        {
            const int pass_bits = min(radix_bits_per_pass, end_bit - begin_bit);

            unsigned int ranks[ItemsPerThread];
            block_rank_type().rank_keys(
                bit_keys,
                ranks,
                storage.get().rank,
                [begin_bit, pass_bits](const bit_key_type& key)
                { return key_codec::extract_digit(key, begin_bit, pass_bits); });
            begin_bit += radix_bits_per_pass;

            exchange_keys(storage, bit_keys, ranks);
            exchange_values(storage, values, ranks);

            if(begin_bit >= end_bit)
                break;

            // Synchronization required to make bock_rank wait on the next iteration.
            ::rocprim::syncthreads();
        }

        if ROCPRIM_IF_CONSTEXPR(ToStriped)
        {
            to_striped_keys(storage, bit_keys);
            to_striped_values(storage, values);
        }

        ROCPRIM_UNROLL
        for(unsigned int i = 0; i < ItemsPerThread; i++)
        {
            keys[i] = key_codec::decode(bit_keys[i]);
        }
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void exchange_keys(storage_type& storage,
                       bit_key_type (&bit_keys)[ItemsPerThread],
                       const unsigned int (&ranks)[ItemsPerThread])
    {
        storage_type_& storage_ = storage.get();
        ::rocprim::syncthreads(); // Storage will be reused (union), synchronization is needed
        bit_keys_exchange_type().scatter_to_blocked(bit_keys, bit_keys, ranks, storage_.bit_keys_exchange);
    }

    template<class SortedValue>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void exchange_values(storage_type& storage,
                         SortedValue (&values)[ItemsPerThread],
                         const unsigned int (&ranks)[ItemsPerThread])
    {
        storage_type_& storage_ = storage.get();
        ::rocprim::syncthreads(); // Storage will be reused (union), synchronization is needed
        values_exchange_type().scatter_to_blocked(values, values, ranks, storage_.values_exchange);
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void exchange_values(storage_type& storage,
                         empty_type (&values)[ItemsPerThread],
                         const unsigned int (&ranks)[ItemsPerThread])
    {
        (void) storage;
        (void) values;
        (void) ranks;
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void to_striped_keys(storage_type& storage,
                         bit_key_type (&bit_keys)[ItemsPerThread])
    {
        storage_type_& storage_ = storage.get();
        ::rocprim::syncthreads(); // Storage will be reused (union), synchronization is needed
        bit_keys_exchange_type().blocked_to_striped(bit_keys, bit_keys, storage_.bit_keys_exchange);
    }

    template<class SortedValue>
    ROCPRIM_DEVICE ROCPRIM_INLINE
    void to_striped_values(storage_type& storage,
                           SortedValue (&values)[ItemsPerThread])
    {
        storage_type_& storage_ = storage.get();
        ::rocprim::syncthreads(); // Storage will be reused (union), synchronization is needed
        values_exchange_type().blocked_to_striped(values, values, storage_.values_exchange);
    }

    ROCPRIM_DEVICE ROCPRIM_INLINE
    void to_striped_values(storage_type& storage,
                           empty_type * values)
    {
        (void) storage;
        (void) values;
    }
};

END_ROCPRIM_NAMESPACE

// end of group blockmodule

#endif // ROCPRIM_BLOCK_BLOCK_RADIX_SORT_HPP_