import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Apply rotary position embedding to incoming keys and update the KV cache.

    This isolates a common decode-time kernel: rotate the fresh key vectors for
    the current positions and store both keys and values into the running cache.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        k_new: torch.Tensor,
        v_new: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        cache_k: torch.Tensor,
        cache_v: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        updated_k = cache_k.clone()
        updated_v = cache_v.clone()
        rot_dim = k_new.shape[-1]

        for b in range(k_new.shape[0]):
            for t in range(k_new.shape[1]):
                pos = int(positions[b, t].item())
                c = cos[pos]
                s = sin[pos]
                x = k_new[b, t]
                x_even = x[..., 0::2]
                x_odd = x[..., 1::2]
                rot_even = x_even * c - x_odd * s
                rot_odd = x_even * s + x_odd * c
                rotated = torch.empty_like(x)
                rotated[..., 0::2] = rot_even
                rotated[..., 1::2] = rot_odd
                updated_k[b, pos] = rotated
                updated_v[b, pos] = v_new[b, t]

        return torch.stack((updated_k, updated_v), dim=0)


batch_size = 2
cache_len = 128
update_len = 16
num_heads = 4
head_dim = 32


def _rope_tables(cache_len: int, half_dim: int):
    positions = torch.arange(cache_len, dtype=torch.float32).unsqueeze(1)
    freqs = torch.pow(10000.0, -torch.arange(half_dim, dtype=torch.float32) / half_dim).unsqueeze(0)
    angles = positions * freqs
    return torch.cos(angles), torch.sin(angles)


