"""Single-pass top-k on RTX PRO 6000 (SM120 Blackwell).

Top-k over the last dim of a 2D fp32 tensor. Memory-bound regime: the input
read dominates, so the win over the builtin selector is a SINGLE kernel launch that
streams each row exactly once with vectorized (float4) coalesced loads, keeps a
per-row top-k in registers, and merges in shared memory. No multi-pass radix
select, no full sort.

One thread block per row. Each thread grabs a strided slice of the row, keeps a
sorted (descending) local top-K in registers, then all threads' local top-Ks are
k-way merged by thread 0 via a max-heap in shared memory.
"""
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Declarations for the pybind module (the definitions live in the .cu file).
# A single tensor-in entry point; no raw-pointer binding from Python ints.
_CPP_DECL = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
void topk_cuda(torch::Tensor x, int k, int B, torch::Tensor values, torch::Tensor indices);
"""

# ---------------------------------------------------------------------------
# C++/CUDA source. Templated on K so nvcc unrolls the per-thread top-k arrays.
# ---------------------------------------------------------------------------
_CPP_SRC = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cstdint>

__device__ __forceinline__ void kinsert(float* __restrict__ lval,
                                        int* __restrict__ lidx,
                                        int K, float v, int idx) {
    if (v > lval[K - 1]) {
        int pos = K - 1;
        while (pos > 0 && v > lval[pos - 1]) {
            lval[pos]     = lval[pos - 1];
            lidx[pos]     = lidx[pos - 1];
            --pos;
        }
        lval[pos] = v;
        lidx[pos] = idx;
    }
}

template <int K>
__global__ void topk_kernel(const float* __restrict__ input,
                            float* __restrict__ out_values,
                            long long* __restrict__ out_indices,
                            int n, int batch, int B) {
    const int row = blockIdx.x;
    if (row >= batch) return;
    const float* __restrict__ rowdata = input + (long long)row * n;
    const int tid = threadIdx.x;

    // Per-thread local top-K, sorted descending. lval[0] = largest.
    float lval[K];
    int   lidx[K];
    for (int i = 0; i < K; ++i) { lval[i] = -1e30f; lidx[i] = -1; }

    // Vectorized float4 strided loads (coalesced across the block).
    const int n4 = n >> 2;
    const float4* __restrict__ rowvec = reinterpret_cast<const float4*>(rowdata);
    for (int i = tid; i < n4; i += B) {
        float4 vv = rowvec[i];
        int base = i << 2;
        kinsert(lval, lidx, K, vv.x, base);
        kinsert(lval, lidx, K, vv.y, base + 1);
        kinsert(lval, lidx, K, vv.z, base + 2);
        kinsert(lval, lidx, K, vv.w, base + 3);
    }
    // Scalar tail (n not multiple of 4).
    const int base = n4 << 2;
    for (int i = base + tid; i < n; i += B) {
        kinsert(lval, lidx, K, rowdata[i], i);
    }

    // Write local top-K into shared memory: B lists of K, each sorted desc.
    extern __shared__ char smem[];
    float* __restrict__ sval = reinterpret_cast<float*>(smem);
    int*   __restrict__ sidx = reinterpret_cast<int*>(smem + (size_t)B * K * sizeof(float));
    for (int i = 0; i < K; ++i) {
        sval[tid * K + i] = lval[i];
        sidx[tid * K + i] = lidx[i];
    }
    __syncthreads();

    // k-way merge: thread 0 merges the B sorted lists with a max-heap.
    if (tid == 0) {
        const int MAXB = 256;
        float hv[MAXB];
        int   hl[MAXB];   // list id
        int   hp[MAXB];   // position within list
        int   hsize = 0;

        for (int j = 0; j < B; ++j) {
            int pos = hsize++;
            hv[pos] = sval[j * K];
            hl[pos] = j;
            hp[pos] = 0;
            while (pos > 0) {
                int par = (pos - 1) >> 1;
                if (hv[par] < hv[pos]) {
                    float t = hv[par]; hv[par] = hv[pos]; hv[pos] = t;
                    int ti = hl[par]; hl[par] = hl[pos]; hl[pos] = ti;
                    ti = hp[par]; hp[par] = hp[pos]; hp[pos] = ti;
                    pos = par;
                } else break;
            }
        }

        float* __restrict__ ov = out_values + (long long)row * K;
        long long* __restrict__ oi = out_indices + (long long)row * K;
        for (int out = 0; out < K; ++out) {
            float bv = hv[0];
            int bl = hl[0], bp = hp[0];
            ov[out] = bv;
            oi[out] = (long long)sidx[(long long)bl * K + bp];
            int np = bp + 1;
            if (np < K) { hv[0] = sval[(long long)bl * K + np]; hp[0] = np; }
            else        { hv[0] = -1e30f; hp[0] = K; }
            int pos = 0;
            while (true) {
                int l = 2 * pos + 1, r = 2 * pos + 2, lg = pos;
                if (l < hsize && hv[l] > hv[lg]) lg = l;
                if (r < hsize && hv[r] > hv[lg]) lg = r;
                if (lg != pos) {
                    float t = hv[pos]; hv[pos] = hv[lg]; hv[lg] = t;
                    int ti = hl[pos]; hl[pos] = hl[lg]; hl[lg] = ti;
                    ti = hp[pos]; hp[pos] = hp[lg]; hp[lg] = ti;
                    pos = lg;
                } else break;
            }
        }
    }
}

// Single tensor-in dispatcher. Switch on k to pick the K-templated kernel.
void topk_cuda(torch::Tensor x, int k, int B, torch::Tensor values, torch::Tensor indices) {
    const float* xi = x.data_ptr<float>();
    float*       vo = values.data_ptr<float>();
    auto*        io = indices.data_ptr<int64_t>();
    const int rows = x.size(0);   // one block per row
    const int n    = x.size(1);   // row length, streamed by B threads
    cudaStream_t s = c10::cuda::getCurrentCUDAStream();
    int sm;
    switch (k) {
        case 1:  sm = B*1*8;  topk_kernel<1> <<<rows,B,sm,s>>>(xi,vo,io,n,rows,B);  break;
        case 8:  sm = B*8*8;  topk_kernel<8> <<<rows,B,sm,s>>>(xi,vo,io,n,rows,B);  break;
        case 16: sm = B*16*8; topk_kernel<16><<<rows,B,sm,s>>>(xi,vo,io,n,rows,B);  break;
        case 32: sm = B*32*8; topk_kernel<32><<<rows,B,sm,s>>>(xi,vo,io,n,rows,B);  break;
        case 64: sm = B*64*8; topk_kernel<64><<<rows,B,sm,s>>>(xi,vo,io,n,rows,B);  break;
    }
}
"""

