block_radix_rank.hpp

// Copyright (c) 2022 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_RANK_HPP_
#define ROCPRIM_BLOCK_BLOCK_RADIX_RANK_HPP_

#include "../config.hpp"
#include "../functional.hpp"

#include "../detail/radix_sort.hpp"

#include "block_scan.hpp"

#include "detail/block_radix_rank_basic.hpp"
#include "detail/block_radix_rank_match.hpp"


BEGIN_ROCPRIM_NAMESPACE

enum class block_radix_rank_algorithm
{
    basic,
    basic_memoize,
    match,
    default_algorithm = basic,
};

namespace detail
{
// Selector for block radix rank algorithm that gives the radix rank implementation.
template<block_radix_rank_algorithm Algorithm>
struct select_block_radix_rank_impl;

template<>
struct select_block_radix_rank_impl<block_radix_rank_algorithm::basic>
{
    template<unsigned int BlockSizeX,
             unsigned int RadixBits,
             unsigned int BlockSizeY,
             unsigned int BlockSizeZ>
    using type = block_radix_rank<BlockSizeX, RadixBits, false, BlockSizeY, BlockSizeZ>;
};

template<>
struct select_block_radix_rank_impl<block_radix_rank_algorithm::basic_memoize>
{
    template<unsigned int BlockSizeX,
             unsigned int RadixBits,
             unsigned int BlockSizeY,
             unsigned int BlockSizeZ>
    using type = block_radix_rank<BlockSizeX, RadixBits, true, BlockSizeY, BlockSizeZ>;
};

template<>
struct select_block_radix_rank_impl<block_radix_rank_algorithm::match>
{
    template<unsigned int BlockSizeX,
             unsigned int RadixBits,
             unsigned int BlockSizeY,
             unsigned int BlockSizeZ>
    using type = block_radix_rank_match<BlockSizeX, RadixBits, BlockSizeY, BlockSizeZ>;
};
} // namespace detail

template<unsigned int               BlockSizeX,
         unsigned int               RadixBits,
         block_radix_rank_algorithm Algorithm  = block_radix_rank_algorithm::default_algorithm,
         unsigned int               BlockSizeY = 1,
         unsigned int               BlockSizeZ = 1>
class block_radix_rank
#ifndef DOXYGEN_SHOULD_SKIP_THIS
    : private detail::select_block_radix_rank_impl<
          Algorithm>::template type<BlockSizeX, RadixBits, BlockSizeY, BlockSizeZ>
#endif
{
    using base_type = typename detail::select_block_radix_rank_impl<
        Algorithm>::template type<BlockSizeX, RadixBits, BlockSizeY, BlockSizeZ>;

public:
    static constexpr unsigned int digits_per_thread = base_type::digits_per_thread;

    using storage_type = typename base_type::storage_type;

    template<typename Key, unsigned ItemsPerThread>
    ROCPRIM_DEVICE void rank_keys(const Key (&keys)[ItemsPerThread],
                                  unsigned int (&ranks)[ItemsPerThread],
                                  storage_type& storage,
                                  unsigned int  begin_bit = 0,
                                  unsigned int  pass_bits = RadixBits)
    {
        base_type::rank_keys(keys, ranks, storage, begin_bit, pass_bits);
    }

    template<typename Key, unsigned ItemsPerThread>
    ROCPRIM_DEVICE void rank_keys(const Key (&keys)[ItemsPerThread],
                                  unsigned int (&ranks)[ItemsPerThread],
                                  unsigned int begin_bit = 0,
                                  unsigned int pass_bits = RadixBits)
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        base_type::rank_keys(keys, ranks, storage, begin_bit, pass_bits);
    }

    template<typename Key, unsigned ItemsPerThread>
    ROCPRIM_DEVICE void rank_keys_desc(const Key (&keys)[ItemsPerThread],
                                       unsigned int (&ranks)[ItemsPerThread],
                                       storage_type& storage,
                                       unsigned int  begin_bit = 0,
                                       unsigned int  pass_bits = RadixBits)
    {
        base_type::rank_keys_desc(keys, ranks, storage, begin_bit, pass_bits);
    }

    template<typename Key, unsigned ItemsPerThread>
    ROCPRIM_DEVICE void rank_keys_desc(const Key (&keys)[ItemsPerThread],
                                       unsigned int (&ranks)[ItemsPerThread],
                                       unsigned int begin_bit = 0,
                                       unsigned int pass_bits = RadixBits)
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        base_type::rank_keys_desc(keys, ranks, storage, begin_bit, pass_bits);
    }

    template<typename Key, unsigned ItemsPerThread, typename DigitExtractor>
    ROCPRIM_DEVICE void rank_keys(const Key (&keys)[ItemsPerThread],
                                  unsigned int (&ranks)[ItemsPerThread],
                                  storage_type&  storage,
                                  DigitExtractor digit_extractor)
    {
        base_type::rank_keys(keys, ranks, storage, digit_extractor);
    }

    template<typename Key, unsigned ItemsPerThread, typename DigitExtractor>
    ROCPRIM_DEVICE void rank_keys(const Key (&keys)[ItemsPerThread],
                                  unsigned int (&ranks)[ItemsPerThread],
                                  DigitExtractor digit_extractor)
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        base_type::rank_keys(keys, ranks, storage, digit_extractor);
    }

    template<typename Key, unsigned ItemsPerThread, typename DigitExtractor>
    ROCPRIM_DEVICE void rank_keys_desc(const Key (&keys)[ItemsPerThread],
                                       unsigned int (&ranks)[ItemsPerThread],
                                       storage_type&  storage,
                                       DigitExtractor digit_extractor)
    {
        base_type::rank_keys_desc(keys, ranks, storage, digit_extractor);
    }

    template<typename Key, unsigned ItemsPerThread, typename DigitExtractor>
    ROCPRIM_DEVICE void rank_keys_desc(const Key (&keys)[ItemsPerThread],
                                       unsigned int (&ranks)[ItemsPerThread],
                                       DigitExtractor digit_extractor)
    {
        ROCPRIM_SHARED_MEMORY storage_type storage;
        base_type::rank_keys_desc(keys, ranks, storage, digit_extractor);
    }

    template<typename Key, unsigned ItemsPerThread, typename DigitExtractor>
    ROCPRIM_DEVICE void rank_keys(const Key (&keys)[ItemsPerThread],
                                  unsigned int (&ranks)[ItemsPerThread],
                                  storage_type&  storage,
                                  DigitExtractor digit_extractor,
                                  unsigned int (&prefix)[digits_per_thread],
                                  unsigned int (&counts)[digits_per_thread])
    {
        base_type::rank_keys(keys, ranks, storage, digit_extractor, prefix, counts);
    }
};

END_ROCPRIM_NAMESPACE

// end of group blockmodule

#endif // ROCPRIM_BLOCK_BLOCK_RADIX_RANK_HPP_