// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// 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
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef PADDLE_PHI_KERNELS_FUNCS_TOP_K_CUDA_KERNEL_H_
#define PADDLE_PHI_KERNELS_FUNCS_TOP_K_CUDA_KERNEL_H_

// GPU TopK kernel implementation using radix-select and multi-tier sorting.

#include <algorithm>
#include <cstdint>
#include <limits>

// Include top_k_function_cuda.h to get CUB NumericTraits for float16/bfloat16.
// This header includes cub/cub.cuh and defines the required traits.
#include "paddle/phi/kernels/funcs/top_k_function_cuda.h"

#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_info.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"

#include "paddle/phi/kernels/argsort_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/take_along_axis_kernel.h"

// ============================================================================
// Helper definitions
// All helpers are placed in an anonymous namespace to avoid ODR conflicts
// with Paddle's existing implementations.
// ============================================================================

namespace topk_detail {

// Stream type alias: gpuStream_t is in phi:: namespace, bring it into scope
using phi::gpuStream_t;

// --- Constants ---
constexpr int MAX_TENSORINFO_DIMS = 25;
constexpr int64_t MAX_GRID_SIZE = 65535LL;

// --- ceil_div and round_up ---
template <typename T>
__host__ __device__ __forceinline__ T topk_ceil_div(T a, T b) {
  return (a + b - 1) / b;
}

template <typename T>
__host__ __device__ __forceinline__ T topk_round_up(T a, T b) {
  return topk_ceil_div(a, b) * b;
}

// --- getGridFromTiles ---
inline bool getGridFromTiles(int64_t gridTiles, dim3* grid) {
  if (gridTiles > MAX_GRID_SIZE * MAX_GRID_SIZE * MAX_GRID_SIZE) {
    return false;
  }
  int64_t gridX = gridTiles > MAX_GRID_SIZE ? MAX_GRID_SIZE : gridTiles;
  int64_t gridY = 1;
  int64_t gridZ = 1;
  if (gridTiles > MAX_GRID_SIZE) {
    gridTiles = topk_ceil_div(gridTiles, (int64_t)MAX_GRID_SIZE);
    gridY = gridTiles > MAX_GRID_SIZE ? MAX_GRID_SIZE : gridTiles;
    if (gridTiles > MAX_GRID_SIZE) {
      gridTiles = topk_ceil_div(gridTiles, (int64_t)MAX_GRID_SIZE);
      gridZ = gridTiles > MAX_GRID_SIZE ? MAX_GRID_SIZE : gridTiles;
    }
  }
  *grid = dim3(gridX, gridY, gridZ);
  return true;
}

// --- getLinearBlockId ---
template <typename index_t>
__device__ __forceinline__ index_t getLinearBlockId() {
  return static_cast<index_t>(blockIdx.z) * gridDim.y * gridDim.x +
         static_cast<index_t>(blockIdx.y) * gridDim.x + blockIdx.x;
}

// --- doLdg ---
// Generic fallback for custom types (phi::float16, phi::bfloat16, etc.)
template <typename T>
__device__ __forceinline__ T doLdg(const T* p) {
  return *p;
}

// Specializations for built-in types that support __ldg
#if !defined(__HIPCC__)
template <>
__device__ __forceinline__ float doLdg(const float* p) {
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
template <>
__device__ __forceinline__ double doLdg(const double* p) {
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
template <>
__device__ __forceinline__ int doLdg(const int* p) {
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
template <>
__device__ __forceinline__ unsigned int doLdg(const unsigned int* p) {
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
template <>
__device__ __forceinline__ long long doLdg(  // NOLINT
    const long long* p) {                    // NOLINT
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
template <>
__device__ __forceinline__ unsigned long long doLdg(  // NOLINT
    const unsigned long long* p) {                    // NOLINT
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
template <>
__device__ __forceinline__ int16_t doLdg(const int16_t* p) {
#if __CUDA_ARCH__ >= 350
  return __ldg(p);
#else
  return *p;
#endif
}
#endif  // !__HIPCC__

// --- Bitfield ---
template <typename T>
struct Bitfield {};

template <>
struct Bitfield<unsigned int> {
  static __device__ __forceinline__ unsigned int getBitfield(unsigned int val,
                                                             int pos,
                                                             int len) {
    unsigned int ret;
#if defined(__HIPCC__)
    ret = (val >> pos) & ((1u << len) - 1u);
#else
    asm("bfe.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(val), "r"(pos), "r"(len));
#endif
    return ret;
  }

  static __device__ __forceinline__ unsigned int setBitfield(
      unsigned int val, unsigned int to_insert, int pos, int len) {
    unsigned int ret;
#if defined(__HIPCC__)
    unsigned int mask = ((1u << len) - 1u) << pos;
    ret = (val & ~mask) | ((to_insert << pos) & mask);
#else
    asm("bfi.b32 %0, %1, %2, %3, %4;"
        : "=r"(ret)
        : "r"(to_insert), "r"(val), "r"(pos), "r"(len));
#endif
    return ret;
  }
};

template <>
struct Bitfield<uint64_t> {
  static __device__ __forceinline__ uint64_t getBitfield(uint64_t val,
                                                         int pos,
                                                         int len) {
    uint64_t ret;
#if defined(__HIPCC__)
    ret = (val >> pos) & ((1ULL << len) - 1ULL);
#else
    asm("bfe.u64 %0, %1, %2, %3;" : "=l"(ret) : "l"(val), "r"(pos), "r"(len));
#endif
    return ret;
  }

  static __device__ __forceinline__ uint64_t setBitfield(uint64_t val,
                                                         uint64_t to_insert,
                                                         int pos,
                                                         int len) {
    uint64_t ret;
#if defined(__HIPCC__)
    uint64_t mask = ((1ULL << len) - 1ULL) << pos;
    ret = (val & ~mask) | ((to_insert << pos) & mask);
#else
    asm("bfi.b64 %0, %1, %2, %3, %4;"
        : "=l"(ret)
        : "l"(to_insert), "l"(val), "r"(pos), "r"(len));
#endif
    return ret;
  }
};

// --- getLaneId / getLaneMaskLe ---
__device__ __forceinline__ int getLaneId() {
#if defined(__HIPCC__)
  return __lane_id();
#else
  int laneId;
  asm("mov.s32 %0, %%laneid;" : "=r"(laneId));
  return laneId;
#endif
}

__device__ __forceinline__ unsigned getLaneMaskLe() {
#if defined(__HIPCC__)
  // HIP warp size is 64, construct mask for lanes <= current lane
  return (getLaneId() == 63) ? 0xFFFFFFFFFFFFFFFFULL
                             : (1ULL << (getLaneId() + 1)) - 1ULL;
#else
  unsigned mask;
  asm("mov.u32 %0, %%lanemask_le;" : "=r"(mask));
  return mask;
#endif
}

__device__ __forceinline__ unsigned getLaneMaskLt() {
#if defined(__HIPCC__)
  return (getLaneId() == 0) ? 0ULL : (1ULL << getLaneId()) - 1ULL;
#else
  unsigned mask;
  asm("mov.u32 %0, %%lanemask_lt;" : "=r"(mask));
  return mask;
#endif
}

// --- WARP macros ---
#ifdef __HIPCC__
#define TOPK_WARP_SIZE 64
#define TOPK_WARP_BALLOT(PREDICATE) __ballot((PREDICATE))
#define TOPK_WARP_BALLOT_MASK(PREDICATE, MASK) __ballot((PREDICATE))
#define TOPK_WARP_SHFL_DOWN(VAL, DELTA) \
  __shfl_down((VAL), static_cast<unsigned int>(DELTA))
#else
#define TOPK_WARP_SIZE 32
#define TOPK_WARP_BALLOT(PREDICATE) __ballot_sync(0xffffffff, (PREDICATE))
#define TOPK_WARP_BALLOT_MASK(PREDICATE, MASK) \
  __ballot_sync((MASK), (PREDICATE))
#define TOPK_WARP_SHFL_DOWN(VAL, DELTA) \
  __shfl_down_sync(0xffffffff, (VAL), static_cast<unsigned int>(DELTA))
#endif

// --- TopKTypeConfig ---
template <typename T>
struct TopKTypeConfig {};

template <>
struct TopKTypeConfig<float> {
  typedef uint32_t RadixType;

  static inline __device__ RadixType convert(float v) {
    RadixType x = __float_as_int(v);
    RadixType mask = (x & 0x80000000) ? 0xffffffff : 0x80000000;
    return (v == v) ? (x ^ mask) : 0xffffffff;
  }

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

template <>
struct TopKTypeConfig<double> {
  typedef uint64_t RadixType;

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

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

template <>
struct TopKTypeConfig<int32_t> {
  typedef uint32_t RadixType;

  static inline __device__ RadixType convert(int32_t v) {
    static_assert(sizeof(int) == 4, "");
    return 2147483648u + v;
  }

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

template <>
struct TopKTypeConfig<int64_t> {
  typedef uint64_t RadixType;

  static inline __device__ RadixType convert(int64_t v) {
    static_assert(sizeof(int64_t) == 8, "");
    return 9223372036854775808ull + v;
  }

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

template <>
struct TopKTypeConfig<phi::dtype::float16> {
  typedef uint32_t RadixType;

  static inline __device__ RadixType convert(phi::dtype::float16 v) {
    RadixType x = __half_as_ushort(v.to_half());
    RadixType mask = (x & 0x00008000) ? 0x0000ffff : 0x00008000;
    half v_h = v.to_half();
    return (v_h == v_h) ? (x ^ mask) : 0xffff;
  }

  static inline __device__ phi::dtype::float16 deconvert(RadixType v) {
    RadixType mask = (v & 0x00008000) ? 0x00008000 : 0x0000ffff;
    return static_cast<phi::dtype::float16>(__ushort_as_half(v ^ mask));
  }
};

template <>
struct TopKTypeConfig<phi::dtype::bfloat16> {
  typedef uint32_t RadixType;

  static inline __device__ RadixType convert(phi::dtype::bfloat16 v) {
    RadixType x = v.x;
    RadixType mask = (x & 0x00008000) ? 0x0000ffff : 0x00008000;
    return (v == v) ? (x ^ mask) : 0xffff;
  }

  static inline __device__ phi::dtype::bfloat16 deconvert(RadixType v) {
    RadixType mask = (v & 0x00008000) ? 0x00008000 : 0x0000ffff;
    phi::dtype::bfloat16 r;
    r.x = (v ^ mask);
    return r;
  }
};

// uint8_t is needed by the radix select
template <>
struct TopKTypeConfig<uint8_t> {
  typedef uint32_t RadixType;

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

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

// --- TensorInfo ---
template <typename T, typename IndexType>
struct TensorInfo {
  T* data;
  IndexType sizes[MAX_TENSORINFO_DIMS];
  IndexType strides[MAX_TENSORINFO_DIMS];
  int dims;

  // collapse_dims: merges contiguous dimensions for efficient indexing
  // See note on [collapse dims].
  int collapseDims(const int excludeDim = -1) {
    int stopDim = (excludeDim == -1) ? dims : excludeDim;
    int newIndex = -1;
    int oldIndex = 0;
    int remappedExcludedDim = -1;

    while (oldIndex < dims) {
      // Finds a dimension to collapse into
      for (; oldIndex < stopDim; ++oldIndex) {
        if (sizes[oldIndex] == 1) {
          continue;
        }

        ++newIndex;
        sizes[newIndex] = sizes[oldIndex];
        strides[newIndex] = strides[oldIndex];
        ++oldIndex;
        break;
      }

      // Collapses dims
      for (; oldIndex < stopDim; ++oldIndex) {
        if (sizes[oldIndex] == 1) {
          continue;
        }

        if (strides[newIndex] == sizes[oldIndex] * strides[oldIndex]) {
          sizes[newIndex] *= sizes[oldIndex];
          strides[newIndex] = strides[oldIndex];
        } else {
          ++newIndex;
          sizes[newIndex] = sizes[oldIndex];
          strides[newIndex] = strides[oldIndex];
        }
      }

      // Handles excludeDim being set (oldIndex == excludeDim)
      if (oldIndex != dims) {
        // Preserves excluded dimension
        ++newIndex;
        sizes[newIndex] = sizes[oldIndex];
        strides[newIndex] = strides[oldIndex];
        remappedExcludedDim = newIndex;

        // Restarts iteration after excludeDim
        ++oldIndex;
        stopDim = dims;
      }
    }

    // Handles special case of all dims size 1
    if (newIndex == -1 || (newIndex == 0 && sizes[0] == 1)) {
      dims = 1;
      sizes[0] = 1;
      strides[0] = 1;
      return 0;
    }

    dims = newIndex + 1;
    return remappedExcludedDim;
  }
};

// --- IndexToOffset ---
template <typename T, typename IndexType, int Dims>
struct IndexToOffset {
  static __host__ __device__ IndexType
  get(IndexType linearId, const TensorInfo<T, IndexType>& info) {
    IndexType offset = 0;
    for (int i = Dims - 1; i > 0; --i) {
      IndexType curDimIndex = linearId % info.sizes[i];
      offset += curDimIndex * info.strides[i];
      linearId /= info.sizes[i];
    }
    return offset + linearId * info.strides[0];
  }
};

// Specialization for Dim == -1 (runtime dims)
template <typename T, typename IndexType>
struct IndexToOffset<T, IndexType, -1> {
  static __host__ __device__ IndexType
  get(IndexType linearId, const TensorInfo<T, IndexType>& info) {
    IndexType offset = 0;
    for (int i = info.dims - 1; i > 0; --i) {
      IndexType curDimIndex = linearId % info.sizes[i];
      offset += curDimIndex * info.strides[i];
      linearId /= info.sizes[i];
    }
    return offset + linearId * info.strides[0];
  }
};

// Specialization for Dim == 1
template <typename T, typename IndexType>
struct IndexToOffset<T, IndexType, 1> {
  static __host__ __device__ IndexType
  get(IndexType linearId, const TensorInfo<T, IndexType>& info) {
    return linearId * info.strides[0];
  }
};

// Specialization for Dim == 2
template <typename T, typename IndexType>
struct IndexToOffset<T, IndexType, 2> {
  static __host__ __device__ IndexType
  get(IndexType linearId, const TensorInfo<T, IndexType>& info) {
    IndexType curDimIndex = linearId % info.sizes[1];
    IndexType offset = curDimIndex * info.strides[1];
    linearId /= info.sizes[1];
    return offset + linearId * info.strides[0];
  }
};

// Specialization for Dim == 3
template <typename T, typename IndexType>
struct IndexToOffset<T, IndexType, 3> {
  static __host__ __device__ IndexType
  get(IndexType linearId, const TensorInfo<T, IndexType>& info) {
    IndexType curDimIndex = linearId % info.sizes[2];
    IndexType offset = curDimIndex * info.strides[2];
    linearId /= info.sizes[2];
    curDimIndex = linearId % info.sizes[1];
    offset += curDimIndex * info.strides[1];
    linearId /= info.sizes[1];
    return offset + linearId * info.strides[0];
  }
};

// --- inclusiveBinaryPrefixScan / exclusiveBinaryPrefixScan ---
// Prefix scan utilities

template <typename T>
__device__ inline void swapVars(T* t1, T* t2) {
  T tmp = *t1;
  *t1 = *t2;
  *t2 = tmp;
}

template <typename Comparator, typename K, typename V>
__device__ inline void bitonicSwap(K* kA,
                                   V* vA,
                                   bool* validA,
                                   K* kB,
                                   V* vB,
                                   bool* validB,
                                   bool dir,
                                   const Comparator& comp) {
  // Invalid entries always sort to the end
  bool swap = (comp(*kA, *kB) && *validA) || !*validB;
  if (swap == dir) {
    swapVars(kA, kB);
    swapVars(vA, vB);
    swapVars(validA, validB);
  }
}

template <int Power2SortSize,
          typename IndexType,
          typename Comparator,
          typename K,
          typename V>
__device__ inline void bitonicSort(K* keys,
                                   V* values,
                                   bool* valid,
                                   const Comparator& comp) {
#pragma unroll
  for (unsigned int size = 2; size < Power2SortSize; size *= 2) {
    bool flag = ((threadIdx.x & (size / 2)) != 0);

#pragma unroll
    for (unsigned int stride = size / 2; stride > 0; stride /= 2) {
      __syncthreads();

      unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      bitonicSwap<Comparator, K, V>(&keys[pos],
                                    &values[pos],
                                    &valid[pos],
                                    &keys[pos + stride],
                                    &values[pos + stride],
                                    &valid[pos + stride],
                                    flag,
                                    comp);
    }
  }

#pragma unroll
  for (unsigned int stride = Power2SortSize / 2; stride > 0; stride /= 2) {
    __syncthreads();

    unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
    bitonicSwap<Comparator, K, V>(&keys[pos],
                                  &values[pos],
                                  &valid[pos],
                                  &keys[pos + stride],
                                  &values[pos + stride],
                                  &valid[pos + stride],
                                  false,
                                  comp);
  }

  __syncthreads();
}

template <int KeyDims,
          int ValueDims,
          int block_dim_x,
          int max_block_dim_y,
          typename K,
          typename V,
          typename Comparator,
          typename IndexType>
__global__ void __launch_bounds__(block_dim_x* max_block_dim_y)
    bitonicSortKVInPlace(TensorInfo<K, IndexType> keys,
                         IndexType keySlices,
                         IndexType keySliceSize,
                         IndexType keySliceStride,
                         TensorInfo<V, IndexType> values,
                         IndexType valueSliceStride,
                         Comparator comp) {
  const IndexType blockIndex = getLinearBlockId<IndexType>();
  const IndexType linearIndex = blockIndex * blockDim.y + threadIdx.y;

  if (blockIndex * blockDim.y >= keySlices) {
    return;
  }
  const bool row_valid = linearIndex < keySlices;

  constexpr int items_per_thread = 2;
  constexpr int Power2SortSize = block_dim_x * items_per_thread;

  __shared__ K blockSharedKeys[max_block_dim_y][Power2SortSize];
  __shared__ V blockSharedValues[max_block_dim_y][Power2SortSize];
  __shared__ bool blockSharedValid[max_block_dim_y][Power2SortSize];

  auto sharedKeys = blockSharedKeys[threadIdx.y];
  auto sharedValues = blockSharedValues[threadIdx.y];
  auto sharedValid = blockSharedValid[threadIdx.y];

  const IndexType keyStartOffset =
      IndexToOffset<K, IndexType, KeyDims>::get(linearIndex, keys);
  const IndexType valueStartOffset =
      IndexToOffset<V, IndexType, ValueDims>::get(linearIndex, values);

#pragma unroll
  for (int k = 0; k < items_per_thread; ++k) {
    auto idx = threadIdx.x + k * blockDim.x;
    bool valid = row_valid && idx < keySliceSize;

    sharedKeys[idx] =
        valid ? keys.data[idx * keySliceStride + keyStartOffset] : K{};
    sharedValues[idx] =
        valid ? values.data[idx * valueSliceStride + valueStartOffset] : V{};
    sharedValid[idx] = valid;
  }

  bitonicSort<Power2SortSize, IndexType>(
      sharedKeys, sharedValues, sharedValid, comp);

  if (!row_valid) {
    return;
  }

#pragma unroll
  for (int k = 0; k < items_per_thread; ++k) {
    auto idx = threadIdx.x + k * blockDim.x;
    if (idx < keySliceSize) {
      keys.data[idx * keySliceStride + keyStartOffset] = sharedKeys[idx];
      values.data[idx * valueSliceStride + valueStartOffset] =
          sharedValues[idx];
    }
  }
}

template <typename scalar_t, bool handleNaN = false>
struct GTOp {
  __device__ bool operator()(const scalar_t& lhs, const scalar_t& rhs) const {
    return (handleNaN && (lhs != lhs) && !(rhs != rhs)) ||
           (static_cast<scalar_t>(lhs) > static_cast<scalar_t>(rhs));
  }
};

template <typename scalar_t, bool handleNaN = false>
struct LTOp {
  __device__ bool operator()(const scalar_t& lhs, const scalar_t& rhs) const {
    return (handleNaN && !(lhs != lhs) && (rhs != rhs)) ||
           (static_cast<scalar_t>(lhs) < static_cast<scalar_t>(rhs));
  }
};

template <typename T, typename IndexType, int Dim>
void launch_bitonic_sort(TensorInfo<T, IndexType> keyInfo,
                         IndexType keySlices,
                         IndexType keySliceSize,
                         IndexType keySliceStride,
                         TensorInfo<int64_t, IndexType> valueInfo,
                         IndexType valueSliceStride,
                         bool largest,
                         gpuStream_t stream) {
  constexpr int sort_size = 32;
  constexpr int max_block_y = 16;
  constexpr int items_per_thread = 2;
  constexpr int block_x = sort_size / items_per_thread;

  const int block_y = std::min(
      static_cast<int>(max_block_y),
      static_cast<int>(std::max(static_cast<IndexType>(1), keySlices)));
  dim3 block(block_x, block_y);

  dim3 grid;
  const int grid_count = (keySlices + block_y - 1) / block_y;
  getGridFromTiles(grid_count, &grid);

  if (largest) {
    bitonicSortKVInPlace<Dim, Dim, block_x, max_block_y>
        <<<grid, block, 0, stream>>>(keyInfo,
                                     keySlices,
                                     keySliceSize,
                                     keySliceStride,
                                     valueInfo,
                                     valueSliceStride,
                                     GTOp<T, true>());
  } else {
    bitonicSortKVInPlace<Dim, Dim, block_x, max_block_y>
        <<<grid, block, 0, stream>>>(keyInfo,
                                     keySlices,
                                     keySliceSize,
                                     keySliceStride,
                                     valueInfo,
                                     valueSliceStride,
                                     LTOp<T, true>());
  }
}

// ============================================================================
// StridedRandomAccessor
// Required by CUB WarpLoad/WarpStore and BlockLoad/BlockStore for strided
// tensor access.
// ============================================================================

template <typename T, typename index_t = int64_t>
class ConstStridedRandomAccessor {
 public:
  using difference_type = index_t;
  using value_type = const T;
  using pointer = const T*;
  using reference = const T&;
  using iterator_category = std::random_access_iterator_tag;
  using PtrType = T*;
  using index_type = index_t;

  __host__ __device__ ConstStridedRandomAccessor(PtrType ptr, index_t stride)
      : ptr_{ptr}, stride_{stride} {}
  __host__ __device__ explicit ConstStridedRandomAccessor(PtrType ptr)
      : ptr_{ptr}, stride_{1} {}
  __host__ __device__ ConstStridedRandomAccessor()
      : ptr_{nullptr}, stride_{1} {}

  __host__ __device__ reference operator*() const { return *ptr_; }
  __host__ __device__ const T* operator->() const {
    return reinterpret_cast<const T*>(ptr_);
  }
  __host__ __device__ reference operator[](index_t idx) const {
    return ptr_[idx * stride_];
  }

  __host__ __device__ ConstStridedRandomAccessor& operator++() {
    ptr_ += stride_;
    return *this;
  }
  __host__ __device__ ConstStridedRandomAccessor operator++(int) {
    ConstStridedRandomAccessor copy(*this);
    ++*this;
    return copy;
  }
  __host__ __device__ ConstStridedRandomAccessor& operator--() {
    ptr_ -= stride_;
    return *this;
  }
  __host__ __device__ ConstStridedRandomAccessor operator--(int) {
    ConstStridedRandomAccessor copy(*this);
    --*this;
    return copy;
  }
  __host__ __device__ ConstStridedRandomAccessor& operator+=(index_t offset) {
    ptr_ += offset * stride_;
    return *this;
  }
  __host__ __device__ ConstStridedRandomAccessor
  operator+(index_t offset) const {
    return ConstStridedRandomAccessor(ptr_ + offset * stride_, stride_);
  }
  __host__ __device__ friend ConstStridedRandomAccessor operator+(
      index_t offset, const ConstStridedRandomAccessor& accessor) {
    return accessor + offset;
  }
  __host__ __device__ ConstStridedRandomAccessor& operator-=(index_t offset) {
    ptr_ -= offset * stride_;
    return *this;
  }
  __host__ __device__ ConstStridedRandomAccessor
  operator-(index_t offset) const {
    return ConstStridedRandomAccessor(ptr_ - offset * stride_, stride_);
  }
  __host__ __device__ difference_type
  operator-(const ConstStridedRandomAccessor& other) const {
    return (ptr_ - other.ptr_) / stride_;
  }
  __host__ __device__ bool operator==(
      const ConstStridedRandomAccessor& other) const {
    return (ptr_ == other.ptr_) && (stride_ == other.stride_);
  }
  __host__ __device__ bool operator!=(
      const ConstStridedRandomAccessor& other) const {
    return !(*this == other);
  }
  __host__ __device__ bool operator<(
      const ConstStridedRandomAccessor& other) const {
    return ptr_ < other.ptr_;
  }
  __host__ __device__ bool operator<=(
      const ConstStridedRandomAccessor& other) const {
    return (*this < other) || (*this == other);
  }
  __host__ __device__ bool operator>(
      const ConstStridedRandomAccessor& other) const {
    return !(*this <= other);
  }
  __host__ __device__ bool operator>=(
      const ConstStridedRandomAccessor& other) const {
    return !(*this < other);
  }

 protected:
  PtrType ptr_;
  index_t stride_;
};

template <typename T, typename index_t = int64_t>
class StridedRandomAccessor : public ConstStridedRandomAccessor<T, index_t> {
 public:
  using difference_type = index_t;
  using value_type = T;
  using pointer = T*;
  using reference = T&;
  using BaseType = ConstStridedRandomAccessor<T, index_t>;
  using PtrType = T*;

  __host__ __device__ StridedRandomAccessor(PtrType ptr, index_t stride)
      : BaseType(ptr, stride) {}
  __host__ __device__ explicit StridedRandomAccessor(PtrType ptr)
      : BaseType(ptr) {}
  __host__ __device__ StridedRandomAccessor() : BaseType() {}

  __host__ __device__ reference operator*() const { return *this->ptr_; }
  __host__ __device__ T* operator->() const {
    return reinterpret_cast<T*>(this->ptr_);
  }
  __host__ __device__ reference operator[](index_t idx) const {
    return this->ptr_[idx * this->stride_];
  }

  __host__ __device__ StridedRandomAccessor& operator++() {
    this->ptr_ += this->stride_;
    return *this;
  }
  __host__ __device__ StridedRandomAccessor operator++(int) {
    StridedRandomAccessor copy(*this);
    ++*this;
    return copy;
  }
  __host__ __device__ StridedRandomAccessor& operator--() {
    this->ptr_ -= this->stride_;
    return *this;
  }
  __host__ __device__ StridedRandomAccessor operator--(int) {
    StridedRandomAccessor copy(*this);
    --*this;
    return copy;
  }
  __host__ __device__ StridedRandomAccessor& operator+=(index_t offset) {
    this->ptr_ += offset * this->stride_;
    return *this;
  }
  __host__ __device__ StridedRandomAccessor operator+(index_t offset) const {
    return StridedRandomAccessor(this->ptr_ + offset * this->stride_,
                                 this->stride_);
  }
  __host__ __device__ friend StridedRandomAccessor operator+(
      index_t offset, const StridedRandomAccessor& accessor) {
    return accessor + offset;
  }
  __host__ __device__ StridedRandomAccessor& operator-=(index_t offset) {
    this->ptr_ -= offset * this->stride_;
    return *this;
  }
  __host__ __device__ StridedRandomAccessor operator-(index_t offset) const {
    return StridedRandomAccessor(this->ptr_ - offset * this->stride_,
                                 this->stride_);
  }
  __host__ __device__ difference_type operator-(const BaseType& other) const {
    return (static_cast<const BaseType&>(*this) - other);
  }
};

// ============================================================================
// CubKeyType mapping - maps Paddle types to CUB-compatible CUDA types
// For BlockRadixSort, CUB needs __half / __nv_bfloat16 instead of
// phi::float16 / phi::bfloat16.
// ============================================================================
template <typename T>
struct CubKeyType {
  using type = T;
};

template <>
struct CubKeyType<phi::dtype::float16> {
  using type = __half;
};

template <>
struct CubKeyType<phi::dtype::bfloat16> {
#if defined(__HIPCC__)
  using type = hip_bfloat16;
#else
  using type = __nv_bfloat16;
#endif
};

// ============================================================================
// Utility functions
// ============================================================================

inline int64_t nextHighestPowerOf2(int64_t n) {
  n--;
  n |= n >> 1;
  n |= n >> 2;
  n |= n >> 4;
  n |= n >> 8;
  n |= n >> 16;
  n |= n >> 32;
  n++;
  return n;
}

template <typename T>
static int minimum_grid_for_occupancy(T kernel, int max_block_size) {
  int minGridSize = 0;
  int blockSize = 0;
  cudaOccupancyMaxPotentialBlockSize(
      &minGridSize, &blockSize, kernel, /*dynamicSMemSize=*/0, max_block_size);
  return minGridSize;
}

template <typename T>
constexpr bool type_has_nan() {
  if constexpr (std::numeric_limits<T>::is_specialized) {
    return std::numeric_limits<T>::has_quiet_NaN;
  } else if constexpr (std::is_same_v<T, phi::dtype::float16> ||  // NOLINT
                       std::is_same_v<T, phi::dtype::bfloat16>) {
    return true;
  } else {
    return false;
  }
}

// ============================================================================
// warpMergeSortKVInPlace kernel
// For sort sizes 33..128, uses CUB WarpMergeSort (one warp per slice,
// multiple slices per block via blockDim.y).
// ============================================================================

template <int KeyDims,
          int ValueDims,
          int sort_size,
          int max_block_dim_y,
          typename K,
          typename V,
          typename Comparator,
          typename IndexType>
__global__ void __launch_bounds__(32 * max_block_dim_y)
    warpMergeSortKVInPlace(TensorInfo<K, IndexType> keys,
                           IndexType keySlices,
                           IndexType keySliceSize,
                           IndexType keySliceStride,
                           TensorInfo<V, IndexType> values,
                           IndexType valueSliceStride,
                           Comparator comp,
                           K invalid_key) {
  const IndexType blockIndex = getLinearBlockId<IndexType>();
  const IndexType linearIndex = blockIndex * blockDim.y + threadIdx.y;

  if (linearIndex >= keySlices) {
    return;
  }

  const IndexType keyStartOffset =
      IndexToOffset<K, IndexType, KeyDims>::get(linearIndex, keys);
  const IndexType valueStartOffset =
      IndexToOffset<V, IndexType, ValueDims>::get(linearIndex, values);

  K* keys_slice = &keys.data[keyStartOffset];
  V* values_slice = &values.data[valueStartOffset];

  StridedRandomAccessor<K, IndexType> keys_iter(keys_slice, keySliceStride);
  StridedRandomAccessor<V, IndexType> values_iter(values_slice,
                                                  valueSliceStride);

  constexpr int warp_size = 32;
  constexpr int kItemsPerThread = sort_size / warp_size;
  static_assert(kItemsPerThread * warp_size == sort_size,
                "sort_size must be a multiple of warp_size (32)");

  using LoadKeys = cub::WarpLoad<K, kItemsPerThread, cub::WARP_LOAD_TRANSPOSE>;
  using LoadValues =
      cub::WarpLoad<V, kItemsPerThread, cub::WARP_LOAD_TRANSPOSE>;
  using Sort = cub::WarpMergeSort<K, kItemsPerThread, warp_size, V>;
  using StoreKeys =
      cub::WarpStore<K, kItemsPerThread, cub::WARP_STORE_TRANSPOSE>;
  using StoreValues =
      cub::WarpStore<V, kItemsPerThread, cub::WARP_STORE_TRANSPOSE>;

  __shared__ union {
    typename LoadKeys::TempStorage load_keys;
    typename LoadValues::TempStorage load_values;
    typename Sort::TempStorage sort;
    typename StoreKeys::TempStorage store_keys;
    typename StoreValues::TempStorage store_values;
  } tmp_storage[max_block_dim_y];

  auto& warp_storage = tmp_storage[threadIdx.y];

  K local_keys[kItemsPerThread];
  V local_values[kItemsPerThread];

  const auto invalid_value = V{};
  LoadKeys(warp_storage.load_keys)
      .Load(keys_iter, local_keys, keySliceSize, invalid_key);
#if !defined(__HIPCC__)
  __syncwarp();
#endif
  LoadValues(warp_storage.load_values)
      .Load(values_iter, local_values, keySliceSize, invalid_value);
#if !defined(__HIPCC__)
  __syncwarp();
#endif

  Sort(warp_storage.sort)
      .StableSort(local_keys, local_values, comp, keySliceSize, invalid_key);
#if !defined(__HIPCC__)
  __syncwarp();
#endif

  StoreKeys(warp_storage.store_keys).Store(keys_iter, local_keys, keySliceSize);
#if !defined(__HIPCC__)
  __syncwarp();
#endif
  StoreValues(warp_storage.store_values)
      .Store(values_iter, local_values, keySliceSize);
}

// ============================================================================
// radixSortKVInPlace kernel
// For sort sizes 129..4096, uses CUB BlockRadixSort (one block per slice).
// ============================================================================

template <int KeyDims,
          int ValueDims,
          int block_size,
          int kItemsPerThread,
          typename K,
          typename V,
          typename IndexType>
__global__ void __launch_bounds__(block_size)
    radixSortKVInPlace(TensorInfo<K, IndexType> keys,
                       IndexType keySlices,
                       IndexType keySliceSize,
                       IndexType keySliceStride,
                       TensorInfo<V, IndexType> values,
                       IndexType valueSliceStride,
                       bool descending) {
  static_assert(block_size > 0, "");

  const IndexType linearIndex = getLinearBlockId<IndexType>();
  if (linearIndex >= keySlices) {
    return;
  }

  const IndexType keyStartOffset =
      IndexToOffset<K, IndexType, KeyDims>::get(linearIndex, keys);
  const IndexType valueStartOffset =
      IndexToOffset<V, IndexType, ValueDims>::get(linearIndex, values);

  K* keys_slice = &keys.data[keyStartOffset];
  V* values_slice = &values.data[valueStartOffset];

  StridedRandomAccessor<K, IndexType> keys_iter(keys_slice, keySliceStride);
  StridedRandomAccessor<V, IndexType> values_iter(values_slice,
                                                  valueSliceStride);

  using key_t = typename CubKeyType<K>::type;
  using LoadKeys =
      cub::BlockLoad<K, block_size, kItemsPerThread, cub::BLOCK_LOAD_TRANSPOSE>;
  using LoadValues =
      cub::BlockLoad<V, block_size, kItemsPerThread, cub::BLOCK_LOAD_TRANSPOSE>;
  using Sort = cub::BlockRadixSort<key_t, block_size, kItemsPerThread, V>;
  using StoreKeys = cub::
      BlockStore<K, block_size, kItemsPerThread, cub::BLOCK_STORE_TRANSPOSE>;
  using StoreValues = cub::
      BlockStore<V, block_size, kItemsPerThread, cub::BLOCK_STORE_TRANSPOSE>;

  __shared__ union {
    typename LoadKeys::TempStorage load_keys;
    typename LoadValues::TempStorage load_values;
    typename Sort::TempStorage sort;
    typename StoreKeys::TempStorage store_keys;
    typename StoreValues::TempStorage store_values;
  } tmp_storage;

  // Compute invalid key: always sorts higher than any valid key
  const K invalid_key = [descending] {
    using radix_t = typename cub::Traits<key_t>::UnsignedBits;
    union {
      K key;
      radix_t radix;
    } tmp;
    tmp.radix = descending ? cub::Traits<key_t>::LOWEST_KEY
                           : cub::Traits<key_t>::MAX_KEY;
    return tmp.key;
  }();
  const V invalid_value = static_cast<V>(0);

  K local_keys[kItemsPerThread];
  V local_values[kItemsPerThread];

  LoadKeys(tmp_storage.load_keys)
      .Load(keys_iter, local_keys, keySliceSize, invalid_key);
  __syncthreads();
  LoadValues(tmp_storage.load_values)
      .Load(values_iter, local_values, keySliceSize, invalid_value);
  __syncthreads();

  if (descending) {
    Sort(tmp_storage.sort)
        .SortDescending(reinterpret_cast<key_t(&)[kItemsPerThread]>(local_keys),
                        local_values);
  } else {
    Sort(tmp_storage.sort)
        .Sort(reinterpret_cast<key_t(&)[kItemsPerThread]>(local_keys),
              local_values);
  }
  __syncthreads();

  StoreKeys(tmp_storage.store_keys).Store(keys_iter, local_keys, keySliceSize);
  __syncthreads();
  StoreValues(tmp_storage.store_values)
      .Store(values_iter, local_values, keySliceSize);
}

// ============================================================================
// launch_warp_merge_sort - wrapper for CUB WarpMergeSort<128>
// ============================================================================

template <typename T, typename IndexType, int Dim>
void launch_warp_merge_sort(TensorInfo<T, IndexType> keyInfo,
                            IndexType keySlices,
                            IndexType keySliceSize,
                            IndexType keySliceStride,
                            TensorInfo<int64_t, IndexType> valueInfo,
                            IndexType valueSliceStride,
                            bool largest,
                            gpuStream_t stream) {
  constexpr int sort_size = 128;
  constexpr int max_block_dim_y = 16;
  constexpr int warp_size = 32;

  // Scale batch size down if the grid would be too small
  const auto min_grid =
      minimum_grid_for_occupancy(warpMergeSortKVInPlace<Dim,
                                                        Dim,
                                                        sort_size,
                                                        max_block_dim_y,
                                                        T,
                                                        int64_t,
                                                        LTOp<T, true>,
                                                        IndexType>,
                                 warp_size * max_block_dim_y);
  const auto max_batch =
      std::max(IndexType{1}, keySlices / (IndexType)min_grid);
  const int block_y = std::min((IndexType)max_block_dim_y, max_batch);
  dim3 block(warp_size, block_y);

  dim3 grid;
  const int grid_count = (keySlices + block_y - 1) / block_y;
  getGridFromTiles(grid_count, &grid);

  if (largest) {
    // Use numeric limits for invalid_key: lower_bound for descending
    const T invalid_key = std::numeric_limits<T>::lowest();
    warpMergeSortKVInPlace<Dim, Dim, sort_size, max_block_dim_y>
        <<<grid, block, 0, stream>>>(keyInfo,
                                     keySlices,
                                     keySliceSize,
                                     keySliceStride,
                                     valueInfo,
                                     valueSliceStride,
                                     GTOp<T, true>(),
                                     invalid_key);
  } else {
    // For ascending: NAN sorts after inf, otherwise use upper_bound
    const T invalid_key = [] {
      if constexpr (type_has_nan<T>()) {
        return T(NAN);
      }
      return std::numeric_limits<T>::max();
    }();
    warpMergeSortKVInPlace<Dim, Dim, sort_size, max_block_dim_y>
        <<<grid, block, 0, stream>>>(keyInfo,
                                     keySlices,
                                     keySliceSize,
                                     keySliceStride,
                                     valueInfo,
                                     valueSliceStride,
                                     LTOp<T, true>(),
                                     invalid_key);
  }
}

// ============================================================================
// launch_medium_radix_sort - wrapper for CUB BlockRadixSort
// ============================================================================

template <int Dim,
          int sort_size,
          int items_per_thread,
          typename K,
          typename IndexType>
void fixed_size_radix_sort(TensorInfo<K, IndexType> keyInfo,
                           IndexType keySlices,
                           IndexType keySliceSize,
                           IndexType keySliceStride,
                           TensorInfo<int64_t, IndexType> valueInfo,
                           IndexType valueSliceStride,
                           bool descending,
                           gpuStream_t stream) {
  static_assert(sort_size % items_per_thread == 0, "");
  constexpr int block = sort_size / items_per_thread;
  dim3 grid;
  getGridFromTiles(keySlices, &grid);

  radixSortKVInPlace<Dim, Dim, block, items_per_thread>
      <<<grid, block, 0, stream>>>(keyInfo,
                                   keySlices,
                                   keySliceSize,
                                   keySliceStride,
                                   valueInfo,
                                   valueSliceStride,
                                   descending);
}

template <typename T, typename IndexType, int Dim>
void launch_medium_radix_sort(TensorInfo<T, IndexType> keyInfo,
                              IndexType keySlices,
                              IndexType keySliceSize,
                              IndexType keySliceStride,
                              TensorInfo<int64_t, IndexType> valueInfo,
                              IndexType valueSliceStride,
                              bool descending,
                              gpuStream_t stream) {
  int64_t ceilPowerOf2 = nextHighestPowerOf2(keySliceSize);
  constexpr int default_ipt = 32;

#define HANDLE_RADIX_CASE(SIZE, IPT)                      \
  fixed_size_radix_sort<Dim, SIZE, IPT>(keyInfo,          \
                                        keySlices,        \
                                        keySliceSize,     \
                                        keySliceStride,   \
                                        valueInfo,        \
                                        valueSliceStride, \
                                        descending,       \
                                        stream)

  switch (ceilPowerOf2) {
    case 4096:
      HANDLE_RADIX_CASE(4096, default_ipt);
      break;
    case 2048:
      HANDLE_RADIX_CASE(2048, default_ipt);
      break;
    case 1024:
    case 512:
    case 256:
      HANDLE_RADIX_CASE(1024, default_ipt);
      break;
    // sizes <= 128 should have been handled by WarpMergeSort
    default:
      break;
  }
#undef HANDLE_RADIX_CASE
}

template <typename T, bool KillWARDependency, class BinaryFunction>
__device__ void inclusiveBinaryPrefixScan(T* smem,
                                          bool in,
                                          T* out,
                                          BinaryFunction binop) {
  T vote = TOPK_WARP_BALLOT(in);
  T index = __popc(getLaneMaskLe() & vote);
  T carry = __popc(vote);

  int warp = threadIdx.x / TOPK_WARP_SIZE;

  if (getLaneId() == 0) {
    smem[warp] = carry;
  }

  __syncthreads();

  if (threadIdx.x == 0) {
    int current = 0;
    for (int i = 0;
         i < topk_ceil_div(static_cast<int>(blockDim.x), TOPK_WARP_SIZE);
         ++i) {
      T v = smem[i];
      smem[i] = binop(smem[i], current);
      current = binop(current, v);
    }
  }

  __syncthreads();

  if (warp >= 1) {
    index = binop(index, smem[warp - 1]);
  }

  *out = index;

  if (KillWARDependency) {
    __syncthreads();
  }
}

template <typename T, bool KillWARDependency, class BinaryFunction>
__device__ void exclusiveBinaryPrefixScan(
    T* smem, bool in, T* out, T* carry, BinaryFunction binop) {
  inclusiveBinaryPrefixScan<T, false, BinaryFunction>(smem, in, out, binop);
  *out -= static_cast<T>(in);
  *carry =
      smem[topk_ceil_div(static_cast<int>(blockDim.x), TOPK_WARP_SIZE) - 1];
  if (KillWARDependency) {
    __syncthreads();
  }
}

// --- AddOp ---
template <typename T>
struct AddOp {
  __device__ __forceinline__ T operator()(T const& lhs, T const& rhs) {
    return (lhs + rhs);
  }
};

// ============================================================================
// SortingRadixSelect.cuh ported content
// ============================================================================

namespace radix_select {

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

// CountType is separate from IndexType — counts always fit in int32
// because indices are limited to integer fp precision.
template <typename T,
          typename RadixType,
          typename IndexType,
          typename CountType,
          int RadixSize,
          int RadixBits>
__device__ void countRadixUsingMask(const T* data,
                                    CountType counts[RadixSize],
                                    CountType* smem,
                                    RadixType desired,
                                    RadixType desiredMask,
                                    int radixDigitPos,
                                    IndexType sliceSize,
                                    IndexType withinSliceStride) {
#pragma unroll
  for (int i = 0; i < RadixSize; ++i) {
    counts[i] = 0;
  }
  if (threadIdx.x < RadixSize) {
    smem[threadIdx.x] = 0;
  }
  __syncthreads();

  // Must be called outside of loop to ensure all threads participate.
  // This creates a dynamic mask of which threads will enter the loop.
  // When sliceSize < blockDim.x, only threads with threadIdx.x < sliceSize
  // will enter the loop body, so we need a mask to avoid deadlock in
  // __ballot_sync.
#if !defined(__HIPCC__)
  unsigned mask = TOPK_WARP_BALLOT(threadIdx.x < sliceSize);
#endif
  for (IndexType i = threadIdx.x; i < sliceSize;) {
    RadixType val =
        TopKTypeConfig<T>::convert(doLdg(&data[i * withinSliceStride]));
    bool hasVal = ((val & desiredMask) == desired);
    RadixType digitInRadix =
        Bitfield<RadixType>::getBitfield(val, radixDigitPos, RadixBits);
#pragma unroll
    for (uint32_t j = 0; j < RadixSize; ++j) {
      bool vote = hasVal && (digitInRadix == j);
#if defined(__HIPCC__)
      counts[j] += __popcll(TOPK_WARP_BALLOT(vote));
#else
      counts[j] += __popc(TOPK_WARP_BALLOT_MASK(vote, mask));
#endif
    }
    i += blockDim.x;
#if !defined(__HIPCC__)
    mask = TOPK_WARP_BALLOT_MASK(i < sliceSize, mask);
#endif
  }

  if (getLaneId() == 0) {
#pragma unroll
    for (uint32_t i = 0; i < RadixSize; ++i) {
      atomicAdd(&smem[i], counts[i]);
    }
  }
  __syncthreads();

#pragma unroll
  for (uint32_t i = 0; i < RadixSize; ++i) {
    counts[i] = smem[i];
  }
  __syncthreads();
}

template <typename T, typename RadixType, typename IndexType>
__device__ T findPattern(const T* data,
                         T* smem,
                         IndexType sliceSize,
                         IndexType withinSliceStride,
                         RadixType desired,
                         RadixType desiredMask) {
  if (threadIdx.x < 2) {
    smem[threadIdx.x] = static_cast<T>(0);
  }
  __syncthreads();

  IndexType numIterations = topk_round_up(sliceSize, (IndexType)blockDim.x);
  for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
    bool inRange = (i < sliceSize);
    T v = inRange ? doLdg(&data[i * withinSliceStride]) : static_cast<T>(0);
    if (inRange && ((TopKTypeConfig<T>::convert(v) & desiredMask) == desired)) {
      smem[0] = static_cast<T>(1);
      smem[1] = v;
    }
    __syncthreads();
    T found = smem[0];
    T val = smem[1];
    __syncthreads();
    if (found != static_cast<T>(0)) {
      return val;
    }
  }
  // should not get here
  assert(false);
  return static_cast<T>(0);
}

template <typename T, typename RadixType, typename IndexType>
__device__ void radixSelect(const T* data,
                            IndexType k,
                            bool largest,
                            IndexType sliceSize,
                            IndexType withinSliceStride,
                            int* smem,
                            T* topKValue) {
  // Indices are limited to integer fp precision, so counts can fit in
  // int32, regardless of IndexType
  int counts[RADIX_SIZE];
  RadixType desired = 0;
  RadixType desiredMask = 0;

  IndexType kToFind = k;

#pragma unroll
  for (int digitPos = sizeof(T) * 8 - RADIX_BITS; digitPos >= 0;
       digitPos -= RADIX_BITS) {
    countRadixUsingMask<T, RadixType, IndexType, int, RADIX_SIZE, RADIX_BITS>(
        data,
        counts,
        smem,
        desired,
        desiredMask,
        digitPos,
        sliceSize,
        withinSliceStride);

    auto found_unique = [&](int i, int count) -> bool {
      if (count == 1 && kToFind == 1) {
        desired =
            Bitfield<RadixType>::setBitfield(desired, i, digitPos, RADIX_BITS);
        desiredMask = Bitfield<RadixType>::setBitfield(
            desiredMask, RADIX_MASK, digitPos, RADIX_BITS);
        *topKValue =
            findPattern<T, RadixType, IndexType>(data,
                                                 reinterpret_cast<T*>(smem),
                                                 sliceSize,
                                                 withinSliceStride,
                                                 desired,
                                                 desiredMask);
        return true;
      }
      return false;
    };
    auto found_non_unique = [&](int i, int count) -> bool {
      if (count >= kToFind) {
        desired =
            Bitfield<RadixType>::setBitfield(desired, i, digitPos, RADIX_BITS);
        desiredMask = Bitfield<RadixType>::setBitfield(
            desiredMask, RADIX_MASK, digitPos, RADIX_BITS);
        return true;
      }
      kToFind -= count;
      return false;
    };

    if (largest) {
#pragma unroll
      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 {
#pragma unroll
      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;
      }
    }
  }
  *topKValue = TopKTypeConfig<T>::deconvert(desired);
}

}  // namespace radix_select

// ============================================================================
// CUDA_KERNEL_LOOP_TYPE macro
// ============================================================================
#define TOPK_CUDA_KERNEL_LOOP_TYPE(i, n, index_type)                  \
  for (index_type i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
       i += blockDim.x * gridDim.x)

// ============================================================================
// CUB_SUPPORTS_SCAN_BY_KEY check
// CUB >= 1.15 supports DeviceScan::InclusiveSumByKey
// ============================================================================
#ifndef __HIPCC__
// CUDA path: check CUB version
#if defined(CUB_VERSION) && CUB_VERSION >= 101500
#define TOPK_CUB_SUPPORTS_SCAN_BY_KEY() 1
#else
// Try to detect based on CUDA version (CUDA 11.6+ bundles CUB >= 1.15)
#if CUDART_VERSION >= 11060
#define TOPK_CUB_SUPPORTS_SCAN_BY_KEY() 1
#else
#define TOPK_CUB_SUPPORTS_SCAN_BY_KEY() 0
#endif
#endif
#else
// HIP/ROCm path
#define TOPK_CUB_SUPPORTS_SCAN_BY_KEY() 0
#endif

}  // namespace topk_detail

// ============================================================================
// Main TopK implementation
// ============================================================================

namespace topk_impl {

using namespace topk_detail;  // NOLINT

// getTensorInfo: builds TensorInfo from DenseTensor
template <typename T, typename IndexType>
TensorInfo<T, IndexType> getTensorInfo(const phi::DenseTensor& tensor) {
  TensorInfo<T, IndexType> info;
  info.data = reinterpret_cast<T*>(const_cast<void*>(tensor.data()));
  info.dims = tensor.dims().size();
  for (int i = 0; i < info.dims; i++) {
    info.sizes[i] = tensor.dims()[i];
    info.strides[i] = tensor.strides()[i];
  }
  return info;
}

// SegmentOffsetIter for sorted output - must be at namespace scope for CUDA
struct SegmentOffsetIter {
  int64_t k;
  __host__ __device__ __forceinline__ int64_t operator()(int64_t idx) const {
    return idx * k;
  }
};

template <typename T, typename Context>
void sortKeyValueInplace(const Context& dev_ctx,
                         phi::DenseTensor* out,
                         phi::DenseTensor* indices,
                         int axis,
                         bool largest) {
  const auto& out_dims = out->dims();
  int dim = axis;
  int64_t sliceSize = out_dims[dim];
  int64_t numSlices = out->numel() / sliceSize;
  auto stream = dev_ctx.stream();

  if (sliceSize <= 1) return;

  auto keyInfo = getTensorInfo<T, uint32_t>(*out);
  auto valueInfo = getTensorInfo<int64_t, uint32_t>(*indices);

  auto strideKey = keyInfo.strides[dim];
  keyInfo.sizes[dim] = 1;
  int collapseKeyDim = keyInfo.collapseDims(dim);
  keyInfo.strides[collapseKeyDim] = strideKey;

  auto strideValue = valueInfo.strides[dim];
  valueInfo.sizes[dim] = 1;
  int collapseValueDim = valueInfo.collapseDims(dim);
  valueInfo.strides[collapseValueDim] = strideValue;

  // Three-tier sort dispatch:
  //   1. sliceSize <= 32:  Bitonic Sort (unstable, fast, no extra memory)
  //   2. sliceSize <= 128: WarpMergeSort (CUB, one slice per warp)
  //   3. sliceSize <= 4096: BlockRadixSort (CUB, one slice per block)
  // Dispatch on the actual number of collapsed dims (keyInfo.dims),
  // NOT on collapseKeyDim (the remapped excluded-dim index).
  // When the excluded dim is in the middle (e.g. dim=1 of a 3-D tensor),
  // collapseKeyDim==1 but keyInfo.dims==3; using DIM=1 would make
  // IndexToOffset ignore the trailing dimensions, producing wrong offsets.

#define TOPK_SORT_DIM_DISPATCH(LAUNCH_FUNC) \
  if (keyInfo.dims == 1) {                  \
    LAUNCH_FUNC(1);                         \
  } else if (keyInfo.dims == 2) {           \
    LAUNCH_FUNC(2);                         \
  } else if (keyInfo.dims == 3) {           \
    LAUNCH_FUNC(3);                         \
  } else {                                  \
    LAUNCH_FUNC(-1);                        \
  }

  if (sliceSize <= 32) {
    // Bitonic sort (unstable)
#define LAUNCH_BITONIC(DIM)                          \
  launch_bitonic_sort<T, uint32_t, DIM>(keyInfo,     \
                                        numSlices,   \
                                        sliceSize,   \
                                        strideKey,   \
                                        valueInfo,   \
                                        strideValue, \
                                        largest,     \
                                        stream)
    TOPK_SORT_DIM_DISPATCH(LAUNCH_BITONIC);
#undef LAUNCH_BITONIC
  } else if (sliceSize <= 128) {
    // WarpMergeSort (stable, uses CUB WarpMergeSort)
#define LAUNCH_WARP(DIM)                                \
  launch_warp_merge_sort<T, uint32_t, DIM>(keyInfo,     \
                                           numSlices,   \
                                           sliceSize,   \
                                           strideKey,   \
                                           valueInfo,   \
                                           strideValue, \
                                           largest,     \
                                           stream)
    TOPK_SORT_DIM_DISPATCH(LAUNCH_WARP);
#undef LAUNCH_WARP
  } else {
    // BlockRadixSort (for sizes up to 4096)
    bool descending = largest;
#define LAUNCH_RADIX(DIM)                                 \
  launch_medium_radix_sort<T, uint32_t, DIM>(keyInfo,     \
                                             numSlices,   \
                                             sliceSize,   \
                                             strideKey,   \
                                             valueInfo,   \
                                             strideValue, \
                                             descending,  \
                                             stream)
    TOPK_SORT_DIM_DISPATCH(LAUNCH_RADIX);
#undef LAUNCH_RADIX
  }

#undef TOPK_SORT_DIM_DISPATCH
}

namespace sbtopk {  // single_block_topk

template <typename T, typename IndexType, int Dim, bool WithKthValues>
__global__ void __launch_bounds__(1024)
    gatherTopK(TensorInfo<const T, IndexType> input,
               IndexType inputSliceSize,
               IndexType outputSliceSize,  // aka `k`
               bool largest,
               IndexType numInputSlices,
               IndexType inputWithinSliceStride,
               TensorInfo<T, IndexType> topK,
               IndexType topKWithinSliceStride,
               TensorInfo<int64_t, IndexType> indices,
               IndexType indicesWithinSliceStride,
               T* kthValues) {
  // Indices are limited to integer fp precision, so counts can fit in
  // int32, regardless of IndexType
#if defined(__HIPCC__)
  __shared__ int smem[64];
#else
  __shared__ int smem[32];  // one per each warp, up to warp limit
#endif
  IndexType slice = getLinearBlockId<IndexType>();
  if (slice >= numInputSlices) {
    return;
  }

  // Find the start offset for our slice
  IndexType sliceStartIndex =
      IndexToOffset<const T, IndexType, Dim>::get(slice, input);
  IndexType topKSliceStartIndex =
      IndexToOffset<T, IndexType, Dim>::get(slice, topK);
  IndexType indicesSliceStartIndex =
      IndexToOffset<int64_t, IndexType, Dim>::get(slice, indices);

  const T* inputSliceStart = &input.data[sliceStartIndex];
  T* topKSliceStart = &topK.data[topKSliceStartIndex];
  int64_t* indicesSliceStart = &indices.data[indicesSliceStartIndex];

  // Find the k-th highest element in our input
  T topKValue;
  if (WithKthValues) {
    topKValue = kthValues[slice];
  } else {
    topKValue = static_cast<T>(0);
    radix_select::radixSelect<T,
                              typename TopKTypeConfig<T>::RadixType,
                              IndexType>(inputSliceStart,
                                         outputSliceSize,
                                         largest,
                                         inputSliceSize,
                                         inputWithinSliceStride,
                                         smem,
                                         &topKValue);
  }
  const auto topKConverted = TopKTypeConfig<T>::convert(topKValue);

  IndexType numIterations =
      topk_round_up(inputSliceSize, (IndexType)blockDim.x);
  IndexType writeIndexStart = 0;

  for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
    bool inRange = (i < inputSliceSize);
    T v = inRange ? doLdg(&inputSliceStart[i * inputWithinSliceStride])
                  : static_cast<T>(0);
    const auto convertedV = TopKTypeConfig<T>::convert(v);
    bool hasTopK;
    if (largest) {
      hasTopK = inRange && (convertedV > topKConverted);
    } else {
      hasTopK = inRange && (convertedV < topKConverted);
    }

    int index;
    int carry;
    exclusiveBinaryPrefixScan<int, true>(
        smem, hasTopK, &index, &carry, AddOp<int>());

    if (hasTopK) {
      int writeIndex = writeIndexStart + index;
      assert(writeIndex < outputSliceSize);
      IndexType topKOffset = writeIndex * topKWithinSliceStride;
      IndexType indexOffset = writeIndex * indicesWithinSliceStride;
      topKSliceStart[topKOffset] = v;
      indicesSliceStart[indexOffset] = i;
    }
    writeIndexStart += carry;
  }

  // Fill in the rest with actual == top-K values.
  assert(outputSliceSize >= writeIndexStart);
  IndexType topKRemaining = (outputSliceSize - writeIndexStart);

  for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
    bool inRange = (i < inputSliceSize);
    T v = inRange ? doLdg(&inputSliceStart[i * inputWithinSliceStride])
                  : static_cast<T>(0);
    const auto convertedV = TopKTypeConfig<T>::convert(v);
    bool hasTopK = inRange && (convertedV == topKConverted);

    int index;
    int carry;
    exclusiveBinaryPrefixScan<int, true>(
        smem, hasTopK, &index, &carry, AddOp<int>());

    if (hasTopK && index < topKRemaining) {
      int writeIndex = writeIndexStart + index;
      assert(writeIndex < outputSliceSize);
      IndexType topKOffset = writeIndex * topKWithinSliceStride;
      IndexType indexOffset = writeIndex * indicesWithinSliceStride;
      topKSliceStart[topKOffset] = v;
      indicesSliceStart[indexOffset] = i;
    }

    if (carry >= topKRemaining) {
      break;
    }
    topKRemaining -= carry;
    writeIndexStart += carry;
  }
}

template <typename T, typename IndexType, int Dim>
void launch(TensorInfo<const T, IndexType> input,
            IndexType inputSliceSize,
            IndexType outputSliceSize,
            bool largest,
            IndexType numInputSlices,
            IndexType inputWithinSliceStride,
            TensorInfo<T, IndexType> topK,
            IndexType topKWithinSliceStride,
            TensorInfo<int64_t, IndexType> indices,
            IndexType indicesWithinSliceStride,
            gpuStream_t stream) {
  dim3 grid;
  bool ok = getGridFromTiles(numInputSlices, &grid);
  assert(ok);
  (void)ok;
  int warp_size = TOPK_WARP_SIZE;
  dim3 block(
      std::min(topk_ceil_div((int64_t)inputSliceSize, (int64_t)warp_size) *
                   (int64_t)warp_size,
               (int64_t)1024));
  gatherTopK<T, IndexType, Dim, /*WithKthValues=*/false>
      <<<grid, block, 0, stream>>>(input,
                                   inputSliceSize,
                                   outputSliceSize,
                                   largest,
                                   numInputSlices,
                                   inputWithinSliceStride,
                                   topK,
                                   topKWithinSliceStride,
                                   indices,
                                   indicesWithinSliceStride,
                                   nullptr);
}

}  // namespace sbtopk

namespace mbtopk {  // multi_block_topk

constexpr int BLOCK_THREADS = 256;
constexpr int RADIX_BITS = 8;
constexpr int RADIX_DIGITS = 1 << RADIX_BITS;  // 256
constexpr int RADIX_MASK = (RADIX_DIGITS - 1);
static_assert(
    RADIX_DIGITS <= BLOCK_THREADS,
    "radixFindKthValues kernel requires RADIX_DIGITS <= BLOCK_THREADS");
constexpr int MIN_ITEMS_PER_THREAD = 4;
constexpr int MAX_ITEMS_PER_THREAD = 64;

template <typename T, typename IndexType>
__global__ void fill(T* x, T value, IndexType size) {
  IndexType idx =
      static_cast<IndexType>(blockIdx.x) * static_cast<IndexType>(blockDim.x) +
      static_cast<IndexType>(threadIdx.x);
  for (IndexType i = idx; i < size; i += static_cast<IndexType>(gridDim.x) *
                                         static_cast<IndexType>(blockDim.x)) {
    x[i] = value;
  }
}

template <typename T, typename IndexType, typename Bitwise, int Dim>
__global__ void __launch_bounds__(BLOCK_THREADS)
    radixFindKthValues(TensorInfo<const T, IndexType> input,
                       uint32_t slice_size,
                       uint32_t* ks_to_find,
                       uint32_t num_slices,
                       IndexType withinSliceStride,
                       int current_bit,
                       int items_per_thread,
                       uint32_t blocks_per_slice,
                       Bitwise desiredMask,
                       Bitwise* desires,
                       int16_t* counts) {
  int items_per_block = items_per_thread * BLOCK_THREADS;
  int tidx = threadIdx.x;
  uint32_t block_idx = getLinearBlockId<uint32_t>();
  uint32_t slice_idx = block_idx / blocks_per_slice;
  uint32_t blk_idx_in_slice = block_idx % blocks_per_slice;
  if (slice_idx >= num_slices) {
    return;
  }

  Bitwise desired = desires[slice_idx];
  IndexType slice_start_index =
      IndexToOffset<const T, IndexType, Dim>::get(slice_idx, input);
  const T* data = &input.data[slice_start_index];

  static_assert(MAX_ITEMS_PER_THREAD * BLOCK_THREADS <
                    std::numeric_limits<int16_t>::max(),
                "blockwise counter too large");
  union __align__(16) TempStorage {
    uint32_t digit_counters[RADIX_DIGITS];
  };
  __shared__ TempStorage temp_storage;

  if (tidx < RADIX_DIGITS) {
    temp_storage.digit_counters[tidx] = 0;
  }
  __syncthreads();

  items_per_thread =
      (blk_idx_in_slice + 1 < blocks_per_slice)
          ? items_per_thread
          : topk_ceil_div(
                (int64_t)(slice_size - blk_idx_in_slice * items_per_block),
                (int64_t)BLOCK_THREADS);

  for (int i = 0; i < items_per_thread; ++i) {
    IndexType idx =
        blk_idx_in_slice * items_per_block + i * BLOCK_THREADS + tidx;
    if (idx < slice_size) {
      idx *= withinSliceStride;
      Bitwise val = TopKTypeConfig<T>::convert(doLdg(&data[idx]));
      bool has_val = ((val & desiredMask) == (desired & desiredMask));
      Bitwise digit =
          Bitfield<Bitwise>::getBitfield(val, current_bit, RADIX_BITS);
      if (has_val) {
        atomicAdd(&temp_storage.digit_counters[digit], 1);
      }
    }
  }

  __syncthreads();

  static_assert(RADIX_DIGITS <= BLOCK_THREADS,
                "this kernel requires RADIX_DIGITS <= BLOCK_THREADS");
  uint32_t digit_count = 0;
  if (tidx < RADIX_DIGITS) {
    digit_count = temp_storage.digit_counters[tidx];
  }

  if (tidx < RADIX_DIGITS) {
    counts[block_idx * RADIX_DIGITS + tidx] = digit_count;
  }
}

template <typename Bitwise, typename T>
__global__ void __launch_bounds__(RADIX_DIGITS)
    computeBlockwiseWithinKCounts(Bitwise* desires_in,
                                  int16_t* counts,
                                  uint32_t* ks_to_find_in,
                                  uint32_t blocks_per_slice,
                                  int current_bit,
                                  bool largest,
                                  uint32_t* withinKCounts,
                                  T* kthValues,
                                  uint32_t* ks_to_find_out,
                                  Bitwise* desires_out,
                                  uint32_t num_blocks) {
  int tidx = threadIdx.x;
  uint32_t block_idx = getLinearBlockId<uint32_t>();
  uint32_t slice_idx = block_idx / blocks_per_slice;

  if (block_idx >= num_blocks) {
    return;
  }

  typedef cub::BlockScan<uint32_t, RADIX_DIGITS> BlockScan;
  union __align__(16) TempStorage {
    uint32_t digit_count_cumsum[RADIX_DIGITS];
    typename BlockScan::TempStorage scan_storage;
  };
  __shared__ TempStorage temp_storage;

  uint32_t digit_count = 0;
  if (tidx < RADIX_DIGITS) {
    for (uint32_t blk = 0; blk < blocks_per_slice; ++blk) {
      digit_count +=
          counts[(slice_idx * blocks_per_slice + blk) * RADIX_DIGITS + tidx];
    }
  }

  uint32_t digit_count_cumsum;
  BlockScan(temp_storage.scan_storage)
      .InclusiveSum(digit_count, digit_count_cumsum);
  __syncthreads();
  if (tidx < RADIX_DIGITS) {
    temp_storage.digit_count_cumsum[tidx] = digit_count_cumsum;
  }
  __syncthreads();

  __shared__ Bitwise desired;
  uint32_t k_to_find = ks_to_find_in[slice_idx];

  if (tidx < RADIX_DIGITS) {
    uint32_t digit_count_cumsum_left =
        (tidx == 0) ? 0 : temp_storage.digit_count_cumsum[tidx - 1];

    if (digit_count_cumsum_left < k_to_find &&
        k_to_find <= digit_count_cumsum) {
      desired = desires_in[slice_idx];
      desired = Bitfield<Bitwise>::setBitfield(
          desired, tidx, current_bit, RADIX_BITS);
      if (block_idx == slice_idx * blocks_per_slice) {
        desires_out[slice_idx] = desired;
        if (current_bit > 0) {
          ks_to_find_out[slice_idx] = k_to_find - digit_count_cumsum_left;
        } else {
          kthValues[slice_idx] = TopKTypeConfig<T>::deconvert(desired);
        }
      }
    }
  }
  __syncthreads();

#if !TOPK_CUB_SUPPORTS_SCAN_BY_KEY()
  return;
#endif

  Bitwise desired_digit =
      Bitfield<Bitwise>::getBitfield(desired, current_bit, RADIX_BITS);

  bool warp_is_active, thread_is_active;
  int warp = tidx / TOPK_WARP_SIZE;
  if (largest) {
    int end_of_warp = warp * TOPK_WARP_SIZE + TOPK_WARP_SIZE - 1;
    warp_is_active = end_of_warp > static_cast<int>(desired_digit);
    thread_is_active = tidx > static_cast<int>(desired_digit);
  } else {
    int start_of_warp = warp * TOPK_WARP_SIZE;
    warp_is_active = start_of_warp < static_cast<int>(desired_digit);
    thread_is_active = tidx < static_cast<int>(desired_digit);
  }
  uint32_t count = 0;
  if (warp_is_active) {
    if (thread_is_active) {
      count = doLdg(counts + block_idx * RADIX_DIGITS + tidx);
    }
    for (int offset = TOPK_WARP_SIZE / 2; offset > 0; offset /= 2) {
      count += TOPK_WARP_SHFL_DOWN(count, offset);
    }
  }

  constexpr int num_warps = RADIX_DIGITS / TOPK_WARP_SIZE;
  __shared__ uint32_t warp_counts[num_warps];
  if (tidx % TOPK_WARP_SIZE == 0) {
    warp_counts[warp] = count;
  }
  __syncthreads();
#ifdef __HIPCC__
  assert(RADIX_DIGITS < TOPK_WARP_SIZE * TOPK_WARP_SIZE);
#else
  static_assert(RADIX_DIGITS < TOPK_WARP_SIZE * TOPK_WARP_SIZE,
                "Assuming only 1 warp is needed for final reduction");
#endif
  if (warp != 0) {
    return;
  }
  count = 0;
  if (tidx < num_warps) {
    count = warp_counts[tidx];
  }
  for (int offset = num_warps / 2; offset > 0; offset /= 2) {
    count += TOPK_WARP_SHFL_DOWN(count, offset);
  }
  if (tidx == 0) {
    withinKCounts[block_idx] += count;
  }
}

#if TOPK_CUB_SUPPORTS_SCAN_BY_KEY()
template <typename Bitwise>
__global__ void computeBlockwiseKthCounts(Bitwise* desires,
                                          int16_t* counts,
                                          uint32_t num_blocks,
                                          uint32_t blocks_per_slice,
                                          uint32_t* kthCounts) {
  TOPK_CUDA_KERNEL_LOOP_TYPE(idx, num_blocks, uint32_t) {
    uint32_t slice_idx = idx / blocks_per_slice;
    Bitwise desired = doLdg(desires + slice_idx);
    Bitwise desired_digit =
        Bitfield<Bitwise>::getBitfield(desired, 0, RADIX_BITS);
    kthCounts[idx] = doLdg(counts + idx * RADIX_DIGITS + desired_digit);
  }
}

template <typename T, typename IndexType, int Dim>
__global__ void __launch_bounds__(BLOCK_THREADS)
    gatherTopK(TensorInfo<const T, IndexType> input,
               IndexType inputSliceSize,
               IndexType outputSliceSize,
               bool largest,
               uint32_t numInputSlices,
               IndexType inputWithinSliceStride,
               TensorInfo<T, IndexType> topK,
               IndexType topKWithinSliceStride,
               TensorInfo<int64_t, IndexType> indices,
               IndexType indicesWithinSliceStride,
               uint32_t items_per_thread,
               uint32_t blocks_per_slice,
               T* kthValues,
               uint32_t* withinKCounts,
               uint32_t* kthCounts,
               uint32_t num_blocks) {
  uint32_t items_per_block = items_per_thread * BLOCK_THREADS;
  uint32_t tidx = threadIdx.x;
  uint32_t block_idx = getLinearBlockId<uint32_t>();

  if (block_idx >= num_blocks) {
    return;
  }

  uint32_t slice_idx = block_idx / blocks_per_slice;
  uint32_t blk_idx_in_slice = block_idx % blocks_per_slice;

  items_per_thread =
      (blk_idx_in_slice + 1 < blocks_per_slice)
          ? items_per_thread
          : topk_ceil_div(
                (int64_t)(inputSliceSize - blk_idx_in_slice * items_per_block),
                (int64_t)BLOCK_THREADS);

  IndexType sliceStartIndex =
      IndexToOffset<const T, IndexType, Dim>::get(slice_idx, input);
  IndexType topKSliceStartIndex =
      IndexToOffset<T, IndexType, Dim>::get(slice_idx, topK);
  IndexType indicesSliceStartIndex =
      IndexToOffset<int64_t, IndexType, Dim>::get(slice_idx, indices);

  const T* inputSliceStart = &input.data[sliceStartIndex];
  T* topKSliceStart = &topK.data[topKSliceStartIndex];
  int64_t* indicesSliceStart = &indices.data[indicesSliceStartIndex];

  T kthValue = kthValues[slice_idx];
  const auto kthValueConverted = TopKTypeConfig<T>::convert(kthValue);

  uint32_t startWithinK = 0;
  if (blk_idx_in_slice > 0) {
    startWithinK = withinKCounts[block_idx - 1];
  }
  uint32_t startKth =
      withinKCounts[slice_idx * blocks_per_slice + blocks_per_slice - 1];
  if (blk_idx_in_slice > 0) {
    startKth += kthCounts[block_idx - 1];
  }

  typedef cub::BlockScan<uint32_t, BLOCK_THREADS> BlockScan;
  __shared__ typename BlockScan::TempStorage temp_storage;
  for (uint32_t i = 0; i < items_per_thread; ++i) {
    IndexType idx =
        blk_idx_in_slice * items_per_block + i * BLOCK_THREADS + tidx;
    T val;
    int withinK = 0;
    int kth = 0;
    if (idx < inputSliceSize) {
      val = doLdg(inputSliceStart + idx * inputWithinSliceStride);
      const auto valConverted = TopKTypeConfig<T>::convert(val);
      withinK = (largest ? valConverted > kthValueConverted
                         : valConverted < kthValueConverted);
      kth = (valConverted == kthValueConverted);
    }

    uint32_t withinKIndex;
    uint32_t numWithinK;
    BlockScan(temp_storage).ExclusiveSum(withinK, withinKIndex, numWithinK);
    __syncthreads();
    if (withinK) {
      uint32_t offset = withinKIndex + startWithinK;
      topKSliceStart[offset * topKWithinSliceStride] = val;
      indicesSliceStart[offset * indicesWithinSliceStride] = idx;
    }
    startWithinK += numWithinK;

    if (startKth < outputSliceSize) {
      uint32_t kthIndex;
      uint32_t numKth;
      BlockScan(temp_storage).ExclusiveSum(kth, kthIndex, numKth);
      __syncthreads();
      if (kth) {
        uint32_t offset = kthIndex + startKth;
        if (offset < outputSliceSize) {
          topKSliceStart[offset * topKWithinSliceStride] = val;
          indicesSliceStart[offset * indicesWithinSliceStride] = idx;
        }
      }
      startKth += numKth;
    }
  }
}
#endif  // TOPK_CUB_SUPPORTS_SCAN_BY_KEY

// get_items_per_thread: compute optimal items per thread based on GPU occupancy
int get_items_per_thread(uint64_t num_slices,
                         uint64_t slice_size,
                         int device_id) {
  constexpr int REGS_PER_THREAD = 40;
  constexpr int REGS_PER_BLOCK = REGS_PER_THREAD * BLOCK_THREADS;
  const auto& prop = phi::backends::gpu::GetDeviceProperties(device_id);
  int mpc = prop.multiProcessorCount;
#ifdef PADDLE_WITH_HIP
  // HIP/DCU: hipDeviceProp_t lacks regsPerMultiprocessor and
  // maxBlocksPerMultiProcessor. Use conservative defaults:
  // 65536 registers per CU is typical for AMD GCN/CDNA architectures.
  // maxThreadsPerMultiProcessor / BLOCK_THREADS as blocks_per_mp estimate.
  int regs_per_mp = 65536;
  int max_blocks_per_mp = prop.maxThreadsPerMultiProcessor / BLOCK_THREADS;
#else
  int regs_per_mp = prop.regsPerMultiprocessor;
  int max_blocks_per_mp = prop.maxBlocksPerMultiProcessor;
#endif
  int blocks_per_mp = std::min(regs_per_mp / REGS_PER_BLOCK, max_blocks_per_mp);
  int64_t items_per_thread =
      topk_ceil_div((int64_t)(slice_size * num_slices),
                    (int64_t)(mpc * blocks_per_mp * BLOCK_THREADS));
  items_per_thread = std::max(
      MIN_ITEMS_PER_THREAD,
      std::min(static_cast<int>(items_per_thread), MAX_ITEMS_PER_THREAD));
  return items_per_thread;
}

class BlockIdxToKey {
  uint32_t blocks_per_slice;

 public:
  explicit BlockIdxToKey(uint32_t blocks_per_slice)
      : blocks_per_slice(blocks_per_slice) {}
  __device__ __forceinline__ uint32_t operator()(uint32_t blk) const {
    return blk / blocks_per_slice;
  }
};

template <typename T, typename IndexType, int Dim>
void launch(TensorInfo<const T, IndexType> input,
            IndexType inputSliceSize,
            IndexType outputSliceSize,
            bool largest,
            uint32_t numInputSlices,
            IndexType inputWithinSliceStride,
            TensorInfo<T, IndexType> topK,
            IndexType topKWithinSliceStride,
            TensorInfo<int64_t, IndexType> indices,
            IndexType indicesWithinSliceStride,
            gpuStream_t stream,
            int device_id,
            const phi::Place& place) {
  int items_per_thread =
      get_items_per_thread(numInputSlices, inputSliceSize, device_id);
  int items_per_block = items_per_thread * BLOCK_THREADS;

  using Bitwise = typename TopKTypeConfig<T>::RadixType;
  uint32_t blocks_per_slice =
      topk_ceil_div((int64_t)inputSliceSize, (int64_t)items_per_block);
  uint32_t num_blocks = numInputSlices * blocks_per_slice;

  // Temporary storage allocation using phi::memory_utils
  auto phi_stream = phi::Stream(reinterpret_cast<phi::StreamId>(stream));

  auto kthValues_buffer =
      phi::memory_utils::Alloc(place, numInputSlices * sizeof(T), phi_stream);
  T* kthValues = reinterpret_cast<T*>(kthValues_buffer->ptr());

  auto semaphores_buffer = phi::memory_utils::Alloc(
      place, numInputSlices * sizeof(uint32_t), phi_stream);
  uint32_t* semaphores = reinterpret_cast<uint32_t*>(semaphores_buffer->ptr());
#ifdef PADDLE_WITH_HIP
  hipMemsetAsync(semaphores, 0, numInputSlices * sizeof(uint32_t), stream);
#else
  cudaMemsetAsync(semaphores, 0, numInputSlices * sizeof(uint32_t), stream);
#endif

  auto ks_to_find_buffer = phi::memory_utils::Alloc(
      place, 2 * numInputSlices * sizeof(uint32_t), phi_stream);
  uint32_t* ks_to_find = reinterpret_cast<uint32_t*>(ks_to_find_buffer->ptr());
  uint32_t k_to_find =
      largest ? inputSliceSize - outputSliceSize + 1 : outputSliceSize;
  fill<uint32_t>
      <<<std::min(((int64_t)numInputSlices + 511) / 512, (int64_t)1073741824),
         512,
         0,
         stream>>>(ks_to_find, k_to_find, numInputSlices);

  auto desired_buffer = phi::memory_utils::Alloc(
      place, 2 * numInputSlices * sizeof(Bitwise), phi_stream);
  Bitwise* desired = reinterpret_cast<Bitwise*>(desired_buffer->ptr());

  auto counts_buffer = phi::memory_utils::Alloc(
      place, num_blocks * RADIX_DIGITS * sizeof(int16_t), phi_stream);
  int16_t* counts = reinterpret_cast<int16_t*>(counts_buffer->ptr());
  static_assert(MAX_ITEMS_PER_THREAD * BLOCK_THREADS <
                    std::numeric_limits<int16_t>::max(),
                "blockwise counter too large");

#if TOPK_CUB_SUPPORTS_SCAN_BY_KEY()
  auto withinKCounts_buffer = phi::memory_utils::Alloc(
      place, num_blocks * sizeof(uint32_t), phi_stream);
  uint32_t* withinKCounts =
      reinterpret_cast<uint32_t*>(withinKCounts_buffer->ptr());
#ifdef PADDLE_WITH_HIP
  hipMemsetAsync(withinKCounts, 0, num_blocks * sizeof(uint32_t), stream);
#else
  cudaMemsetAsync(withinKCounts, 0, num_blocks * sizeof(uint32_t), stream);
#endif

  auto kthCounts_buffer = phi::memory_utils::Alloc(
      place, num_blocks * sizeof(uint32_t), phi_stream);
  uint32_t* kthCounts = reinterpret_cast<uint32_t*>(kthCounts_buffer->ptr());
#else
  uint32_t* withinKCounts = nullptr;
#endif

  Bitwise desiredMask = 0;
  dim3 grid;
  bool ok = getGridFromTiles(num_blocks, &grid);
  assert(ok);
  (void)ok;
  dim3 block(BLOCK_THREADS);

  uint32_t* ks_to_find_in = ks_to_find;
  uint32_t* ks_to_find_out = ks_to_find + numInputSlices;
  Bitwise* desired_in = desired;
  Bitwise* desired_out = desired + numInputSlices;

  for (int current_bit = sizeof(T) * 8 - RADIX_BITS; current_bit >= 0;
       current_bit -= RADIX_BITS) {
    radixFindKthValues<T, IndexType, Bitwise, Dim>
        <<<grid, block, 0, stream>>>(input,
                                     inputSliceSize,
                                     ks_to_find_in,
                                     numInputSlices,
                                     inputWithinSliceStride,
                                     current_bit,
                                     items_per_thread,
                                     blocks_per_slice,
                                     desiredMask,
                                     desired_in,
                                     counts);

    computeBlockwiseWithinKCounts<Bitwise, T>
        <<<grid, RADIX_DIGITS, 0, stream>>>(desired_in,
                                            counts,
                                            ks_to_find_in,
                                            blocks_per_slice,
                                            current_bit,
                                            largest,
                                            withinKCounts,
                                            kthValues,
                                            ks_to_find_out,
                                            desired_out,
                                            num_blocks);

    auto tmp_desired = desired_in;
    desired_in = desired_out;
    desired_out = tmp_desired;
    auto tmp_ks = ks_to_find_in;
    ks_to_find_in = ks_to_find_out;
    ks_to_find_out = tmp_ks;
    // Host-side equivalent of Bitfield<Bitwise>::setBitfield(desiredMask,
    // RADIX_MASK, current_bit, RADIX_BITS) Cannot use Bitfield::setBitfield
    // here because it's __device__-only (uses PTX asm)
    {
      Bitwise mask = ((Bitwise(1) << RADIX_BITS) - 1) << current_bit;
      desiredMask =
          (desiredMask & ~mask) | ((Bitwise(RADIX_MASK) << current_bit) & mask);
    }
  }
  desired = desired_in;

#if TOPK_CUB_SUPPORTS_SCAN_BY_KEY()
  computeBlockwiseKthCounts<Bitwise>
      <<<std::min(((int64_t)numInputSlices + 255) / 256, (int64_t)1073741824),
         256,
         0,
         stream>>>(desired, counts, num_blocks, blocks_per_slice, kthCounts);

  // Use cub::DeviceScan::InclusiveSumByKey
  using counting_iter_t = cub::CountingInputIterator<uint32_t>;
  using slice_idx_iter_t =
      cub::TransformInputIterator<uint32_t, BlockIdxToKey, counting_iter_t>;
  slice_idx_iter_t slice_idx_iter(counting_iter_t(0),
                                  BlockIdxToKey(blocks_per_slice));

  // InclusiveSumByKey for withinKCounts
  {
    size_t temp_storage_bytes = 0;
    cub::DeviceScan::InclusiveSumByKey(nullptr,
                                       temp_storage_bytes,
                                       slice_idx_iter,
                                       withinKCounts,
                                       withinKCounts,
                                       num_blocks,
                                       cub::Equality(),
                                       stream);
    auto temp_buf =
        phi::memory_utils::Alloc(place, temp_storage_bytes, phi_stream);
    cub::DeviceScan::InclusiveSumByKey(temp_buf->ptr(),
                                       temp_storage_bytes,
                                       slice_idx_iter,
                                       withinKCounts,
                                       withinKCounts,
                                       num_blocks,
                                       cub::Equality(),
                                       stream);
  }
  // InclusiveSumByKey for kthCounts
  {
    size_t temp_storage_bytes = 0;
    cub::DeviceScan::InclusiveSumByKey(nullptr,
                                       temp_storage_bytes,
                                       slice_idx_iter,
                                       kthCounts,
                                       kthCounts,
                                       num_blocks,
                                       cub::Equality(),
                                       stream);
    auto temp_buf =
        phi::memory_utils::Alloc(place, temp_storage_bytes, phi_stream);
    cub::DeviceScan::InclusiveSumByKey(temp_buf->ptr(),
                                       temp_storage_bytes,
                                       slice_idx_iter,
                                       kthCounts,
                                       kthCounts,
                                       num_blocks,
                                       cub::Equality(),
                                       stream);
  }

  gatherTopK<T, IndexType, Dim>
      <<<grid, block, 0, stream>>>(input,
                                   inputSliceSize,
                                   outputSliceSize,
                                   largest,
                                   numInputSlices,
                                   inputWithinSliceStride,
                                   topK,
                                   topKWithinSliceStride,
                                   indices,
                                   indicesWithinSliceStride,
                                   items_per_thread,
                                   blocks_per_slice,
                                   kthValues,
                                   withinKCounts,
                                   kthCounts,
                                   num_blocks);
#else
  // Fallback: use single-block gatherTopK with kthValues
  {
    dim3 grid2;
    bool ok2 = getGridFromTiles(numInputSlices, &grid2);
    assert(ok2);
    (void)ok2;
    int warp_size = TOPK_WARP_SIZE;
    dim3 block2(
        std::min(topk_ceil_div((int64_t)inputSliceSize, (int64_t)warp_size) *
                     (int64_t)warp_size,
                 (int64_t)1024));
    sbtopk::gatherTopK<T, IndexType, Dim, /*WithKthValues=*/true>
        <<<grid2, block2, 0, stream>>>(input,
                                       inputSliceSize,
                                       outputSliceSize,
                                       largest,
                                       numInputSlices,
                                       inputWithinSliceStride,
                                       topK,
                                       topKWithinSliceStride,
                                       indices,
                                       indicesWithinSliceStride,
                                       kthValues);
  }
#endif
}

}  // namespace mbtopk

bool should_use_multiblock(int64_t num_slices, int64_t slice_size) {
  if (num_slices > std::numeric_limits<uint32_t>::max() ||
      slice_size > std::numeric_limits<uint32_t>::max())
    return false;
#if TOPK_CUB_SUPPORTS_SCAN_BY_KEY()
  return (num_slices <= 20 && slice_size >= 20000) ||
         (num_slices > 20 && num_slices <= 40 && slice_size >= 10000) ||
         (num_slices > 40 && num_slices <= 80 && slice_size >= 8000) ||
         (num_slices > 80 && num_slices < 200 && slice_size >= 5000) ||
         (num_slices >= 200 && num_slices < 800 && slice_size >= 3000) ||
         (num_slices >= 800 && num_slices <= 4000 && slice_size >= 800) ||
         (num_slices > 4000 && slice_size >= 400);
#else
  return (num_slices <= 400 && slice_size >= 5000) ||
         (num_slices > 400 && num_slices < 4000 && slice_size >= 1000) ||
         (num_slices >= 4000 && slice_size >= 300);
#endif
}

// canUse32BitIndexMath: check if tensor indexing fits in 32-bit integers
bool canUse32BitIndexMath(
    const phi::DenseTensor& t,
    int64_t max_elem = std::numeric_limits<int32_t>::max()) {
  int64_t elements = t.numel();
  if (elements >= max_elem) {
    return false;
  }
  if (elements == 0) {
    return max_elem > 0;
  }

  int64_t offset = 0;
  int64_t linearId = elements - 1;

  for (int i = t.dims().size() - 1; i >= 0; --i) {
    int64_t curDimIndex = linearId % t.dims()[i];
    int64_t curDimOffset = curDimIndex * t.strides()[i];
    offset += curDimOffset;
    linearId /= t.dims()[i];
  }

  if (offset >= max_elem) {
    return false;
  }

  return true;
}

}  // namespace topk_impl
#endif  // PADDLE_PHI_KERNELS_FUNCS_TOP_K_CUDA_KERNEL_H_