_cuda_ext = None


def _get_ext():
    global _cuda_ext
    if _cuda_ext is None:
        _cuda_ext = load_inline(
            name="topk_bitonic",
            cpp_sources=_CPP_DECL,     # declarations so pybind sees the dispatchers
            cuda_sources=_CPP_SRC,
            functions=["topk_cuda"],
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "-gencode", "arch=compute_120,code=sm_120",
                "-maxrregcount=128",
            ],
            verbose=False,
        )
    return _cuda_ext


# Block size (threads per row). 64 keeps shared mem within the 48 KB default
# carveout for all K<=64 (64*64*8 = 32 KB).
_BLOCK_SIZE = 64


class Model(nn.Module):
    def __init__(self, batch: int, n: int, k: int):
        super().__init__()
        self.batch, self.n, self.k = batch, n, k
        self.register_buffer("_dummy", torch.zeros(1))
        if k not in {1, 8, 16, 32, 64}:
            raise ValueError(f"k={k} not in supported {{1,8,16,32,64}}")
        _get_ext()  # compile eagerly so first forward isn't penalised

    def forward(self, x: torch.Tensor):
        assert x.dim() == 2 and x.dtype == torch.float32
        batch, n, k = self.batch, self.n, self.k
        values = torch.empty((batch, k), dtype=torch.float32, device=x.device)
        indices = torch.empty((batch, k), dtype=torch.int64, device=x.device)
        _get_ext().topk_cuda(x, k, _BLOCK_SIZE, values, indices)
        return values, indices


# Module-level shims rebuilt by check.py / benchmark.py per shape.
batch = 64
n = 8192
k = 8


def get_inputs():
    x = torch.randn(batch, n, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [batch, n, k]