"""Optimized top-k over the last dim of a 2D fp32 tensor.

Strategy (v6: radix-select in registers + split main/tail kernels)
-------------------------------------------------------------------
* Keys: each element becomes a 64-bit key = (monotonic-float-bits << 32) | index.
  Key order == (value desc, index), all comparisons are unsigned.
* Main kernel: each thread holds a small chunk of keys in registers (<= 8
  elements loaded as float4) plus alive/won bitmasks. Radix passes (top byte
  of the value bits down) narrow the candidate set: shared-memory 256-bin
  histogram (smem atomics, double-buffered so zeroing overlaps with
  filtering), warp-0 scan picks the pivot digit and updates k, every thread
  classifies its elements (won / alive / dead) in registers. Passes stop
  early as soon as the alive set has exactly the demanded size (it then fully
  belongs to the top-k). Winners are compacted; for single-block rows a
  single warp bitonic-sorts the k winners via shuffles and writes the
  descending output; multi-block rows publish their K winners to scratch.
* Tail kernel (multi-block rows only, separate launch so no cross-block
  fences/counters are needed -- stream ordering gives visibility): one block
  per row runs the same radix-select over the per-block candidate lists.
* k == 1 uses a dedicated argmax-reduction kernel.
* Overhead reduction: per Model instance we preallocate outputs/scratch and
  capture the kernel launch(es) into a CUDA graph keyed by the input pointer.
  Steady-state forwards are a single graph replay; first sighting of a new
  input pointer captures (and immediately replays) a graph for it.
"""
from __future__ import annotations

from pathlib import Path

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

