// -*- C++ -*-
//===-- parallel_backend_sycl_radix_sort.h --------------------------------===//
//
// Copyright (C) Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// This file incorporates work covered by the following copyright and permission
// notice:
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
//
//===----------------------------------------------------------------------===//

#ifndef _ONEDPL_PARALLEL_BACKEND_SYCL_RADIX_SORT_H
#define _ONEDPL_PARALLEL_BACKEND_SYCL_RADIX_SORT_H

#include <limits>
#include <type_traits>
#include <utility>
#include <cstdint>
#include <algorithm>
#include <functional> // for std::invoke

#include "sycl_defs.h"
#include "parallel_backend_sycl_utils.h"
#include "execution_sycl_defs.h"

#include "sycl_traits.h" //SYCL traits specialization for some oneDPL types.

namespace oneapi
{
namespace dpl
{
namespace __par_backend_hetero
{
//------------------------------------------------------------------------
// radix sort: bitwise order-preserving conversions to unsigned integrals
//------------------------------------------------------------------------

template <bool __is_ascending>
bool
__order_preserving_cast(bool __val)
{
    if constexpr (__is_ascending)
        return __val;
    else
        return !__val;
}

template <bool __is_ascending, typename _UInt, ::std::enable_if_t<::std::is_unsigned_v<_UInt>, int> = 0>
_UInt
__order_preserving_cast(_UInt __val)
{
    if constexpr (__is_ascending)
        return __val;
    else
        return ~__val; //bitwise inversion
}

template <bool __is_ascending, typename _Int,
          ::std::enable_if_t<::std::is_integral_v<_Int> && ::std::is_signed_v<_Int>, int> = 0>
::std::make_unsigned_t<_Int>
__order_preserving_cast(_Int __val)
{
    using _UInt = ::std::make_unsigned_t<_Int>;
    // mask: 100..0 for ascending, 011..1 for descending
    constexpr _UInt __mask =
        (__is_ascending) ? _UInt(1) << ::std::numeric_limits<_Int>::digits : ::std::numeric_limits<_UInt>::max() >> 1;
    return __val ^ __mask;
}

template <bool __is_ascending>
::std::uint16_t
__order_preserving_cast(sycl::half __val)
{
    ::std::uint16_t __uint16_val = oneapi::dpl::__internal::__dpl_bit_cast<::std::uint16_t>(__val);
    ::std::uint16_t __mask;
    // __uint16_val >> 15 takes the sign bit of the original value
    if constexpr (__is_ascending)
        __mask = (__uint16_val >> 15 == 0) ? 0x8000u : 0xFFFFu;
    else
        __mask = (__uint16_val >> 15 == 0) ? 0x7FFFu : ::std::uint16_t(0);
    return __uint16_val ^ __mask;
}

template <bool __is_ascending, typename _Float,
          ::std::enable_if_t<::std::is_floating_point_v<_Float> && sizeof(_Float) == sizeof(::std::uint32_t), int> = 0>
::std::uint32_t
__order_preserving_cast(_Float __val)
{
    ::std::uint32_t __uint32_val = oneapi::dpl::__internal::__dpl_bit_cast<::std::uint32_t>(__val);
    ::std::uint32_t __mask;
    // __uint32_val >> 31 takes the sign bit of the original value
    if constexpr (__is_ascending)
        __mask = (__uint32_val >> 31 == 0) ? 0x80000000u : 0xFFFFFFFFu;
    else
        __mask = (__uint32_val >> 31 == 0) ? 0x7FFFFFFFu : ::std::uint32_t(0);
    return __uint32_val ^ __mask;
}

template <bool __is_ascending, typename _Float,
          ::std::enable_if_t<::std::is_floating_point_v<_Float> && sizeof(_Float) == sizeof(::std::uint64_t), int> = 0>
::std::uint64_t
__order_preserving_cast(_Float __val)
{
    ::std::uint64_t __uint64_val = oneapi::dpl::__internal::__dpl_bit_cast<::std::uint64_t>(__val);
    ::std::uint64_t __mask;
    // __uint64_val >> 63 takes the sign bit of the original value
    if constexpr (__is_ascending)
        __mask = (__uint64_val >> 63 == 0) ? 0x8000000000000000u : 0xFFFFFFFFFFFFFFFFu;
    else
        __mask = (__uint64_val >> 63 == 0) ? 0x7FFFFFFFFFFFFFFFu : ::std::uint64_t(0);
    return __uint64_val ^ __mask;
}

//------------------------------------------------------------------------
// radix sort: bucket functions
//------------------------------------------------------------------------

// get number of buckets (size of radix bits) in T
template <typename _T>
constexpr ::std::uint32_t
__get_buckets_in_type(::std::uint32_t __radix_bits)
{
    return oneapi::dpl::__internal::__dpl_ceiling_div(sizeof(_T) * ::std::numeric_limits<unsigned char>::digits,
                                                      __radix_bits);
}

// get bits value (bucket) in a certain radix position
template <std::uint32_t __radix_mask, typename _T>
std::uint32_t
__get_bucket(_T __value, std::uint32_t __radix_offset)
{
    return (__value >> __radix_offset) & _T(__radix_mask);
}

//------------------------------------------------------------------------
// radix sort: kernel names
//------------------------------------------------------------------------

template <std::uint32_t, bool, typename... _Name>
class __radix_sort_count_kernel;

template <std::uint32_t, typename... _Name>
class __radix_sort_scan_kernel;

template <std::uint32_t, bool, typename... _Name>
class __radix_sort_reorder_kernel;

template <typename... _Name>
class __radix_sort_copy_back_kernel;

// Helper for SLM indexing in count kernel. Layout stores radix states contiguously per work-item
// for cache locality during counting, then reorganizes for reduction.
template <std::uint32_t __packing_ratio, std::uint32_t __radix_states>
struct __index_views
{
    //Number of elements of larger CounterT which each thread is responsible for resetting
    static constexpr std::uint32_t __counter_lanes = __radix_states / __packing_ratio;
    // Number of elements from the previous layout of work-items that each work item is summing in the initial reduction
    static constexpr std::uint32_t __reduction_factor =
        (__packing_ratio > __counter_lanes) ? __packing_ratio : __counter_lanes;
    // Number of radix states handled by each group
    static constexpr std::uint32_t __radix_states_per_group = __radix_states / __reduction_factor;

