#include <stdint.h>
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
namespace cg = cooperative_groups;

__device__ __inline__ uint64_t make_key(float prob, int64_t col) {
    return (((uint64_t)__float_as_uint(prob)) << 32) |
           (uint64_t)(0xffffffffu - (uint32_t)col);
}

__device__ __inline__ float block_reduce_sum(float val) {
    cg::thread_block block = cg::this_thread_block();
    cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);

    // 第一步：warp 内归约
    val = cg::reduce(tile, val, cg::plus<float>());

    __shared__ float warp_sums[32];
    const int warp_id = threadIdx.x >> 5;
    const int lane_id = threadIdx.x & 31;
    const int num_warps = blockDim.x >> 5;

    if (lane_id == 0) warp_sums[warp_id] = val;
    __syncthreads();

    // 第二步：第一个 warp 做二次归约
    if (warp_id == 0) {
        val = (lane_id < num_warps) ? warp_sums[lane_id] : 0.0f;
        val = cg::reduce(tile, val, cg::plus<float>());
        if (lane_id == 0) warp_sums[0] = val;
    }
    __syncthreads();

    return warp_sums[0];   // 广播给所有线程
}

// ============================== 主 kernel ==============================
__global__ void topp(const float* __restrict__ probs,
                     const float* __restrict__ top_p,
                     float* __restrict__ renorm_probs,
                     int64_t batch_size,
                     int64_t num_classes) {
    int tid = threadIdx.x;

    __shared__ uint64_t mask_share;
    __shared__ uint64_t current_mask_share;
    __shared__ float   accept_sum;

    if (tid == 0) {
        mask_share = 0ull;
        current_mask_share = 0ull;
        accept_sum = 0.0f;
    }
    __syncthreads();

    const int64_t row_offset = (int64_t)blockIdx.x * num_classes;
    const float row_top_p = top_p[blockIdx.x];

    // ---------- 1. 从高位到低位寻找分界点 ----------
    for (int bit = 63; bit >= 0; --bit) {
        const uint64_t current_mask = current_mask_share;
        const uint64_t mask = mask_share;
        const uint64_t now_flag = 1ull << bit;

        float local_sum = 0.0f;
        for (int64_t col = tid; col < num_classes; col += blockDim.x) {
            float p = probs[row_offset + col];
            uint64_t key = make_key(p, col);
            if (((key & mask) == current_mask) && ((key & now_flag) != 0)) {
                local_sum += p;
            }
        }
        const float total_sum = block_reduce_sum(local_sum);

        if (tid == 0) {
            if (total_sum > 0.0f && accept_sum + total_sum <= row_top_p) {
                accept_sum += total_sum;              // 包含该位
            } else if (total_sum > 0.0f) {
                current_mask_share |= now_flag;       // 分界点在此位
            }
            mask_share |= now_flag;
        }
        __syncthreads();
    }

    // ---------- 2. 重归一化 ----------
    const uint64_t threshold = current_mask_share;

    float final_sum = 0.0f;
    for (int64_t col = tid; col < num_classes; col += blockDim.x) {
        float p = probs[row_offset + col];
        uint64_t key = make_key(p, col);
        if (key >= threshold) final_sum += p;
    }
    final_sum = block_reduce_sum(final_sum);

    const float inv = (final_sum > 0.0f) ? (1.0f / final_sum) : 0.0f;
    for (int64_t col = tid; col < num_classes; col += blockDim.x) {
        float p = probs[row_offset + col];
        uint64_t key = make_key(p, col);
        renorm_probs[row_offset + col] = (key >= threshold) ? (p * inv) : 0.0f;
    }
}

extern "C" void run_kernel(const float* probs,
                           const float* top_p,
                           float* renorm_probs,
                           int64_t batch_size,
                           int64_t num_classes) {
    topp<<<batch_size, 1024>>>(probs, top_p, renorm_probs, batch_size, num_classes);
}