"""Custom CUDA top-k kernel for H100 PCIe (SM90).

Memory-bound: the input read dominates. Strategy:
  - Partial kernel: a grid of (B blocks/row x batch) tiles. Each block's threads
    cooperatively scan a tile, each maintaining a sorted top-K buffer in shared
    memory (filtered by a register-resident running minimum so most elements cost
    one register compare). A ping-pong merge tree in shared memory reduces the
    per-thread buffers to the block's top-K.
  - Reduce kernel (only when B>1): merges the B per-tile top-K buffers per row
    into the final sorted top-k.

K = next power of two >= k (all evaluated k are already powers of two). Values
are copied verbatim from the input (never recomputed), so numeric error is ~0.
"""
import os

os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "9.0")

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_CPP_DECL = r"""
#include <torch/extension.h>
#include <vector>
std::vector<at::Tensor> topk_forward(at::Tensor x, int64_t k, int64_t B,
                                     int64_t THREADS, int64_t RTHREADS);
"""

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

static __device__ __constant__ float NEG_INF_C = -INFINITY;

template <int K>
__device__ __forceinline__ void insert_sorted(float* val, int64_t* idx,
                                              float x, int64_t gidx) {
    // val[0..K-1] sorted DESCENDING (val[0] max, val[K-1] min). Insert x if it
    // beats the current min, shifting smaller elements down.
    if (x > val[K - 1]) {
        int p = K - 1;
        while (p > 0 && val[p - 1] < x) {
            val[p] = val[p - 1];
            idx[p] = idx[p - 1];
            --p;
        }
        val[p] = x;
        idx[p] = gidx;
    }
}

// Ping-pong merge tree. On entry each thread tid has written its K-length sorted
// buffer to sv[tid*K..], si[tid*K..]. After the tree the block's top-K (sorted
// desc) lives in sv[0..K-1], si[0..K-1].
template <int K>
__device__ __forceinline__ void reduce_tree(float* sv, int64_t* si, int THREADS) {
    float* bufs[2]   = {sv, sv + (size_t)THREADS * K};
    int64_t* bufi[2] = {si, si + (size_t)THREADS * K};
    int src = 0;
    int tid = threadIdx.x;
    for (int stride = THREADS / 2; stride > 0; stride >>= 1) {
        int dst = 1 - src;
        if (tid < stride) {
            float* A  = bufs[src] + (size_t)tid * K;
            float* B  = bufs[src] + (size_t)(tid + stride) * K;
            float* O  = bufs[dst] + (size_t)tid * K;
            int64_t* Ai = bufi[src] + (size_t)tid * K;
            int64_t* Bi = bufi[src] + (size_t)(tid + stride) * K;
            int64_t* Oi = bufi[dst] + (size_t)tid * K;
            int i = 0, j = 0;
            for (int t = 0; t < K; ++t) {
                bool takeA = (j >= K) || (i < K && A[i] >= B[j]);
                if (takeA) { O[t] = A[i]; Oi[t] = Ai[i]; ++i; }
                else       { O[t] = B[j]; Oi[t] = Bi[j]; ++j; }
            }
        }
        __syncthreads();
        src = dst;
    }
    if (src == 1 && tid < K) {
        sv[tid] = bufs[1][tid];
        si[tid] = bufi[1][tid];
    }
    __syncthreads();
}

// Partial kernel: top-K of one tile -> scratch (or output when B==1 & K==k).
template <int K>
__global__ void partial_topk(const float* __restrict__ x, int n,
                             int elts_per_block, int B,
                             float* __restrict__ sval,
                             int64_t* __restrict__ sidx) {
    int row = blockIdx.y;
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int THREADS = blockDim.x;
    const float* xrow = x + (int64_t)row * n;
    int tile_start = b * elts_per_block;
    int tile_end = min(tile_start + elts_per_block, n);

    extern __shared__ char smem_raw[];
    float* sv = reinterpret_cast<float*>(smem_raw);
    int64_t* si = reinterpret_cast<int64_t*>(sv + 2 * (size_t)THREADS * K);
    float* myval = sv + (size_t)tid * K;
    int64_t* myidx = si + (size_t)tid * K;
    #pragma unroll
    for (int t = 0; t < K; ++t) { myval[t] = -INFINITY; myidx[t] = 0; }

    float cur_min = -INFINITY;
    for (int gi = tile_start + tid; gi < tile_end; gi += THREADS) {
        float v = xrow[gi];
        if (v > cur_min) {
            insert_sorted<K>(myval, myidx, v, (int64_t)gi);
            cur_min = myval[K - 1];
        }
    }
    __syncthreads();
    reduce_tree<K>(sv, si, THREADS);

    int out_off = (row * B + b) * K;
    if (tid < K) {
        sval[out_off + tid] = sv[tid];
        sidx[out_off + tid] = si[tid];
    }
}