    // Index for uint8 bucket counters: [wg_id][radix_id]
    std::uint32_t
    __get_bucket_idx(std::uint32_t __radix_id, std::uint32_t __wg_id) const
    {
        return __wg_id * __radix_states + __radix_id;
    }

    // Index for packed uint32 counters (4 uint8s packed): [wg_id][radix_id_lane]
    std::uint32_t
    __get_bucket32_idx(std::uint32_t __radix_id_lane, std::uint32_t __wg_id) const
    {
        return __wg_id * (__radix_states / __packing_ratio) + __radix_id_lane;
    }

    // Index for reduction phase: [radix_id][partial_sum_id]
    // Stride per radix state = __wg_size / __reduction_factor
    std::uint32_t
    __get_count_idx(std::uint32_t __workgroup_size, std::uint32_t __radix_id, std::uint32_t __wg_id) const
    {
        return __radix_id * (__workgroup_size / __reduction_factor) + __wg_id;
    }
};

//-----------------------------------------------------------------------
// radix sort: count kernel helper
// Template helper function to count elements from a given input range.
// Called with a single runtime branch to select the input range.
//-----------------------------------------------------------------------
template <std::uint32_t __radix_bits, bool __is_ascending, typename _InputRange, typename _Proj, typename _IndexViews>
void
__radix_sort_count_impl(_InputRange& __input, _Proj __proj, std::uint32_t __radix_offset, std::size_t __sg_chunk_start,
                        std::size_t __sg_chunk_end, std::size_t __full_end, std::uint32_t __sg_size,
                        std::uint32_t __sg_local_id, std::size_t __wg_size, std::size_t __self_lidx,
                        std::uint8_t* __slm_buckets, _IndexViews __views)
{
    constexpr std::uint32_t __unroll_elements = 8;

    // Full iterations - no bounds checking needed
    for (std::size_t __base_idx = __sg_chunk_start + __sg_local_id; __base_idx < __full_end;
         __base_idx += __sg_size * __unroll_elements)
    {
        _ONEDPL_PRAGMA_UNROLL
        for (std::size_t __unroll = 0; __unroll < __unroll_elements; ++__unroll)
        {
            auto __val = __order_preserving_cast<__is_ascending>(
                std::invoke(__proj, __input[__base_idx + __unroll * __sg_size]));
            std::uint32_t __bucket = __get_bucket<(1 << __radix_bits) - 1>(__val, __radix_offset);
            ++__slm_buckets[__views.__get_bucket_idx(__bucket, __self_lidx)];
        }
    }

    // Remainder - at most one partial block per subgroup
    if (__full_end < __sg_chunk_end)
    {
        const std::size_t __base_idx = __full_end + __sg_local_id;
        _ONEDPL_PRAGMA_UNROLL
        for (std::size_t __unroll = 0; __unroll < __unroll_elements; ++__unroll)
        {
            std::size_t __curr_idx = __base_idx + __unroll * __sg_size;
            if (__curr_idx < __sg_chunk_end)
            {
                auto __val = __order_preserving_cast<__is_ascending>(std::invoke(__proj, __input[__curr_idx]));
                std::uint32_t __bucket = __get_bucket<(1 << __radix_bits) - 1>(__val, __radix_offset);
                ++__slm_buckets[__views.__get_bucket_idx(__bucket, __self_lidx)];
            }
        }
    }
}

//-----------------------------------------------------------------------
// radix sort: count kernel (per iteration)
//-----------------------------------------------------------------------

template <typename _KernelName, std::uint32_t __radix_bits, bool __is_ascending, std::uint32_t __unroll_elements = 8,
          typename _ValRange1, typename _ValRange2, typename _CountBuf, typename _Proj>
sycl::event
__radix_sort_count_submit(sycl::queue& __q, std::size_t __segments, std::size_t __wg_size, std::uint32_t __radix_offset,
                          bool __input_is_first, _ValRange1&& __val_rng1, _ValRange2&& __val_rng2,
                          _CountBuf& __count_buf, sycl::event __dependency_event, _Proj __proj)
{
    using _CountT = typename _CountBuf::value_type;

    // radix states used for an array storing bucket state counters
    constexpr std::uint32_t __radix_states = 1 << __radix_bits;
    // multiple uint8_t (1 byte) counters are packed in the footprint of 1 _CountT, we reuse SLM in different phases to
    // allow more parallel work to start with, still avoiding overflow as we accumulate bins from more elements together
    constexpr std::uint32_t __packing_ratio = sizeof(_CountT);
    constexpr std::uint32_t __counter_lanes = __radix_states / __packing_ratio;

    // iteration space info
    const std::size_t __n = oneapi::dpl::__ranges::__size(__val_rng1);
    const std::size_t __elem_per_segment = oneapi::dpl::__internal::__dpl_ceiling_div(__n, __segments);
    const std::size_t __no_op_flag_idx = __count_buf.size() - 1;

    //assert that we cannot overflow uint8 accumulation
    assert(oneapi::dpl::__internal::__dpl_ceiling_div(__elem_per_segment, __wg_size) < 256 &&
           "Segment size per work-group is too large to count in uint8");
    auto __count_rng =
        oneapi::dpl::__ranges::all_view<_CountT, __par_backend_hetero::access_mode::read_write>(__count_buf);

    // submit to compute arrays with local count values
    sycl::event __count_levent = __q.submit([&](sycl::handler& __hdl) {
        __hdl.depends_on(__dependency_event);

        // ensure the input data and the space for counters are accessible
        oneapi::dpl::__ranges::__require_access(__hdl, __val_rng1, __val_rng2, __count_rng);
        // an accessor per work-group with value counters from each work-item
        __dpl_sycl::__local_accessor<_CountT> __count_lacc(
            oneapi::dpl::__internal::__dpl_ceiling_div(__radix_states * __wg_size, __packing_ratio), __hdl);
        __hdl.parallel_for<_KernelName>(
            sycl::nd_range<1>(__segments * __wg_size, __wg_size), [=](sycl::nd_item<1> __self_item) {
                // item info
                const std::size_t __self_lidx = __self_item.get_local_id(0);
                const std::size_t __wgroup_idx = __self_item.get_group(0);
                const std::size_t __seg_start = __elem_per_segment * __wgroup_idx;

                // Subgroup info for SG-strided memory access pattern
                __dpl_sycl::__sub_group __sub_group = __self_item.get_sub_group();
                const std::uint32_t __sg_size = __sub_group.get_local_range()[0];
                const std::uint32_t __sg_local_id = __sub_group.get_local_linear_id();
                const std::uint32_t __num_subgroups = oneapi::dpl::__internal::__dpl_ceiling_div(__wg_size, __sg_size);
                // Compute subgroup base from work item ID to handle variable subgroup sizes
                const std::size_t __sg_base = (__self_lidx - __sg_local_id);

                _CountT* __slm_counts = &__count_lacc[0];
                std::uint8_t* __slm_buckets = reinterpret_cast<std::uint8_t*>(__slm_counts);
                __index_views<__packing_ratio, __radix_states> __views;

                constexpr std::uint32_t __radix_states_per_group = decltype(__views)::__radix_states_per_group;
                _CountT __count_arr[__radix_states_per_group] = {0};
                const std::size_t __seg_end = sycl::min(__seg_start + __elem_per_segment, __n);

                // reset SLM buckets
                _ONEDPL_PRAGMA_UNROLL
                for (std::uint32_t __i = 0; __i < __counter_lanes; ++__i)
                {
                    __slm_counts[__views.__get_bucket32_idx(__i, __self_lidx)] = 0;
                }

                // SG-strided reads: each subgroup handles a contiguous chunk of the segment
                // Within each chunk, work items access memory with stride __sg_size for coalescing
                const std::size_t __seg_size = __seg_end - __seg_start;
                const std::size_t __elems_per_sg = __seg_size / __num_subgroups;
                const std::size_t __sg_id = __sg_base / __sg_size;
                const std::size_t __sg_chunk_start = __seg_start + __sg_id * __elems_per_sg;
                // Last subgroup gets any remainder elements
                const std::size_t __sg_chunk_end =
                    (__sg_id == __num_subgroups - 1) ? __seg_end : __sg_chunk_start + __elems_per_sg;

                const std::size_t __sg_chunk_size = __sg_chunk_end - __sg_chunk_start;
                const std::size_t __full_rounds = __sg_chunk_size / (__sg_size * __unroll_elements);
                const std::size_t __full_end = __sg_chunk_start + __full_rounds * __sg_size * __unroll_elements;

                // 1. COUNTING PHASE: SLM accessed as uint8_t array (__slm_buckets)
                //    Each Work-Item (WI) has a contiguous block of 16 uint8_t counters.
                //
                //    SLM Address:    0                   16                  32                  N bytes
                //    uint8_t View:   | r0 r1 ... r14 r15 | r0 r1 ... r14 r15 | ...               |
                //    Owned by WI:    +-------WI 0--------+-------WI 1--------+ ... +---Last WI---+
                //                               |
                //                               | (Group barrier)
                //                               v

                // Single branch to select input range, then count without per-element branching
                if (__input_is_first)
                    __radix_sort_count_impl<__radix_bits, __is_ascending>(
                        __val_rng1, __proj, __radix_offset, __sg_chunk_start, __sg_chunk_end, __full_end, __sg_size,
                        __sg_local_id, __wg_size, __self_lidx, __slm_buckets, __views);
                else
                    __radix_sort_count_impl<__radix_bits, __is_ascending>(
                        __val_rng2, __proj, __radix_offset, __sg_chunk_start, __sg_chunk_end, __full_end, __sg_size,
                        __sg_local_id, __wg_size, __self_lidx, __slm_buckets, __views);

                __dpl_sycl::__group_barrier(__self_item);

                // 2. ACCUMULATION PHASE: Read uint8_t columns into _CountT Registers (__count_arr)
                //    WIs are grouped to prevent uint8_t overflow. Example: __reduction_factor = 4.
                //    WI 0 handles Radix 0-3 for WIs 0,1,2,3 (a column across WIs).
                //
                //    uint8_t reads:  WI 0's r0-r3  \     Summed into WI 0's registers:
                //                    WI 1's r0-r3  |-->  __count_arr[4] = {sum_r0, sum_r1, sum_r2, sum_r3}
                //                    WI 2's r0-r3  |     (These are 32-bit sums, safely preventing overflow)
                //                    WI 3's r0-r3  /
                //                               |
                //                               | (Group barrier)
                //                               v

                // Each work item sums __reduction_factor elements, and handles __radix_states_per_group radix states.
                // This enables full thread utilization in the first stage of reduction
                constexpr std::uint32_t __reduction_factor = decltype(__views)::__reduction_factor;
                const std::uint32_t __wis_per_radix_group = __wg_size / __reduction_factor;
                const std::uint32_t __radix_group = __self_lidx / __wis_per_radix_group;
                const std::uint32_t __radix_base = __radix_group * __radix_states_per_group;
                const std::uint32_t __wi_in_group = __self_lidx % __wis_per_radix_group;
                const std::uint32_t __col_base = __wi_in_group * __reduction_factor;

                // Accumulate __radix_states_per_group radix states from __reduction_factor columns
                _ONEDPL_PRAGMA_UNROLL
                for (std::uint32_t __c = 0; __c < __reduction_factor; ++__c)
                {
                    _ONEDPL_PRAGMA_UNROLL
                    for (std::uint32_t __r = 0; __r < __radix_states_per_group; ++__r)
                    {
                        __count_arr[__r] += static_cast<_CountT>(
                            __slm_buckets[__views.__get_bucket_idx(__radix_base + __r, __col_base + __c)]);
                    }
                }
                __dpl_sycl::__group_barrier(__self_item);

                // 3. REDUCTION PHASE: Write to SLM accessed as _CountT array (__slm_counts)
                //    The same N bytes of SLM are now viewed as (N / __packing_ratio) _CountT elements.
                //    Memory is reorganized by Radix state for a standard tree reduction.
                //
                //    SLM Address:    0 bytes                                                    N bytes
                //    _CountT View:   | Sums for Radix 0  | Sums for Radix 1  | ...              |
                //    Written by:     +--- WI partials ---+--- WI partials ---+ ...              |
                //
                //    (Tree reduction then processes these partial sums down to 1 total per Radix State)
                //    Active sums:    +-------------------+-------------------+ ...              |
                //                    +----------         +----------         + ...              |
                //                    +-----              +-----              + ...              |
                //                    +---                +---                + ...              |
                //                    +--                 +--                 + ...              |
                //                    +-                  +-                  + ...              |
                _ONEDPL_PRAGMA_UNROLL
                for (std::uint32_t __r = 0; __r < __radix_states_per_group; ++__r)
                {
                    __slm_counts[__views.__get_count_idx(__wg_size, __radix_base + __r, __wi_in_group)] =
                        __count_arr[__r];
                }
                __dpl_sycl::__group_barrier(__self_item);

                // Tree reduction: reduce partial sums down to 1 per radix state
                std::uint32_t __num_partial_sums = __wis_per_radix_group;
                for (std::uint32_t __stride = __num_partial_sums >> 1; __stride > 0; __stride >>= 1)
                {
                    // Each WI reduces its assigned radix states
                    if (__wi_in_group < __stride)
                    {
                        _ONEDPL_PRAGMA_UNROLL
                        for (std::uint32_t __r = 0; __r < __radix_states_per_group; ++__r)
                        {
                            __slm_counts[__views.__get_count_idx(__wg_size, __radix_base + __r, __wi_in_group)] +=
                                __slm_counts[__views.__get_count_idx(__wg_size, __radix_base + __r,
                                                                     __wi_in_group + __stride)];
                        }
                    }
                    __dpl_sycl::__group_barrier(__self_item);
                }

                // Write final count to global memory (only first 16 WIs, one per radix state)
                if (__self_lidx < __radix_states)
                {
                    __count_rng[(__segments + 1) * __self_lidx + __wgroup_idx] =
                        __slm_counts[__views.__get_count_idx(__wg_size, __self_lidx, 0)];
                }

                // Reset 'no operation flag' (indicates all keys are in one bin, skip reorder)
                if (__wgroup_idx == 0 && __self_lidx == 0)
                {
                    auto& __no_op_flag = __count_rng[__no_op_flag_idx];
                    __no_op_flag = 0;
                }
            });
    });

    return __count_levent;
}

//-----------------------------------------------------------------------
// radix sort: scan kernel (per iteration)
//-----------------------------------------------------------------------

template <typename _KernelName, std::uint32_t __radix_bits, typename _CountBuf>
sycl::event
__radix_sort_scan_submit(sycl::queue& __q, std::size_t __scan_wg_size, std::size_t __segments, _CountBuf& __count_buf,
                         std::size_t __n, sycl::event __dependency_event)
{
    using _CountT = typename _CountBuf::value_type;

    const std::size_t __no_op_flag_idx = __count_buf.size() - 1;
    auto __count_rng =
        oneapi::dpl::__ranges::all_view<_CountT, __par_backend_hetero::access_mode::read_write>(__count_buf);

    // Scan produces local offsets using count values.
    // There are no local offsets for the first segment, but the rest segments should be scanned
    // with respect to the count value in the first segment what requires n + 1 positions
    const std::size_t __scan_size = __segments + 1;
    __scan_wg_size = std::min(__scan_size, __scan_wg_size);

    const std::uint32_t __radix_states = 1 << __radix_bits;

    // compilation of the kernel prevents out of resources issue, which may occur due to usage of
    // collective algorithms such as joint_exclusive_scan even if local memory is not explicitly requested
    sycl::event __scan_event = __q.submit([&](sycl::handler& __hdl) {
        __hdl.depends_on(__dependency_event);
        // access the counters for all work groups
        oneapi::dpl::__ranges::__require_access(__hdl, __count_rng);
        __hdl.parallel_for<_KernelName>(
            sycl::nd_range<1>(__radix_states * __scan_wg_size, __scan_wg_size), [=](sycl::nd_item<1> __self_item) {
                // find borders of a region with a specific bucket id
                sycl::global_ptr<_CountT> __begin = __count_rng.begin() + __scan_size * __self_item.get_group(0);
                // TODO: consider another approach with use of local memory
                __dpl_sycl::__joint_exclusive_scan(__self_item.get_group(), __begin, __begin + __scan_size, __begin,
                                                   _CountT(0), __dpl_sycl::__plus<_CountT>{});
                const auto __wi = __self_item.get_local_linear_id();
                //That condition may be truth (by algo semantic) just on one WG, one WI, so there is no race here.
                if (__wi == __scan_wg_size - 1 && *(__begin + __scan_size - 1) == __n)
                {
                    auto& __no_op_flag = __count_rng[__no_op_flag_idx];
                    __no_op_flag = 1; //set flag if the all values got into one bin
                }
            });
    });
    return __scan_event;
}

//-----------------------------------------------------------------------
// radix sort: group level reorder algorithms
//-----------------------------------------------------------------------

template <typename _InputRange, typename _OutputRange>
void
__copy_kernel_for_radix_sort(sycl::nd_item<1> __self_item, const std::size_t __seg_start, std::size_t __seg_end,
                             const std::size_t __wg_size, _InputRange& __input, _OutputRange& __output)
{
    const std::size_t __self_lidx = __self_item.get_local_id(0);
    const std::uint16_t __residual = (__seg_end - __seg_start) % __wg_size;
    __seg_end -= __residual;

    for (std::size_t __val_idx = __seg_start + __self_lidx; __val_idx < __seg_end; __val_idx += __wg_size)
        __output[__val_idx] = std::move(__input[__val_idx]);

    if (__residual > 0 && __self_lidx < __residual)
    {
        const std::size_t __val_idx = __seg_end + __self_lidx;
        __output[__val_idx] = std::move(__input[__val_idx]);
    }
}
template <typename _ValueType>
_ValueType
__radix_sort_exclusive_scan(sycl::sub_group __sub_group, _ValueType __val, [[maybe_unused]] std::uint32_t __sg_local_id,
                            [[maybe_unused]] std::uint32_t __sg_size)
{
    // With icpx version 2024.2, exclusive_scan_over_group encounters some crashes which
    // only seem to affect CPU openCL targets, but this is not possible to detect at compile time
#if defined(__INTEL_LLVM_COMPILER) && __INTEL_LLVM_COMPILER >= 20240200 && __INTEL_LLVM_COMPILER < 20250000
    _ValueType __inclusive = __val;
    for (std::uint32_t __shift = 1; __shift < __sg_size; __shift <<= 1)
    {
        _ValueType __partial = sycl::shift_group_right(__sub_group, __inclusive, __shift);
        if (__sg_local_id >= __shift)
            __inclusive += __partial;
    }
    return __inclusive - __val;
#else
    return __dpl_sycl::__exclusive_scan_over_group(__sub_group, __val, __dpl_sycl::__plus<_ValueType>());
#endif
}

//-----------------------------------------------------------------------
// radix sort: reorder kernel helper
// Template helper function to reorder elements from input to output.
// Called with a single runtime branch to select the input/output ranges.
//-----------------------------------------------------------------------
template <std::uint32_t __radix_bits, bool __is_ascending, typename _OffsetT, typename _InputRange,
          typename _OutputRange, typename _OffsetRange, typename _Proj>
void
__radix_sort_reorder_impl(_InputRange& __input, _OutputRange& __output, _OffsetRange& __offset_rng,
                          std::uint16_t* __slm_counts, sycl::nd_item<1> __self_item, sycl::sub_group __sub_group,
                          _Proj __proj, std::uint32_t __radix_offset, std::size_t __segments, std::size_t __segment_idx,
                          std::size_t __wi_start, std::size_t __wi_end, std::uint32_t __sg_id,
                          std::uint32_t __sg_local_id, std::uint32_t __sg_size, std::uint32_t __num_subgroups)
{
    constexpr std::uint32_t __radix_states = 1 << __radix_bits;

    std::uint32_t __wi_prefix[__radix_states];
    {
        // Phase 1: Count pass - each work-item counts its contiguous elements
        std::uint32_t __local_counts[__radix_states] = {0};
        for (std::size_t __idx = __wi_start; __idx < __wi_end; ++__idx)
        {
            auto __val = __order_preserving_cast<__is_ascending>(std::invoke(__proj, __input[__idx]));
            ++__local_counts[__get_bucket<(1 << __radix_bits) - 1>(__val, __radix_offset)];
        }

        // Subgroup scan to get work-item prefix within subgroup
        // Last work-item writes totals directly to SLM (avoids broadcast)
        const bool __is_last_in_sg = (__sg_local_id == __sg_size - 1);
        for (std::uint32_t __b = 0; __b < __radix_states; ++__b)
        {
            __wi_prefix[__b] = __radix_sort_exclusive_scan(__sub_group, __local_counts[__b], __sg_local_id, __sg_size);
            if (__is_last_in_sg)
                __slm_counts[__sg_id * __radix_states + __b] = __wi_prefix[__b] + __local_counts[__b];
        }
    }

    __dpl_sycl::__group_barrier(__self_item);

    // Phase 2: Compute subgroup prefix (subgroups loop through radix states)
    // Reuses the same SLM region: reads totals, computes prefix, writes back in-place
    for (std::uint32_t __radix_state = __sg_id; __radix_state < __radix_states; __radix_state += __num_subgroups)
    {
        std::uint16_t __running_sum = 0;

        // Process counts in chunks of subgroup size
        for (std::uint32_t __base = 0; __base < __num_subgroups; __base += __sg_size)
        {
            const std::uint32_t __sg_idx = __base + __sg_local_id;

            // Load count (0 if out of bounds)
            std::uint16_t __val =
                (__sg_idx < __num_subgroups) ? __slm_counts[__sg_idx * __radix_states + __radix_state] : 0;

            // Exclusive scan within chunk
            std::uint16_t __local_prefix = __radix_sort_exclusive_scan(__sub_group, __val, __sg_local_id, __sg_size);

            // Add running sum from previous chunks
            std::uint16_t __prefix = __running_sum + __local_prefix;

            // Write prefix back to same location (safe: all reads complete before any writes)
            if (__sg_idx < __num_subgroups)
                __slm_counts[__sg_idx * __radix_states + __radix_state] = __prefix;

            // Update running sum: broadcast the last element's total
            std::uint16_t __chunk_total = __local_prefix + __val;
            __running_sum += __dpl_sycl::__group_broadcast(__sub_group, __chunk_total, __sg_size - 1);
        }
    }

    __dpl_sycl::__group_barrier(__self_item);

    // Phase 3: Compute final offsets = global_base + sg_prefix + wi_prefix
    _OffsetT __offsets[__radix_states];
    const std::size_t __scan_size = __segments + 1;
    _OffsetT __scanned_bin = 0;
    __offsets[0] = __offset_rng[__segment_idx] + __slm_counts[__sg_id * __radix_states] + __wi_prefix[0];

    for (std::uint32_t __b = 1; __b < __radix_states; ++__b)
    {
        __scanned_bin += __offset_rng[__b * __scan_size - 1];
        __offsets[__b] = __scanned_bin + __offset_rng[__segment_idx + __scan_size * __b] +
                         __slm_counts[__sg_id * __radix_states + __b] + __wi_prefix[__b];
    }

    // Phase 4: Scatter pass - re-read and write to output
    for (std::size_t __idx = __wi_start; __idx < __wi_end; ++__idx)
    {
        auto __in_val = __input[__idx];
        std::uint32_t __bucket = __get_bucket<(1 << __radix_bits) - 1>(
            __order_preserving_cast<__is_ascending>(std::invoke(__proj, __in_val)), __radix_offset);
        __output[__offsets[__bucket]++] = std::move(__in_val);
    }
}

//-----------------------------------------------------------------------
// radix sort: reorder kernel (per iteration)
//-----------------------------------------------------------------------
template <typename _KernelName, std::uint32_t __radix_bits, bool __is_ascending, typename _Range1, typename _Range2,
          typename _OffsetBuf, typename _Proj>
sycl::event
__radix_sort_reorder_submit(sycl::queue& __q, std::size_t __segments, std::size_t __wg_size, std::size_t __min_sg_size,
                            std::uint32_t __radix_offset, bool __input_is_first, _Range1&& __rng1, _Range2&& __rng2,
                            _OffsetBuf& __offset_buf, sycl::event __dependency_event, _Proj __proj)
{
    constexpr std::uint32_t __radix_states = 1 << __radix_bits;

    using _OffsetT = typename _OffsetBuf::value_type;

    assert(oneapi::dpl::__ranges::__size(__rng1) == oneapi::dpl::__ranges::__size(__rng2));

    // iteration space info
    const std::size_t __n = oneapi::dpl::__ranges::__size(__rng1);
    const std::size_t __elem_per_segment = oneapi::dpl::__internal::__dpl_ceiling_div(__n, __segments);
    const std::size_t __max_num_subgroups = oneapi::dpl::__internal::__dpl_ceiling_div(__wg_size, __min_sg_size);

    const std::size_t __no_op_flag_idx = __offset_buf.size() - 1;

    auto __offset_rng =
        oneapi::dpl::__ranges::all_view<_OffsetT, __par_backend_hetero::access_mode::read>(__offset_buf);

    // submit to reorder values
    sycl::event __reorder_event = __q.submit([&](sycl::handler& __hdl) {
        __hdl.depends_on(__dependency_event);
        // access the offsets for all work groups
        oneapi::dpl::__ranges::__require_access(__hdl, __offset_rng);
        // access the input and output data
        oneapi::dpl::__ranges::__require_access(__hdl, __rng1, __rng2);

        // Minimal SLM: only for subgroup coordination (no value buffering)
        // Single region reused: first stores subgroup totals, then overwritten with prefix sums
        auto __slm_counts = __dpl_sycl::__local_accessor<std::uint16_t>(__max_num_subgroups * __radix_states, __hdl);

        __hdl.parallel_for<_KernelName>(
            sycl::nd_range<1>(__segments * __wg_size, __wg_size), [=](sycl::nd_item<1> __self_item) {
                const std::size_t __segment_idx = __self_item.get_group(0);
                const std::size_t __seg_start = __elem_per_segment * __segment_idx;
                const std::size_t __seg_end = sycl::min(__seg_start + __elem_per_segment, __n);

                auto& __no_op_flag = __offset_rng[__no_op_flag_idx];
                if (__no_op_flag)
                {
                    // Single branch for copy: select input/output once, then copy without per-element branching
                    if (__input_is_first)
                        __copy_kernel_for_radix_sort(__self_item, __seg_start, __seg_end, __wg_size, __rng1, __rng2);
                    else
                        __copy_kernel_for_radix_sort(__self_item, __seg_start, __seg_end, __wg_size, __rng2, __rng1);
                    return;
                }

                auto __sub_group = __self_item.get_sub_group();
                const std::uint32_t __sg_id = __sub_group.get_group_linear_id();
                const std::uint32_t __sg_local_id = __sub_group.get_local_linear_id();
                const std::uint32_t __sg_size = __sub_group.get_local_range()[0];
                const std::uint32_t __num_subgroups = oneapi::dpl::__internal::__dpl_ceiling_div(__wg_size, __sg_size);

                // Compute this subgroup's contiguous chunk of the segment
                const std::size_t __elems_per_sg =
                    oneapi::dpl::__internal::__dpl_ceiling_div(__elem_per_segment, __num_subgroups);
                const std::size_t __sg_start = sycl::min(__seg_start + __sg_id * __elems_per_sg, __seg_end);
                const std::size_t __sg_end = sycl::min(__sg_start + __elems_per_sg, __seg_end);

                // Each work-item owns a contiguous block within its subgroup's chunk
                const std::size_t __sg_items = __sg_end - __sg_start;
                const std::size_t __items_per_wi = __sg_items / __sg_size;
                const std::size_t __wi_start = __sg_start + __sg_local_id * __items_per_wi;
                const std::size_t __wi_end =
                    (__sg_local_id == __sg_size - 1) ? __sg_end : (__wi_start + __items_per_wi);

                // Single branch to select input/output ranges, then reorder without per-element branching
                if (__n < std::numeric_limits<std::uint32_t>::max())
                {
                    using _LocalOffsetT = std::uint32_t;
                    if (__input_is_first)
                        __radix_sort_reorder_impl<__radix_bits, __is_ascending, _LocalOffsetT>(
                            __rng1, __rng2, __offset_rng, &__slm_counts[0], __self_item, __sub_group, __proj,
                            __radix_offset, __segments, __segment_idx, __wi_start, __wi_end, __sg_id, __sg_local_id,
                            __sg_size, __num_subgroups);
                    else
                        __radix_sort_reorder_impl<__radix_bits, __is_ascending, _LocalOffsetT>(
                            __rng2, __rng1, __offset_rng, &__slm_counts[0], __self_item, __sub_group, __proj,
                            __radix_offset, __segments, __segment_idx, __wi_start, __wi_end, __sg_id, __sg_local_id,
                            __sg_size, __num_subgroups);
                }
                else
                {
                    using _LocalOffsetT = std::uint64_t;
                    if (__input_is_first)
                        __radix_sort_reorder_impl<__radix_bits, __is_ascending, _LocalOffsetT>(
                            __rng1, __rng2, __offset_rng, &__slm_counts[0], __self_item, __sub_group, __proj,
                            __radix_offset, __segments, __segment_idx, __wi_start, __wi_end, __sg_id, __sg_local_id,
                            __sg_size, __num_subgroups);
                    else
                        __radix_sort_reorder_impl<__radix_bits, __is_ascending, _LocalOffsetT>(
                            __rng2, __rng1, __offset_rng, &__slm_counts[0], __self_item, __sub_group, __proj,
                            __radix_offset, __segments, __segment_idx, __wi_start, __wi_end, __sg_id, __sg_local_id,
                            __sg_size, __num_subgroups);
                }
            });
    });

    return __reorder_event;
}

//-----------------------------------------------------------------------
// radix sort: copy back kernel (when odd number of iterations)
//-----------------------------------------------------------------------

template <typename _KernelName, typename _Range1, typename _Range2>
sycl::event
__radix_sort_copy_back_submit(sycl::queue& __q, _Range1&& __in_rng, _Range2&& __out_rng, sycl::event __dependency_event)
{
    const std::size_t __n = oneapi::dpl::__ranges::__size(__in_rng);
    return __q.submit([&](sycl::handler& __hdl) {
        __hdl.depends_on(__dependency_event);
        oneapi::dpl::__ranges::__require_access(__hdl, __in_rng, __out_rng);
        __hdl.parallel_for<_KernelName>(sycl::range<1>(__n),
                                        [=](sycl::item<1> __item) { __in_rng[__item] = std::move(__out_rng[__item]); });
    });
}

//-----------------------------------------------------------------------
// radix sort: one iteration
//-----------------------------------------------------------------------

template <typename _CustomName, std::uint32_t __radix_bits, bool __is_ascending>
struct __parallel_multi_group_radix_sort
{
    template <typename... _Name>
    using __count_phase = __radix_sort_count_kernel<__radix_bits, __is_ascending, _Name...>;
    template <typename... _Name>
    using __local_scan_phase = __radix_sort_scan_kernel<__radix_bits, _Name...>;
    template <typename... _Name>
    using __reorder_phase = __radix_sort_reorder_kernel<__radix_bits, __is_ascending, _Name...>;
    template <typename... _Name>
    using __copy_back_phase = __radix_sort_copy_back_kernel<_Name...>;