def get_inputs():
    p = popcorn_pri
    k_new = torch.randn(p.jitter_int(batch_size), p.jitter_int(update_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    v_new = torch.randn(p.jitter_int(batch_size), p.jitter_int(update_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    cos, sin = _rope_tables(p.jitter_int(cache_len), p.jitter_int(head_dim, align=8) // 2)
    cache_k = torch.randn(p.jitter_int(batch_size), p.jitter_int(cache_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    cache_v = torch.randn(p.jitter_int(batch_size), p.jitter_int(cache_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    base = torch.arange(p.jitter_int(update_len), dtype=torch.int32).unsqueeze(0).repeat(p.jitter_int(batch_size), 1)
    positions = (base + torch.tensor([[7], [31]], dtype=torch.int32)) % p.jitter_int(cache_len)
    return [k_new, v_new, cos, sin, cache_k, cache_v, positions]


def get_init_inputs():
    return []

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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor rope_cache_update(torch::Tensor k_new,
                                torch::Tensor v_new,
                                torch::Tensor cos,
                                torch::Tensor sin,
                                torch::Tensor cache_k,
                                torch::Tensor cache_v,
                                torch::Tensor positions);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>

// Specialized for the problem sizes in the benchmark:
// B=2, cache_len=128, update_len=16, num_heads=4, head_dim=32.
// One kernel writes the final stacked tensor directly.  For cache locations not
// updated it copies the cache; for updated locations it writes the rotated new
// key or the new value.  This avoids clone+stack and avoids inter-kernel races.
__global__ void rope_cache_update_kernel(const float* __restrict__ k_new,
                                         const float* __restrict__ v_new,
                                         const float* __restrict__ cos,
                                         const float* __restrict__ sin,
                                         const float* __restrict__ cache_k,
                                         const float* __restrict__ cache_v,
                                         const int* __restrict__ positions,
                                         float* __restrict__ out) {
    const int n_per = 2 * 128 * 4 * 32;   // B*L*H*D = 32768
    const int n_tot = 2 * n_per;          // stacked k,v

    for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n_tot;
         idx += blockDim.x * gridDim.x) {
        const int which = idx >> 15;      // 0 => k, 1 => v
        const int inner = idx & 32767;

        const int d   = inner & 31;
        const int h   = (inner >> 5) & 3;
        const int pos = (inner >> 7) & 127;
        const int b   = inner >> 14;

        int tmatch = -1;
        const int pbase = b << 4;
        #pragma unroll
        for (int t = 0; t < 16; ++t) {
            if (positions[pbase + t] == pos) tmatch = t;  // last update wins
        }

        float val;
        if (which == 0) {
            if (tmatch >= 0) {
                const int pair = d >> 1;
                const int base = (((b << 4) + tmatch) * 4 + h) * 32;
                const float xe = k_new[base + (pair << 1)];
                const float xo = k_new[base + (pair << 1) + 1];
                const float c = cos[(pos << 4) + pair];
                const float s = sin[(pos << 4) + pair];
                if ((d & 1) == 0) {
                    val = xe * c - xo * s;
                } else {
                    val = xe * s + xo * c;
                }
            } else {
                val = cache_k[inner];
            }
        } else {
            if (tmatch >= 0) {
                const int base = (((b << 4) + tmatch) * 4 + h) * 32;
                val = v_new[base + d];
            } else {
                val = cache_v[inner];
            }
        }
        out[idx] = val;
    }
}

// Generic fallback kernel for the same operation for contiguous fp32 tensors.
// It is used only if shapes differ from the benchmark constants.
__global__ void rope_cache_update_kernel_generic(const float* __restrict__ k_new,
                                                 const float* __restrict__ v_new,
                                                 const float* __restrict__ cos,
                                                 const float* __restrict__ sin,
                                                 const float* __restrict__ cache_k,
                                                 const float* __restrict__ cache_v,
                                                 const int* __restrict__ positions,
                                                 float* __restrict__ out,
                                                 int B, int T, int L, int H, int D) {
    const int n_per = B * L * H * D;
    const int n_tot = 2 * n_per;
    const int half = D >> 1;
    for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n_tot;
         idx += blockDim.x * gridDim.x) {
        const int which = idx / n_per;
        int rem = idx - which * n_per;
        const int d = rem % D; rem /= D;
        const int h = rem % H; rem /= H;
        const int pos = rem % L;
        const int b = rem / L;

        int tmatch = -1;
        for (int t = 0; t < T; ++t) {
            if (positions[b * T + t] == pos) tmatch = t;
        }

        float val;
        if (which == 0) {
            if (tmatch >= 0) {
                const int pair = d >> 1;
                const int base = ((b * T + tmatch) * H + h) * D;
                const float xe = k_new[base + 2 * pair];
                const float xo = k_new[base + 2 * pair + 1];
                const float c = cos[pos * half + pair];
                const float s = sin[pos * half + pair];
                val = ((d & 1) == 0) ? (xe * c - xo * s) : (xe * s + xo * c);
            } else {
                val = cache_k[idx];
            }
        } else {
            if (tmatch >= 0) {
                const int base = ((b * T + tmatch) * H + h) * D;
                val = v_new[base + d];
            } else {
                val = cache_v[idx - n_per];
            }
        }
        out[idx] = val;
    }
}

torch::Tensor rope_cache_update(torch::Tensor k_new,
                                torch::Tensor v_new,
                                torch::Tensor cos,
                                torch::Tensor sin,
                                torch::Tensor cache_k,
                                torch::Tensor cache_v,
                                torch::Tensor positions) {
    const at::cuda::CUDAGuard device_guard(k_new.device());
    const int64_t B = k_new.size(0);
    const int64_t T = k_new.size(1);
    const int64_t H = k_new.size(2);
    const int64_t D = k_new.size(3);
    const int64_t L = cache_k.size(1);

    auto out = torch::empty({2, B, L, H, D}, cache_k.options());

    const int threads = 256;
    const int64_t total = 2 * B * L * H * D;
    const int blocks = (int)((total + threads - 1) / threads);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (B == 2 && T == 16 && L == 128 && H == 4 && D == 32) {
        rope_cache_update_kernel<<<blocks, threads, 0, stream>>>(
            k_new.data_ptr<float>(), v_new.data_ptr<float>(),
            cos.data_ptr<float>(), sin.data_ptr<float>(),
            cache_k.data_ptr<float>(), cache_v.data_ptr<float>(),
            positions.data_ptr<int>(), out.data_ptr<float>());
    } else {
        rope_cache_update_kernel_generic<<<blocks, threads, 0, stream>>>(
            k_new.data_ptr<float>(), v_new.data_ptr<float>(),
            cos.data_ptr<float>(), sin.data_ptr<float>(),
            cache_k.data_ptr<float>(), cache_v.data_ptr<float>(),
            positions.data_ptr<int>(), out.data_ptr<float>(),
            (int)B, (int)T, (int)L, (int)H, (int)D);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="rope_cache_update_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["rope_cache_update"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, k_new, v_new, cos, sin, cache_k, cache_v, positions):
        return _ext.rope_cache_update(k_new, v_new, cos, sin, cache_k, cache_v, positions)