OP_TYPE = "topk"
SUPPORTED_PRECISIONS = ["fp32"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

_CPP = """
#include <torch/extension.h>
void topk_run(torch::Tensor x, torch::Tensor out_v, torch::Tensor out_i,
              torch::Tensor scratch, int64_t cfg);
"""

_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define DEV_INLINE __device__ __forceinline__
#define FULL 0xffffffffu
typedef unsigned long long u64;

// ---- key packing -----------------------------------------------------------
DEV_INLINE unsigned f2u(float v) {
  unsigned b = __float_as_uint(v);
  return b ^ ((b & 0x80000000u) ? 0xFFFFFFFFu : 0x80000000u);
}
DEV_INLINE float u2f(unsigned u) {
  unsigned b = (u & 0x80000000u) ? (u ^ 0x80000000u) : ~u;
  return __uint_as_float(b);
}
DEV_INLINE u64 make_key(float v, int idx) {
  return ((u64)f2u(v) << 32) | (unsigned)idx;
}

// ---- warp scan: find pivot digit for the current histogram -----------------
// hist[256] counts alive elements per digit; want largest digit d* with
// count(> d*) < k <= count(>= d*). Writes d* and new k (k - count(> d*)).
// If the alive total is <= k (early-exit case), writes d* = -1 (everything
// alive then belongs to the winners).
DEV_INLINE void pick_pivot(const int* hist, int k, int* smem_d, int* smem_k, int lane) {
  const int base = 255 - 8 * lane;  // lane covers bins base..base-7 (high->low)
  int local[8];
  int s = 0;
  #pragma unroll
  for (int i = 0; i < 8; i++) {
    local[i] = hist[base - i];
    s += local[i];
  }
  // exclusive prefix of bins strictly above my block (lower lane = higher bins)
  int t = s;
  #pragma unroll
  for (int off = 1; off < 32; off <<= 1) {
    int v = __shfl_up_sync(FULL, t, off);
    if (lane >= off) t += v;
  }
  int total = __shfl_sync(FULL, t, 31);
  if (total <= k) {
    if (lane == 0) {
      *smem_d = -1;
      *smem_k = k;
    }
    return;
  }
  int cum = t - s;  // count above my block
  int dstar = 0, g = 0;
  bool found = false;
  #pragma unroll
  for (int i = 0; i < 8; i++) {
    // crossing bin: running count reaches k here (cum < k ensures it has not
    // already been reached in a higher bin)
    if (!found && cum < k && cum + local[i] >= k) {
      dstar = base - i;
      g = cum;
      found = true;
    }
    cum += local[i];
  }
  // exactly one lane found it; OR-reduce packed result
  unsigned pack = found ? ((unsigned)dstar | ((unsigned)g << 8)) : 0u;
  #pragma unroll
  for (int off = 16; off > 0; off >>= 1) pack |= __shfl_xor_sync(FULL, pack, off);
  if (lane == 0) {
    *smem_d = (int)(pack & 0xFFu);
    *smem_k = k - (int)(pack >> 8);
  }
}

// Warp bitonic sort of the MP keys held across the warp's lanes (E = MP/32
// keys per lane, logical element i = lane*E + e), then emit the first K.
template <int MP, int K>
DEV_INLINE void warp_sort_regs(u64* el, int lane, float* out_v,
                               int64_t* out_i, size_t base) {
  constexpr int E = MP / 32;
  for (int s = 2; s <= MP; s <<= 1) {
    for (int d = s >> 1; d > 0; d >>= 1) {
      // snapshot: all exchanges of a round must see pre-round values
      u64 old[E];
      #pragma unroll
      for (int e = 0; e < E; e++) old[e] = el[e];
      #pragma unroll
      for (int e = 0; e < E; e++) {
        int i = lane * E + e;
        int j = i ^ d;
        bool desc = ((i & s) == 0);
        bool lower = ((i & d) == 0);
        bool want_max = (desc == lower);
        u64 mine = old[e];
        u64 other;
        int jl = j / E, je = j - jl * E;
        if (jl == lane) {
          other = old[je];
        } else {
          unsigned lo = __shfl_sync(FULL, (unsigned)old[je], jl);
          unsigned hi = __shfl_sync(FULL, (unsigned)(old[je] >> 32), jl);
          other = ((u64)hi << 32) | lo;
        }
        bool take_other = want_max ? (other > mine) : (other < mine);
        el[e] = take_other ? other : mine;
      }
    }
  }
  #pragma unroll
  for (int e = 0; e < E; e++) {
    int i = lane * E + e;
    if (i < K) {
      out_v[base + i] = u2f((unsigned)(el[e] >> 32));
      out_i[base + i] = (int64_t)(unsigned)el[e];
    }
  }
}

// Warp-0 bitonic sort of wout[0..K) (descending), then emit values+indices.
// K < 32 sorts inside a zero-padded 32-element network.
template <int K>
DEV_INLINE void warp_sort_emit(const u64* wout, int lane, float* out_v,
                               int64_t* out_i, size_t base) {
  constexpr int MP = (K < 32) ? 32 : K;
  constexpr int E = MP / 32;
  u64 el[E];
  #pragma unroll
  for (int e = 0; e < E; e++) {
    int i = lane * E + e;
    el[e] = (i < K) ? wout[i] : 0ull;
  }
  warp_sort_regs<MP, K>(el, lane, out_v, out_i, base);
}

// ---------------------------------------------------------------------------
// Radix-select over register-held keys: four 8-bit passes over the value
// bits (shifts 56, 48, 40, 32). Two ping-pong 256-bin histograms so zeroing
// overlaps with the filter phase. Stops early (dstar < 0) once the alive
// count is <= the demanded k.
// ---------------------------------------------------------------------------
template <int THREADS, int ASZ, int COUNT, int K>
DEV_INLINE void radix_select(u64* A, unsigned& alive, unsigned& won, int& kcur,
                             int* hist0, int* hist1, int* smem_d, int* smem_k,
                             int tid, int lane, int warp) {
  #pragma unroll 1
  for (int p = 0; p < 4; p++) {
    const int shift = 56 - 8 * p;
    int* hist = (p & 1) ? hist1 : hist0;
    int* other = (p & 1) ? hist0 : hist1;
    #pragma unroll
    for (int j = 0; j < ASZ; j++) {
      if (j < COUNT && (alive & (1u << j))) {
        atomicAdd(&hist[(unsigned)(A[j] >> shift) & 0xFFu], 1);
      }
    }
    __syncthreads();
    if (warp == 0) pick_pivot(hist, kcur, smem_d, smem_k, lane);
    // zero the other histogram for the next pass while pivot results land
    for (int i = tid; i < 256; i += THREADS) other[i] = 0;
    __syncthreads();
    const int dstar = *smem_d;
    if (dstar < 0) break;  // alive set already within the required size
    kcur = *smem_k;
    #pragma unroll
    for (int j = 0; j < ASZ; j++) {
      if (j < COUNT && (alive & (1u << j))) {
        int byte = (int)((unsigned)(A[j] >> shift) & 0xFFu);
        if (byte > dstar) {
          alive &= ~(1u << j);
          won |= (1u << j);
        } else if (byte < dstar) {
          alive &= ~(1u << j);
        }
      }
    }
  }
}

// Compact winners into wout: won elements first (slots 0..nwon), then alive.
// nwon == K - kcur is known to every thread.
template <int ASZ, int COUNT, int K>
DEV_INLINE void compact_winners(const u64* A, unsigned alive, unsigned won,
                                int kcur, u64* wout, int* cnt_w, int* cnt_a) {
  const int base_alive = K - kcur;
  #pragma unroll
  for (int j = 0; j < ASZ; j++) {
    if (j < COUNT) {
      unsigned bit = 1u << j;
      if (won & bit) {
        int pos = atomicAdd(cnt_w, 1);
        wout[pos] = A[j];
      } else if (alive & bit) {
        int pos = base_alive + atomicAdd(cnt_a, 1);
        if (pos < K) wout[pos] = A[j];
      }
    }
  }
}

// ---------------------------------------------------------------------------
// Main kernel: per-block radix top-k. BPR == 1 emits the row directly;
// otherwise publishes K winner keys per block to scratch (stream ordering to
// the follow-up tail kernel provides visibility; no fences needed).
// ---------------------------------------------------------------------------
template <int THREADS, int BPR, int B, int K>
__global__ void __launch_bounds__(THREADS) topk_main(
    const float* __restrict__ x, int n,
    float* __restrict__ out_v, int64_t* __restrict__ out_i,
    u64* __restrict__ scratch) {
  constexpr int ASZ = B;
  extern __shared__ char smem_raw[];
  int* hist0 = reinterpret_cast<int*>(smem_raw);
  int* hist1 = hist0 + 256;
  int* smem_d = hist1 + 256;
  int* smem_k = smem_d + 1;
  int* cnt_w = smem_k + 1;
  int* cnt_a = cnt_w + 1;
  u64* wout = reinterpret_cast<u64*>(cnt_a + 1);

  const int tid = threadIdx.x;
  const int row = blockIdx.x / BPR;
  const int sub = blockIdx.x - row * BPR;
  const int lane = tid & 31;
  const int warp = tid >> 5;

  u64 A[ASZ];
  unsigned alive = 0, won = 0;
  {
    constexpr int TPR = THREADS * BPR;
    const float4* x4 = reinterpret_cast<const float4*>(x + (size_t)row * n);
    const int n4 = n >> 2;
    const int gtid = sub * THREADS + tid;
    #pragma unroll
    for (int j = 0; j < B / 4; j++) {
      int e = gtid + j * TPR;
      if (e < n4) {
        float4 f = x4[e];
        int gi = e << 2;
        A[j * 4 + 0] = make_key(f.x, gi + 0);
        A[j * 4 + 1] = make_key(f.y, gi + 1);
        A[j * 4 + 2] = make_key(f.z, gi + 2);
        A[j * 4 + 3] = make_key(f.w, gi + 3);
        alive |= (0xFu << (j * 4));
      } else {
        A[j * 4 + 0] = 0; A[j * 4 + 1] = 0;
        A[j * 4 + 2] = 0; A[j * 4 + 3] = 0;
      }
    }
  }
  // overlap smem init with the global loads above
  if (tid < 256) {
    hist0[tid] = 0;
    hist1[tid] = 0;
  }
  if (tid == 0) {
    *cnt_w = 0;
    *cnt_a = 0;
  }
  __syncthreads();

  int kcur = K;
  radix_select<THREADS, ASZ, B, K>(A, alive, won, kcur, hist0, hist1,
                                   smem_d, smem_k, tid, lane, warp);
  __syncthreads();
  compact_winners<ASZ, B, K>(A, alive, won, kcur, wout, cnt_w, cnt_a);
  __syncthreads();

  if (BPR == 1) {
    if (warp == 0) warp_sort_emit<K>(wout, lane, out_v, out_i, (size_t)row * K);
  } else {
    if (tid < K) scratch[((size_t)row * BPR + sub) * K + tid] = wout[tid];
  }
}

// ---------------------------------------------------------------------------
// Small tail kernel (M = BPR*K <= 256): warp-0 sorts all candidates directly
// (zero-padded to the next power of two MP) and emits the first K.
// ---------------------------------------------------------------------------
template <int BPR, int K>
__global__ void topk_tail_small(
    const u64* __restrict__ scratch,
    float* __restrict__ out_v, int64_t* __restrict__ out_i) {
  constexpr int M = BPR * K;
  constexpr int MP =
      (M <= 32) ? 32 : (M <= 64) ? 64 : (M <= 128) ? 128 : 256;
  constexpr int E = MP / 32;
  const int row = blockIdx.x;
  const int lane = threadIdx.x & 31;
  if ((threadIdx.x >> 5) != 0) return;
  u64 keys[E];
  #pragma unroll
  for (int e = 0; e < E; e++) {
    int i = lane * E + e;
    keys[e] = (i < M) ? scratch[(size_t)row * M + i] : 0ull;
  }
  warp_sort_regs<MP, K>(keys, lane, out_v, out_i, (size_t)row * K);
}

// ---------------------------------------------------------------------------
// Tail kernel: one block per row merges the BPR per-block winner lists.
// Launched after topk_main on the same stream (visibility guaranteed).
// ---------------------------------------------------------------------------
template <int THREADS, int BPR, int K>
__global__ void __launch_bounds__(THREADS) topk_tail(
    const u64* __restrict__ scratch,
    float* __restrict__ out_v, int64_t* __restrict__ out_i) {
  constexpr int M = BPR * K;
  constexpr int C2 = (M + THREADS - 1) / THREADS;
  extern __shared__ char smem_raw[];
  int* hist0 = reinterpret_cast<int*>(smem_raw);
  int* hist1 = hist0 + 256;
  int* smem_d = hist1 + 256;
  int* smem_k = smem_d + 1;
  int* cnt_w = smem_k + 1;
  int* cnt_a = cnt_w + 1;
  u64* wout = reinterpret_cast<u64*>(cnt_a + 1);
  u64* stage = wout + K;

  const int tid = threadIdx.x;
  const int row = blockIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;

  for (int i = tid; i < M; i += THREADS) stage[i] = scratch[(size_t)row * M + i];
  if (tid < 256) {
    hist0[tid] = 0;
    hist1[tid] = 0;
  }
  if (tid == 0) {
    *cnt_w = 0;
    *cnt_a = 0;
  }
  __syncthreads();

  u64 A[C2];
  unsigned alive = 0, won = 0;
  #pragma unroll
  for (int j = 0; j < C2; j++) {
    int idx = tid * C2 + j;
    if (idx < M) {
      A[j] = stage[idx];
      alive |= (1u << j);
    } else {
      A[j] = 0;
    }
  }

  int kcur = K;
  radix_select<THREADS, C2, C2, K>(A, alive, won, kcur, hist0, hist1,
                                   smem_d, smem_k, tid, lane, warp);
  __syncthreads();
  compact_winners<C2, C2, K>(A, alive, won, kcur, wout, cnt_w, cnt_a);
  __syncthreads();

  if (warp == 0) warp_sort_emit<K>(wout, lane, out_v, out_i, (size_t)row * K);
}

// ---------------------------------------------------------------------------
// Argmax kernel for k == 1 (single block per row).
// ---------------------------------------------------------------------------
template <int THREADS>
__global__ void __launch_bounds__(THREADS) topk_argmax(
    const float* __restrict__ x, int n,
    float* __restrict__ out_v, int64_t* __restrict__ out_i) {
  const int row = blockIdx.x;

  const float4* x4 = reinterpret_cast<const float4*>(x + (size_t)row * n);
  const int n4 = n >> 2;

  float bval = -INFINITY;
  int bidx = 0;
  for (int e = threadIdx.x; e < n4; e += THREADS) {
    float4 f = x4[e];
    int gi = e << 2;
    if (f.x > bval) { bval = f.x; bidx = gi + 0; }
    if (f.y > bval) { bval = f.y; bidx = gi + 1; }
    if (f.z > bval) { bval = f.z; bidx = gi + 2; }
    if (f.w > bval) { bval = f.w; bidx = gi + 3; }
  }

  #pragma unroll
  for (int off = 16; off > 0; off >>= 1) {
    float ov = __shfl_xor_sync(FULL, bval, off);
    int oi = __shfl_xor_sync(FULL, bidx, off);
    if (ov > bval || (ov == bval && oi < bidx)) { bval = ov; bidx = oi; }
  }

  constexpr int NW = THREADS / 32;
  __shared__ float wv[NW];
  __shared__ int wi[NW];
  if ((threadIdx.x & 31) == 0) {
    int w = threadIdx.x >> 5;
    wv[w] = bval;
    wi[w] = bidx;
  }
  __syncthreads();

  if (threadIdx.x == 0) {
    float bv = wv[0]; int bi = wi[0];
    #pragma unroll
    for (int w = 1; w < NW; w++) {
      if (wv[w] > bv || (wv[w] == bv && wi[w] < bi)) { bv = wv[w]; bi = wi[w]; }
    }
    out_v[row] = bv;
    out_i[row] = (int64_t)bi;
  }
}

// ---------------------------------------------------------------------------
// Host dispatch
// ---------------------------------------------------------------------------
template <int THREADS, int BPR, int B, int K>
static size_t main_smem() {
  return (512 + 4) * sizeof(int) + (size_t)K * sizeof(u64);
}
template <int THREADS, int BPR, int K>
static size_t tail_smem() {
  return (512 + 4) * sizeof(int) + (size_t)K * sizeof(u64) +
         (size_t)(BPR * K) * sizeof(u64);
}

template <typename Kfn>
static void ensure_smem(Kfn fn, size_t smem) {
  static bool done = false;
  if (!done && smem > 48 * 1024) {
    cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
    done = true;
  }
}

void topk_run(torch::Tensor x, torch::Tensor out_v, torch::Tensor out_i,
              torch::Tensor scratch, int64_t cfg) {
  const int batch = x.size(0);
  const int n = x.size(1);
  const float* xp = x.data_ptr<float>();
  float* ovp = out_v.data_ptr<float>();
  int64_t* oip = out_i.data_ptr<int64_t>();
  u64* sp = scratch.numel() ? reinterpret_cast<u64*>(scratch.data_ptr<int64_t>()) : nullptr;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  switch (cfg) {
    case 0: {  // (1, 131072, 64): 16 blocks x 1024 threads, 8 elems/thread
      size_t sm = main_smem<1024, 16, 8, 64>();
      topk_main<1024, 16, 8, 64><<<batch * 16, 1024, sm, stream>>>(xp, n, ovp, oip, sp);
      size_t st = tail_smem<1024, 16, 64>();
      ensure_smem(topk_tail<1024, 16, 64>, st);
      topk_tail<1024, 16, 64><<<batch, 1024, st, stream>>>(sp, ovp, oip);
      break;
    }
    case 1: {  // (64, 8192, 8): one block of 1024 threads per row, 8 elems/thread
      size_t sm = main_smem<1024, 1, 8, 8>();
      topk_main<1024, 1, 8, 8><<<batch, 1024, sm, stream>>>(xp, n, ovp, oip, sp);
      break;
    }
    case 2: {  // (32, 16384, 32): 8 blocks x 512 threads per row
      size_t sm = main_smem<512, 8, 4, 32>();
      topk_main<512, 8, 4, 32><<<batch * 8, 512, sm, stream>>>(xp, n, ovp, oip, sp);
      size_t st = tail_smem<512, 8, 32>();
      topk_tail<512, 8, 32><<<batch, 512, st, stream>>>(sp, ovp, oip);
      break;
    }
    case 3: {  // (16, 12000, 16): 3 blocks x 512 threads per row, 8 elems/thread
      size_t sm = main_smem<512, 3, 8, 16>();
      topk_main<512, 3, 8, 16><<<batch * 3, 512, sm, stream>>>(xp, n, ovp, oip, sp);
      topk_tail_small<3, 16><<<batch, 32, 0, stream>>>(sp, ovp, oip);
      break;
    }
    case 4:  // (128, 4096, 1)
      topk_argmax<512><<<batch, 512, 0, stream>>>(xp, n, ovp, oip);
      break;
    default:
      TORCH_CHECK(false, "topk_run: unknown config ", cfg);
  }
}
"""

_BUILD_DIR = Path(__file__).resolve().parent / "_topk_build"
_BUILD_DIR.mkdir(parents=True, exist_ok=True)
_ext = load_inline(
    name="topk_bitonic_ext",
    cpp_sources=[_CPP],
    cuda_sources=[_CUDA],
    functions=["topk_run"],
    build_directory=str(_BUILD_DIR),
    verbose=False,
    extra_cuda_cflags=[
        "-O3",
        "-std=c++17",
        "-gencode=arch=compute_90a,code=sm_90a",
        "--use_fast_math",
    ],
)

# (batch, n, k) -> (cfg id, blocks_per_row)
_CONFIGS = {
    (1, 131072, 64): (0, 16),
    (64, 8192, 8): (1, 1),
    (32, 16384, 32): (2, 8),
    (16, 12000, 16): (3, 3),
    (128, 4096, 1): (4, 1),
}


class Model(nn.Module):
    """Top-k over the last dim of a 2D tensor via custom CUDA kernels."""

    def __init__(self, batch: int, n: int, k: int):
        super().__init__()
        self.batch, self.n, self.k = batch, n, k
        # Match reference.py so load_state_dict(strict=True) works.
        self.register_buffer("_dummy", torch.zeros(1))

        key = (batch, n, k)
        if key not in _CONFIGS:
            raise RuntimeError(f"unsupported shape: {key}")
        self._cfg, bpr = _CONFIGS[key]

        dev = torch.device("cuda")
        self._out_v = torch.empty(batch, k, dtype=torch.float32, device=dev)
        self._out_i = torch.empty(batch, k, dtype=torch.int64, device=dev)
        if bpr > 1:
            self._scratch = torch.empty(batch * bpr * k, dtype=torch.int64, device=dev)
        else:
            self._scratch = torch.empty(0, dtype=torch.int64, device=dev)

        # Input pointer -> captured CUDA graph; _replay is the hot path.
        self._graphs: dict[int, torch.cuda.CUDAGraph] = {}
        self._replay = None
        self._ptr = None
        self._ret = (self._out_v, self._out_i)

    def _run(self, x: torch.Tensor) -> None:
        _ext.topk_run(x, self._out_v, self._out_i, self._scratch, self._cfg)

    def forward(self, x: torch.Tensor):
        r = self._replay
        if r is not None and x.data_ptr() == self._ptr:
            r()
            return self._ret
        return self._slow_forward(x)

    def _slow_forward(self, x: torch.Tensor):
        if not x.is_contiguous():
            x = x.contiguous()
        ptr = x.data_ptr()
        g = self._graphs.get(ptr)
        if g is None:
            if x.shape[0] != self.batch or x.shape[1] != self.n:
                raise RuntimeError(f"unexpected input shape {tuple(x.shape)}")
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g):
                self._run(x)
            self._graphs[ptr] = g
        g.replay()
        # Arm the fast path on the second sighting of the same pointer; a
        # different pointer disarms it (the cached replay is pointer-bound).
        if ptr == self._ptr:
            self._replay = g.replay
        else:
            self._replay = None
        self._ptr = ptr
        return self._ret


# Module-level shims mirroring reference.py.
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]