    template <typename _InRange, typename _Proj>
    sycl::event
    operator()(sycl::queue& __q, _InRange&& __in_rng, _Proj __proj)
    {
        using _RadixCountKernel =
            __internal::__kernel_name_generator<__count_phase, _CustomName, std::decay_t<_InRange>, _Proj>;
        using _RadixLocalScanKernel = __internal::__kernel_name_generator<__local_scan_phase, _CustomName>;
        using _RadixReorderKernel =
            __internal::__kernel_name_generator<__reorder_phase, _CustomName, std::decay_t<_InRange>, _Proj>;
        using _RadixCopyBackKernel =
            __internal::__kernel_name_generator<__copy_back_phase, _CustomName, std::decay_t<_InRange>>;

        using _ValueT = oneapi::dpl::__internal::__value_t<_InRange>;
        using _KeyT = oneapi::dpl::__internal::__key_t<_Proj, _InRange>;

        constexpr std::uint32_t __radix_iters = __get_buckets_in_type<_KeyT>(__radix_bits);
        const std::uint32_t __radix_states = 1 << __radix_bits;
        const std::size_t __n = __in_rng.size();

        using _CounterType = std::uint64_t;
        std::size_t __wg_size_count = oneapi::dpl::__internal::__slm_adjusted_work_group_size(
            __q, sizeof(_CounterType) * __radix_states, std::size_t(128));
        // work-group size must be a power of 2 because of the tree reduction
        __wg_size_count =
            sycl::max(oneapi::dpl::__internal::__dpl_bit_floor(__wg_size_count), std::size_t(__radix_states));
        std::size_t __wg_size_scan = oneapi::dpl::__internal::__max_work_group_size(__q, 1024);
        std::size_t __wg_size_reorder = oneapi::dpl::__internal::__max_work_group_size(__q, 256);
        std::size_t __reorder_min_sg_size = oneapi::dpl::__internal::__min_sub_group_size(__q);

        constexpr std::uint32_t __unroll_elements = 8;
        // Keys per work-item in counting phase, recalculates based upon workgroup size for reorder phase.
        // We target a segment size of ~24KB to keep the scatter region in the reorder phase
        // within the L1 cache, preventing thrashing during uncoalesced memory accesses and scattered writes.
        constexpr std::size_t __target_segment_size_bytes = 24 * 1024;
        constexpr std::size_t __absolute_max_keys_per_wi = 255;
        constexpr std::size_t __absolute_min_keys_per_wi = __unroll_elements;
        constexpr std::size_t __target_minimum_segments = 128;

        static_assert(__absolute_max_keys_per_wi <= std::numeric_limits<unsigned char>::max(),
                      "Too large keys per work-item may cause overflow in counting phase");

        // target a number of segments
        std::size_t __keys_per_wi_count =
            oneapi::dpl::__internal::__dpl_ceiling_div(__n, (__wg_size_count * __target_minimum_segments));

        // cap total segment size
        __keys_per_wi_count =
            std::min(__keys_per_wi_count, __target_segment_size_bytes / (sizeof(_ValueT) * __wg_size_count));

        // apply hard limits
        __keys_per_wi_count = std::max(__keys_per_wi_count, __absolute_min_keys_per_wi);
        __keys_per_wi_count = std::min(__keys_per_wi_count, __absolute_max_keys_per_wi);

        const std::size_t __segments =
            oneapi::dpl::__internal::__dpl_ceiling_div(__n, __wg_size_count * __keys_per_wi_count);

        assert(__wg_size_count * __keys_per_wi_count < std::numeric_limits<std::uint16_t>::max() &&
               "Segment size too large for reorder implementation");

        // Additional __radix_states elements are used for getting local offsets from count values + no_op flag;
        // 'No operation' flag specifies whether to skip re-order phase if the all keys are the same (lie in one bin)
        const std::size_t __tmp_buf_size = __segments * __radix_states + __radix_states + 1 /*no_op flag*/;
        // memory for storing count and offset values
        sycl::buffer<_CounterType, 1> __tmp_buf{sycl::range<1>(__tmp_buf_size)};

        // memory for storing values sorted for an iteration
        oneapi::dpl::__par_backend_hetero::__buffer<_ValueT> __out_buffer_holder{__n};
        auto __out_rng = oneapi::dpl::__ranges::all_view<_ValueT, __par_backend_hetero::access_mode::read_write>(
            __out_buffer_holder.get_buffer());

        // iterations per each bucket
        // TODO: radix for bool can be made using 1 iteration (x2 speedup against current implementation)
        sycl::event __dependency_event;
        for (std::uint32_t __radix_iter = 0; __radix_iter < __radix_iters; ++__radix_iter)
        {
            // TODO: convert to ordered type once at the first iteration and convert back at the last one
            bool __input_is_first = (__radix_iter % 2 == 0);
            // Compute the radix position for the given iteration
            std::uint32_t __radix_offset = __radix_iter * __radix_bits;

            // 1. Count Phase
            __dependency_event =
                __radix_sort_count_submit<_RadixCountKernel, __radix_bits, __is_ascending, __unroll_elements>(
                    __q, __segments, __wg_size_count, __radix_offset, __input_is_first, __in_rng, __out_rng, __tmp_buf,
                    __dependency_event, __proj);

            // 2. Scan Phase
            std::size_t __scan_size =
                __input_is_first ? oneapi::dpl::__ranges::__size(__in_rng) : oneapi::dpl::__ranges::__size(__out_rng);
            __dependency_event = __radix_sort_scan_submit<_RadixLocalScanKernel, __radix_bits>(
                __q, __wg_size_scan, __segments, __tmp_buf, __scan_size, __dependency_event);

            // 3. Reorder Phase
            __dependency_event = __radix_sort_reorder_submit<_RadixReorderKernel, __radix_bits, __is_ascending>(
                __q, __segments, __wg_size_reorder, __reorder_min_sg_size, __radix_offset, __input_is_first, __in_rng,
                __out_rng, __tmp_buf, __dependency_event, __proj);
        }

        // If odd number of iterations, the result is in __out_rng; copy back to __in_rng
        if constexpr (__radix_iters % 2 != 0)
        {
            __dependency_event =
                __radix_sort_copy_back_submit<_RadixCopyBackKernel>(__q, __in_rng, __out_rng, __dependency_event);
        }

        return __dependency_event;
    }
}; // struct __parallel_multi_group_radix_sort

// sorting by just one work group
#include "parallel_backend_sycl_radix_sort_one_wg.h"

//-----------------------------------------------------------------------
// radix sort: main function
//-----------------------------------------------------------------------
template <bool __is_ascending, typename _Range, typename _ExecutionPolicy, typename _Proj>
__future<sycl::event>
__parallel_radix_sort(oneapi::dpl::__internal::__device_backend_tag, _ExecutionPolicy&& __exec, _Range&& __in_rng,
                      _Proj __proj)
{
    const std::size_t __n = oneapi::dpl::__ranges::__size(__in_rng);
    assert(__n > 1);

    // radix bits represent number of processed bits in each value during one iteration
    constexpr std::uint32_t __radix_bits = 4;

    sycl::event __event;

    sycl::queue __q_local = __exec.queue();

    // Limit the work-group size to prevent large sizes on CPUs. Empirically found value.
    // This value exceeds the current practical limit for GPUs, but may need to be re-evaluated in the future.
    const std::size_t __max_wg_size = oneapi::dpl::__internal::__max_work_group_size(__q_local, (std::size_t)4096);
    const auto __subgroup_sizes = __q_local.get_device().template get_info<sycl::info::device::sub_group_sizes>();
    const bool __dev_has_sg16 = std::find(__subgroup_sizes.begin(), __subgroup_sizes.end(),
                                          static_cast<std::size_t>(16)) != __subgroup_sizes.end();

    using _RadixSortKernel = oneapi::dpl::__internal::__policy_kernel_name<_ExecutionPolicy>;

    // Select block size based on input size (block_size = elements per work-item)
    // Larger block sizes reduce register spills but require more registers per work-item
    if (__n <= std::min<std::size_t>(1024, __max_wg_size * 4))
        __event = __subgroup_radix_sort<_RadixSortKernel, 4, __radix_bits, __is_ascending>{}(
            __q_local, std::forward<_Range>(__in_rng), __proj, __max_wg_size);
    else if (__n <= std::min<std::size_t>(2048, __max_wg_size * 8))
        __event = __subgroup_radix_sort<_RadixSortKernel, 8, __radix_bits, __is_ascending>{}(
            __q_local, std::forward<_Range>(__in_rng), __proj, __max_wg_size);
    else if (__n <= std::min<std::size_t>(4096, __max_wg_size * 16))
        __event = __subgroup_radix_sort<_RadixSortKernel, 16, __radix_bits, __is_ascending>{}(
            __q_local, std::forward<_Range>(__in_rng), __proj, __max_wg_size);
    // In __subgroup_radix_sort, we request a sub-group size of 16 via _ONEDPL_SYCL_REQD_SUB_GROUP_SIZE_IF_SUPPORTED
    // for compilation targets that support this option. For the below cases, register spills that result in
    // runtime exceptions have been observed on accelerators that do not support the requested sub-group size of 16.
    // For the above cases that request but may not receive a sub-group size of 16, inputs are small enough to avoid
    // register spills on assessed hardware.
    else if (__n <= std::min<std::size_t>(16384, __max_wg_size * 32) && __dev_has_sg16)
        __event = __subgroup_radix_sort<_RadixSortKernel, 32, __radix_bits, __is_ascending>{}(
            __q_local, std::forward<_Range>(__in_rng), __proj, __max_wg_size);
    else
    {
        __event = __parallel_multi_group_radix_sort<_RadixSortKernel, __radix_bits, __is_ascending>{}(
            __q_local, std::forward<_Range>(__in_rng), __proj);
    }

    return __future{std::move(__event)};
}

} // namespace __par_backend_hetero
} // namespace dpl
} // namespace oneapi

#endif // _ONEDPL_PARALLEL_BACKEND_SYCL_RADIX_SORT_H
