"""Custom top-k for B200 (SM100). Two-kernel bitonic:
  stage1: each block reads a power-of-2 tile of a row, bitonic-sorts it,
          writes the top-k (descending) to a temp buffer.
  stage2: one block per row merges (num_tiles * k) candidates via bitonic sort,
          writes the final top-k with int64 indices.
Internal indices are int32 (n <= 131072); converted to int64 on output.
"""
from __future__ import annotations
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"]

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

// Bitonic sort descending in shared memory. N must be a power of 2.
// v: values, idx: companion indices. Threads tid in [0,T).
__device__ __forceinline__ void bitonic_desc(float* __restrict__ v, int* __restrict__ idx,
                                             int N, int tid, int T) {
  for (int k = 2; k <= N; k <<= 1) {
    for (int j = k >> 1; j > 0; j >>= 1) {
      for (int i = tid; i < N; i += T) {
        int ij = i ^ j;
        if (ij > i) {
          bool asc = (i & k) == 0;
          float a = v[i], b = v[ij];
          bool sw = asc ? (a < b) : (a > b);   // descending overall
          if (sw) {
            v[i] = b; v[ij] = a;
            int t = idx[i]; idx[i] = idx[ij]; idx[ij] = t;
          }
        }
      }
      __syncthreads();
    }
  }
}

// stage1: grid (num_tiles, batch), block T. tile = power of 2.
__global__ void topk_stage1(const float* __restrict__ X,
                            float* __restrict__ OV, int* __restrict__ OI,
                            int n, int k, int num_tiles, int tile) {
  int row = blockIdx.y;
  int t = blockIdx.x;
  int tid = threadIdx.x;
  int T = blockDim.x;
  extern __shared__ char sm[];
  float* v = (float*)sm;
  int* idx = (int*)(v + tile);
  int start = t * tile;
  int cnt = (start + tile <= n) ? tile : (n - start);
  if (cnt < 0) cnt = 0;
  const float* Xrow = X + (size_t)row * n;
  for (int i = tid; i < tile; i += T) {
    if (i < cnt) { v[i] = Xrow[start + i]; idx[i] = start + i; }
    else { v[i] = -INFINITY; idx[i] = 0; }
  }
  __syncthreads();
  bitonic_desc(v, idx, tile, tid, T);
  float* OVrow = OV + ((size_t)row * num_tiles + t) * k;
  int*   OIrow = OI + ((size_t)row * num_tiles + t) * k;
  for (int i = tid; i < k; i += T) {
    OVrow[i] = v[i];
    OIrow[i] = idx[i];
  }
}

// stage2: grid (batch,), block T. merges num_tiles*k candidates.
__global__ void topk_stage2(const float* __restrict__ V, const int* __restrict__ I,
                            float* __restrict__ OV, int64_t* __restrict__ OI,
                            int num_tiles, int k, int padN) {
  int row = blockIdx.x;
  int tid = threadIdx.x;
  int T = blockDim.x;
  int total = num_tiles * k;
  extern __shared__ char sm[];
  float* v = (float*)sm;
  int* idx = (int*)(v + padN);
  const float* Vrow = V + (size_t)row * total;
  const int*   Irow = I + (size_t)row * total;
  for (int i = tid; i < padN; i += T) {
    if (i < total) { v[i] = Vrow[i]; idx[i] = Irow[i]; }
    else { v[i] = -INFINITY; idx[i] = 0; }
  }
  __syncthreads();
  bitonic_desc(v, idx, padN, tid, T);
  float* OVrow = OV + (size_t)row * k;
  int64_t* OIrow = OI + (size_t)row * k;
  for (int i = tid; i < k; i += T) {
    OVrow[i] = v[i];
    OIrow[i] = (int64_t)idx[i];
  }
}

void topk_launch(const torch::Tensor& x, torch::Tensor& tmp_v, torch::Tensor& tmp_i,
                 torch::Tensor& out_v, torch::Tensor& out_i,
                 int64_t k, int64_t tile, int64_t threads1, int64_t threads2) {
  int batch = x.size(0);
  int n = x.size(1);
  int num_tiles = (n + tile - 1) / tile;
  const float* Xp = x.data_ptr<float>();
  float* TV = tmp_v.data_ptr<float>();
  int*   TI = tmp_i.data_ptr<int>();
  float* OV = out_v.data_ptr<float>();
  int64_t* OI = out_i.data_ptr<int64_t>();

  // stage1
  {
    dim3 grid(num_tiles, batch);
    dim3 block(threads1);
    size_t smem = (size_t)tile * (sizeof(float) + sizeof(int));
    if (smem > 49152u)
      cudaFuncSetAttribute(topk_stage1, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
    topk_stage1<<<grid, block, smem>>>(Xp, TV, TI, n, (int)k, num_tiles, (int)tile);
  }
  // stage2
  {
    int total = num_tiles * (int)k;
    int padN = 1; while (padN < total) padN <<= 1;
    dim3 grid(batch);
    dim3 block(threads2);
    size_t smem = (size_t)padN * (sizeof(float) + sizeof(int));
    if (smem > 49152u)
      cudaFuncSetAttribute(topk_stage2, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
    topk_stage2<<<grid, block, smem>>>(TV, TI, OV, OI, num_tiles, (int)k, padN);
  }
}
"""

_CPP_SRC = r"""
void topk_launch(const torch::Tensor& x, torch::Tensor& tmp_v, torch::Tensor& tmp_i,
                 torch::Tensor& out_v, torch::Tensor& out_i,
                 int64_t k, int64_t tile, int64_t threads1, int64_t threads2);
"""

_ext = load_inline(
    name="topk_ext_v2",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["topk_launch"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_100,code=sm_100a"],
    verbose=False,
)


def _choose_tile(batch: int, n: int, k: int):
    # aim for many blocks in stage1 for SM occupancy; tile power of 2.
    if n >= 65536:        # very long single rows (shape0-like)
        tile = 1024
    elif n >= 16384:
        tile = 1024
    elif n >= 8192:
        tile = 1024
    elif n >= 4096:
        tile = 512
    else:
        tile = max(64, 1 << (n - 1).bit_length())
    threads1 = 256
    threads2 = 256
    return tile, threads1, threads2


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.tile, self.t1, self.t2 = _choose_tile(batch, n, k)
        self.num_tiles = (n + self.tile - 1) // self.tile
        total = self.num_tiles * k
        self._padN = 1
        while self._padN < total:
            self._padN <<= 1

    def forward(self, x: torch.Tensor):
        batch = x.size(0)
        n = x.size(1)
        k = self.k
        tmp_v = torch.empty((batch, self.num_tiles, k), dtype=torch.float32, device=x.device)
        tmp_i = torch.empty((batch, self.num_tiles, k), dtype=torch.int32, device=x.device)
        out_v = torch.empty((batch, k), dtype=torch.float32, device=x.device)
        out_i = torch.empty((batch, k), dtype=torch.int64, device=x.device)
        _ext.topk_launch(x.contiguous(), tmp_v, tmp_i, out_v, out_i,
                         k, self.tile, self.t1, self.t2)
        return out_v, out_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]