// Reduce kernel: merge B per-tile top-K buffers per row -> final top-k.
template <int K>
__global__ void reduce_topk(const float* __restrict__ sval,
                            const int64_t* __restrict__ sidx, int B, int k,
                            float* __restrict__ out_val,
                            int64_t* __restrict__ out_idx) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int THREADS = blockDim.x;
    int total = B * K;
    const float* rv = sval + (int64_t)row * total;
    const int64_t* ri = sidx + (int64_t)row * total;

    extern __shared__ char smem_raw[];
    float* sv = reinterpret_cast<float*>(smem_raw);
    int64_t* si = reinterpret_cast<int64_t*>(sv + 2 * (size_t)THREADS * K);
    float* myval = sv + (size_t)tid * K;
    int64_t* myidx = si + (size_t)tid * K;
    #pragma unroll
    for (int t = 0; t < K; ++t) { myval[t] = -INFINITY; myidx[t] = 0; }

    float cur_min = -INFINITY;
    for (int i = tid; i < total; i += THREADS) {
        float v = rv[i];
        if (v > cur_min) {
            insert_sorted<K>(myval, myidx, v, ri[i]);
            cur_min = myval[K - 1];
        }
    }
    __syncthreads();
    reduce_tree<K>(sv, si, THREADS);

    if (tid < k) {
        out_val[row * k + tid] = sv[tid];
        out_idx[row * k + tid] = si[tid];
    }
}

template <int K>
static void topk_impl(const float* xp, int batch, int n, int k, int B,
                      int THREADS, int RTHREADS, at::Tensor& out_val,
                      at::Tensor& out_idx, const at::TensorOptions& opts_f,
                      const at::TensorOptions& opts_l) {
    int elts_per_block = (n + B - 1) / B;
    size_t bytes_pair = (size_t)2 * THREADS * K;
    size_t smem_p = sizeof(float) * bytes_pair + sizeof(int64_t) * bytes_pair;
    size_t smem_r = sizeof(float) * (size_t)2 * RTHREADS * K +
                    sizeof(int64_t) * (size_t)2 * RTHREADS * K;
    cudaFuncSetAttribute((void*)partial_topk<K>,
                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                         (int)smem_p);
    bool direct = (B == 1) && (K == k);
    if (direct) {
        dim3 grid(B, batch);
        partial_topk<K><<<grid, THREADS, smem_p>>>(
            xp, n, elts_per_block, B,
            out_val.data_ptr<float>(), out_idx.data_ptr<int64_t>());
    } else {
        auto sval = at::empty({batch, B, K}, opts_f);
        auto sidx = at::empty({batch, B, K}, opts_l);
        dim3 grid(B, batch);
        partial_topk<K><<<grid, THREADS, smem_p>>>(
            xp, n, elts_per_block, B,
            sval.data_ptr<float>(), sidx.data_ptr<int64_t>());
        cudaFuncSetAttribute((void*)reduce_topk<K>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)smem_r);
        reduce_topk<K><<<batch, RTHREADS, smem_r>>>(
            sval.data_ptr<float>(), sidx.data_ptr<int64_t>(), B, k,
            out_val.data_ptr<float>(), out_idx.data_ptr<int64_t>());
    }
}

std::vector<at::Tensor> topk_forward(at::Tensor x, int64_t k, int64_t B,
                                     int64_t THREADS, int64_t RTHREADS) {
    TORCH_CHECK(x.dim() == 2, "expected 2D");
    TORCH_CHECK(x.is_cuda(), "expected cuda");
    TORCH_CHECK(x.is_contiguous(), "expected contiguous");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "expected fp32");
    int batch = x.size(0);
    int n = x.size(1);
    int K = 1;
    while (K < (int)k) K <<= 1;
    auto opts_f = x.options();
    auto opts_l = x.options().dtype(at::kLong);
    auto out_val = at::empty({batch, k}, opts_f);
    auto out_idx = at::empty({batch, k}, opts_l);
    const float* xp = x.data_ptr<float>();
    switch (K) {
        case 1:   topk_impl<1>  (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 2:   topk_impl<2>  (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 4:   topk_impl<4>  (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 8:   topk_impl<8>  (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 16:  topk_impl<16> (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 32:  topk_impl<32> (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 64:  topk_impl<64> (xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        case 128: topk_impl<128>(xp, batch, n, (int)k, (int)B, (int)THREADS, (int)RTHREADS, out_val, out_idx, opts_f, opts_l); break;
        default: TORCH_CHECK(false, "K too large");
    }
    return {out_val, out_idx};
}
"""

_ext = load_inline(
    name="kbh_topk_ext",
    cpp_sources=[_CPP_DECL],
    cuda_sources=[_CUDA_SRC],
    functions=["topk_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_90,code=sm_90"],
    verbose=False,
)


def _next_pow2(x):
    k = 1
    while k < x:
        k <<= 1
    return k


def _choose_config(batch, n, k):
    """Pick (blocks-per-row B, partial THREADS, reduce THREADS)."""
    K = _next_pow2(k)
    # Threads: smem budget 2*THREADS*K*12 <= ~200KB. K>=64 -> 128 threads.
    if K >= 64:
        threads = 128
    else:
        threads = 256
    rthreads = 128 if K >= 64 else 256
    # Target enough total blocks to fill ~132 SMs with some slack.
    target = 132
    if batch >= 100:
        B = 1
    else:
        B = max(1, (target + batch - 1) // batch)
    B = min(B, n)  # never more blocks than elements
    return B, threads, rthreads


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))
        self.B, self.THREADS, self.RTHREADS = _choose_config(batch, n, k)

    def forward(self, x: torch.Tensor):
        x = x.contiguous()
        v, i = _ext.topk_forward(x, self.k, self.B, self.THREADS, self.RTHREADS)
        return v, i


# 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]