#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
 * Copyright 2020-2026 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Portions of this file are derived from PyTorch
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
 * Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
 * Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
 * Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
 * Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
 * Copyright (c) 2011-2013 NYU                      (Clement Farabet)
 * Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon
 * Bottou, Iain Melvin, Jason Weston) Copyright (c) 2006      Idiap Research
 * Institute (Samy Bengio) Copyright (c) 2001-2004 Idiap Research Institute
 * (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
 *
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC
 * Laboratories America and IDIAP Research Institute nor the names of its
 * contributors may be used to endorse or promote products derived from this
 * software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 *
 * pytroch's implementation using radix-selection
 * - from pytorch/aten/src/THC/THCTensorTopK.cuh
 * https://github.com/pytorch/pytorch/blob/d792c21f72b6bcace80face2a4157528d977bfa9/aten/src/THC/THCTensorTopK.cuh
 * - warning: ignore some types, for eg., does not include THCNumerics for
 * T==half
 */

#pragma once
#include <ATen/ceil_div.h>
#include <ATen/native/xpu/sycl/Atomics.h>
#include <comm/SYCLContext.h>
#include <comm/SYCLHelpers.h>

namespace at::native::xpu {

template <typename scalar_t>
struct TopKTypeConfig {};

template <>
struct TopKTypeConfig<float> {
  using RadixType = uint32_t;

  // Converts a float to an integer representation with the same
  // sorting; i.e., for floats f1, f2:
  // if f1 < f2 then convert(f1) < convert(f2)
  // We use this to enable radix selection of floating-point values.
  // This also gives a relative order for NaNs, but that's ok, as they
  // will all be adjacent
  // neg inf: signbit=1 exp=ff fraction=0 --> radix = 0 00 ff..
  // pos inf: signbit=0 exp=ff fraction=0 --> radix = 1 ff 00..
  // pos nan: signbit=0 exp=ff fraction>0 --> radix = 1 ff x>0
  // neg nan: signbit=1 exp=ff fraction>0 --> radix = 0 00 x<ff...
  static inline RadixType convert(float v) {
    RadixType x = *((uint32_t*)&v);
    RadixType mask = (x & 0x80000000) ? 0xffffffff : 0x80000000;
    return (v == v) ? (x ^ mask) : 0xffffffff;
  }

  static inline float deconvert(RadixType v) {
    RadixType mask = (v & 0x80000000) ? 0x80000000 : 0xffffffff;
    return __int_as_float(v ^ mask);
  }
};

template <>
struct TopKTypeConfig<uint8_t> {
  using RadixType = uint32_t;

  static inline RadixType convert(uint8_t v) {
    return v;
  }

  static inline uint8_t deconvert(RadixType v) {
    return v;
  }
};

template <>
struct TopKTypeConfig<int8_t> {
  using RadixType = uint32_t;

  static inline RadixType convert(int8_t v) {
    return 128u + v;
  }

  static inline int8_t deconvert(RadixType v) {
    return v - 128;
  }
};

template <>
struct TopKTypeConfig<int16_t> {
  using RadixType = uint32_t;

  static inline RadixType convert(int16_t v) {
    return 32768u + v;
  }

  static inline int16_t deconvert(RadixType v) {
    return v - 32768;
  }
};

template <>
struct TopKTypeConfig<int32_t> {
  using RadixType = uint32_t;

  static inline RadixType convert(int32_t v) {
    return 2147483648u + v;
  }

  static inline int32_t deconvert(RadixType v) {
    return v - 2147483648u;
  }
};

template <>
struct TopKTypeConfig<int64_t> {
  using RadixType = uint64_t;

  static inline RadixType convert(int64_t v) {
    return 9223372036854775808ull + v;
  }

  static inline int64_t deconvert(RadixType v) {
    return v - 9223372036854775808ull;
  }
};

template <>
struct TopKTypeConfig<double> {
  using RadixType = uint64_t;

  static inline RadixType convert(double v) {
    RadixType x = *((uint64_t*)&v);
    RadixType mask = -((x >> 63)) | 0x8000000000000000;
    return (v == v) ? (x ^ mask) : 0xffffffffffffffff;
  }

  static inline double deconvert(RadixType v) {
    RadixType mask = ((v >> 63) - 1) | 0x8000000000000000;
    return __long_long_as_double(v ^ mask);
  }
};

template <>
struct TopKTypeConfig<at::Half> {
  using RadixType = uint32_t;

  static inline RadixType convert(at::Half v) {
    RadixType x = *((uint16_t*)&v);
    RadixType mask = (x & 0x00008000) ? 0x0000ffff : 0x00008000;
    return (v == v) ? (x ^ mask) : 0xffff;
  }

  static inline at::Half deconvert(RadixType v) {
    RadixType mask = (v & 0x00008000) ? 0x00008000 : 0x0000ffff;
    return __ushort_as_half(v ^ mask);
  }
};

template <>
struct TopKTypeConfig<at::BFloat16> {
  using RadixType = uint32_t;

  static inline RadixType convert(at::BFloat16 v) {
    RadixType x = v.x;
    RadixType mask = (x & 0x00008000) ? 0x0000ffff : 0x00008000;
    return (v == v) ? (x ^ mask) : 0xffff;
  }

  static inline at::BFloat16 deconvert(RadixType v) {
    RadixType mask = (v & 0x00008000) ? 0x00008000 : 0x0000ffff;
    at::BFloat16 r;
    r.x = (v ^ mask);
    return r;
  }
};

template <typename T>
struct Bitfield {};

template <>
struct Bitfield<unsigned int> {
  static inline unsigned int getBitfield(unsigned int val, int pos, int len) {
    pos &= 0xff;
    len &= 0xff;
    unsigned int m = (1u << len) - 1u;
    return (val >> pos) & m;
  }

  static inline unsigned int setBitfield(
      unsigned int val,
      unsigned int toInsert,
      int pos,
      int len) {
    pos &= 0xff;
    len &= 0xff;
    unsigned int m = (1u << len) - 1u;
    toInsert &= m;
    toInsert <<= pos;
    m <<= pos;
    return (val & ~m) | toInsert;
  }
};

template <>
struct Bitfield<uint64_t> {
  static inline uint64_t getBitfield(uint64_t val, int pos, int len) {
    pos &= 0xff;
    len &= 0xff;

    uint64_t m = (1u << len) - 1u;
    return (val >> pos) & m;
  }

  static inline uint64_t setBitfield(
      uint64_t val,
      uint64_t toInsert,
      int pos,
      int len) {
    pos &= 0xff;
    len &= 0xff;

    uint64_t m = (1u << len) - 1u;
    toInsert &= m;
    toInsert <<= pos;
    m <<= pos;
    return (val & ~m) | toInsert;
  }
};

// This function counts the distribution of all input values in a
// slice we are selecting by radix digit at `radixDigitPos`, but only
// those that pass the filter `((v & desiredMask) == desired)`.
// This produces and broadcasts the seen counts for a single block only.
// `smem` must have at least `RadixSize` elements.
template <
    typename scalar_t,
    typename bitwise_t,
    typename index_t,
    typename CountType,
    int RadixSize,
    int RadixBits>
void countRadixUsingMask(
    CountType counts[RadixSize],
    const sycl_local_acc_t<int>& smem,
    bitwise_t desired,
    bitwise_t desiredMask,
    int radixDigitPos,
    index_t sliceSize,
    index_t withinSliceStride,
    const scalar_t* data,
    sycl::nd_item<1>& item_id) {
  // Clear out per-thread counts from a previous round
  for (int i = 0; i < RadixSize; ++i) {
    counts[i] = 0;
  }

  auto local_id = item_id.get_local_id(0);
  if (local_id < RadixSize) {
    smem[local_id] = 0;
  }

  sycl::group_barrier(item_id.get_group());
  // Scan over all the data. Upon a read, the warp will accumulate
  // counts per each digit in the radix using warp voting.
  for (index_t i = local_id; i < sliceSize; i += item_id.get_local_range(0)) {
    bitwise_t val =
        TopKTypeConfig<scalar_t>::convert(data[i * withinSliceStride]);

    bool hasVal = ((val & desiredMask) == desired);
    bitwise_t digitInRadix =
        Bitfield<bitwise_t>::getBitfield(val, radixDigitPos, RadixBits);
    if (hasVal)
      counts[digitInRadix]++;
  }

  auto smem_ptr = (sycl_local_ptr<int>)(smem.template get_multi_ptr<
                                                sycl::access::decorated::no>()
                                            .get());
  for (uint32_t i = 0; i < RadixSize; ++i) {
    atomicAdd(smem_ptr + i, counts[i]);
  }

  sycl::group_barrier(item_id.get_group());

  // For each thread, read in the total counts
  for (uint32_t i = 0; i < RadixSize; ++i) {
    counts[i] = smem[i];
  }

  sycl::group_barrier(item_id.get_group());
}

// Over what radix we are selecting values
constexpr int RADIX_BITS = 2; // digits are base-(2 ^ RADIX_BITS)
constexpr int RADIX_SIZE = 4; // 2 ^ RADIX_BITS
constexpr int RADIX_MASK = (RADIX_SIZE - 1);

// This finds the unique value `v` that matches the pattern
// ((v & desired) == desiredMask) in our sorted int format
template <typename scalar_t, typename bitwise_t, typename index_t>
scalar_t findPattern(
    const sycl_local_acc_t<int>& smem,
    const scalar_t* data,
    index_t sliceSize,
    index_t withinSliceStride,
    bitwise_t desired,
    bitwise_t desiredMask,
    sycl::nd_item<1>& item_id) {
  auto local_id = item_id.get_local_id(0);
  auto smem_ptr = static_cast<scalar_t*>(static_cast<void*>(
      smem.template get_multi_ptr<sycl::access::decorated::no>().get()));
  if (local_id < RADIX_SIZE) {
    smem_ptr[RADIX_SIZE] = static_cast<scalar_t>(0);
  }

  sycl::group_barrier(item_id.get_group());

  // All threads participate in the loop, in order to sync on the flag
  index_t numIterations =
      round_up(sliceSize, (index_t)item_id.get_local_range(0));
  for (index_t i = local_id; i < numIterations;
       i += item_id.get_local_range(0)) {
    bool inRange = (i < sliceSize);
    scalar_t v =
        inRange ? data[i * withinSliceStride] : static_cast<scalar_t>(0);

    if (inRange &&
        ((TopKTypeConfig<scalar_t>::convert(v) & desiredMask) == desired)) {
      // There should not be conflicts if we are using findPattern,
      // since the result is unique
      smem_ptr[0] = static_cast<scalar_t>(1);
      smem_ptr[1] = v; // can't use val as the flag, since it could be 0
    }

    sycl::group_barrier(item_id.get_group());

    scalar_t found = smem_ptr[0];
    scalar_t val = smem_ptr[1];

    sycl::group_barrier(item_id.get_group());

    // Check to see if a thread found the value
    if (found != static_cast<scalar_t>(0)) {
      // all threads return this value
      return val;
    }
  }

  // should not get here
  return static_cast<scalar_t>(0);
}

// Returns the top-Kth element found in the data using radix selection
template <typename scalar_t, typename bitwise_t, typename index_t, bool Order>
void radixSelect(
    const scalar_t* data,
    index_t k,
    index_t sliceSize,
    index_t withinSliceStride,
    const sycl_local_acc_t<int>& smem,
    scalar_t* topK,
    sycl::nd_item<1>& item_id) {
  // Per-thread buckets into which we accumulate digit counts in our
  // radix
  int counts[RADIX_SIZE];

  // We only consider elements x such that (x & desiredMask) == desired
  // Initially, we consider all elements of the array, so the above
  // statement is true regardless of input.
  bitwise_t desired = 0;
  bitwise_t desiredMask = 0;

  // We are looking for the top kToFind-th element when iterating over
  // digits; this count gets reduced by elimination when counting
  // successive digits
  int kToFind = k;

  // We start at the most significant digit in our radix, scanning
  // through to the least significant digit
  for (int digitPos = sizeof(scalar_t) * 8 - RADIX_BITS; digitPos >= 0;
       digitPos -= RADIX_BITS) {
    // Count radix distribution for the current position and reduce
    // across all threads
    countRadixUsingMask<
        scalar_t,
        bitwise_t,
        index_t,
        int,
        RADIX_SIZE,
        RADIX_BITS>(
        counts,
        smem,
        desired,
        desiredMask,
        digitPos,
        sliceSize,
        withinSliceStride,
        data,
        item_id);

    auto found_unique = [&](int i, int count) -> bool {
      /* All threads have the same value in counts here, so all */
      /* threads will return from the function. */
      if (count == 1 && kToFind == 1) {
        /* There is a unique answer. */
        desired =
            Bitfield<bitwise_t>::setBitfield(desired, i, digitPos, RADIX_BITS);
        desiredMask = Bitfield<bitwise_t>::setBitfield(
            desiredMask, RADIX_MASK, digitPos, RADIX_BITS);

        /* The answer is now the unique element v such that: */
        /* (v & desiredMask) == desired */
        /* However, we do not yet know what the actual element is. We */
        /* need to perform a search through the data to find the */
        /* element that matches this pattern. */
        *topK = findPattern<scalar_t, bitwise_t, index_t>(
            smem,
            data,
            sliceSize,
            withinSliceStride,
            desired,
            desiredMask,
            item_id);
        return true;
      }
      return false;
    };
    auto found_non_unique = [&](int i, int count) -> bool {
      if (count >= kToFind) {
        desired =
            Bitfield<bitwise_t>::setBitfield(desired, i, digitPos, RADIX_BITS);
        desiredMask = Bitfield<bitwise_t>::setBitfield(
            desiredMask, RADIX_MASK, digitPos, RADIX_BITS);

        /* The top-Kth element v must now be one such that: */
        /* (v & desiredMask == desired) */
        /* but we haven't narrowed it down; we must check the next */
        /* least-significant digit */
        return true;
      }
      kToFind -= count;
      return false; // continue the loop
    };

    // All threads participate in the comparisons below to know the
    // final result
    if (Order) {
      // Process in descending order
      for (int i = RADIX_SIZE - 1; i >= 0; --i) {
        int count = counts[i];
        if (found_unique(i, count)) {
          return;
        }
        if (found_non_unique(i, count)) {
          break;
        }
      }
    } else {
      // Process in ascending order
      for (int i = 0; i < RADIX_SIZE; ++i) {
        int count = counts[i];
        if (found_unique(i, count)) {
          return;
        }
        if (found_non_unique(i, count)) {
          break;
        }
      }
    }
  } // end digitPos for

  // There is no unique result, but there is a non-unique result
  // matching `desired` exactly
  *topK = TopKTypeConfig<scalar_t>::deconvert(desired);
}

} // namespace at::native::xpu
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
