import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Channel-wise gated delta attention recurrence.

    The recurrent state decays independently for each output channel, which
    makes the update closer to channel-gated delta attention variants used in
    fast long-context baselines.
    """

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

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        gate: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, num_heads, head_dim = q.shape
        state = torch.zeros(batch_size, num_heads, head_dim, head_dim, dtype=q.dtype, device=q.device)
        outputs = []

        for t in range(seq_len):
            gate_t = gate[:, t].unsqueeze(-2)
            kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2)
            state = gate_t * state + kv_outer
            y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2)
            outputs.append(y_t)

        return torch.stack(outputs, dim=1)


batch_size = 2
seq_len = 128
num_heads = 4
head_dim = 24


def get_inputs():
    p = popcorn_pri
    q = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    k = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    v = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    gate = torch.sigmoid(torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32))
    return [q, k, v, gate]


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 gated_delta_forward(torch::Tensor q,
                                  torch::Tensor k,
                                  torch::Tensor v,
                                  torch::Tensor gate);
"""

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

__global__ void gated_delta_kernel(const float* __restrict__ q,
                                   const float* __restrict__ k,
                                   const float* __restrict__ v,
                                   const float* __restrict__ gate,
                                   float* __restrict__ out,
                                   int T, int H, int D) {
    const int j = blockIdx.x;      // output/value channel
    const int h = blockIdx.y;
    const int b = blockIdx.z;
    const int lane = threadIdx.x;  // key/query channel within a warp
    const unsigned mask = 0xffffffffu;

    float s = 0.0f;  // state[lane, j] for this (b,h,j)

    for (int t = 0; t < T; ++t) {
        const int base = ((b * T + t) * H + h) * D;

        float gt = 0.0f;
        float vv = 0.0f;
        if (lane == 0) {
            gt = gate[base + j];
            vv = v[base + j];
        }
        gt = __shfl_sync(mask, gt, 0);
        vv = __shfl_sync(mask, vv, 0);

        float acc = 0.0f;
        if (lane < D) {
            const float kval = k[base + lane];
            const float qval = q[base + lane];
            const float kv = kval * vv;
            const float decayed = gt * s;
            s = decayed + kv;
            acc = qval * s;
        }

        // Sum over the key/query dimension to produce y[t, j].
        acc += __shfl_down_sync(mask, acc, 16);
        acc += __shfl_down_sync(mask, acc, 8);
        acc += __shfl_down_sync(mask, acc, 4);
        acc += __shfl_down_sync(mask, acc, 2);
        acc += __shfl_down_sync(mask, acc, 1);

        if (lane == 0) {
            out[base + j] = acc;
        }
    }
}

torch::Tensor gated_delta_forward(torch::Tensor q,
                                  torch::Tensor k,
                                  torch::Tensor v,
                                  torch::Tensor gate) {
    const at::cuda::CUDAGuard device_guard(q.device());
    const int B = (int)q.size(0);
    const int T = (int)q.size(1);
    const int H = (int)q.size(2);
    const int D = (int)q.size(3);

    auto out = torch::empty_like(q);

    dim3 grid(D, H, B);
    dim3 block(32);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gated_delta_kernel<<<grid, block, 0, stream>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), gate.data_ptr<float>(),
        out.data_ptr<float>(), T, H, D);
    return out;
}
"""

_ext = load_inline(
    name="gated_delta_attention_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["gated_delta_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--fmad=false"],
    with_cuda=True,
    verbose=False,
)


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

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
        return _ext.gated_delta_forward(q, k, v, gate)