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/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define TCONST 128
#define HCONST 4
#define DCONST 24

__global__ __launch_bounds__(32, 16)
void gated_delta_kernel(const float* __restrict__ q,
                        const float* __restrict__ k,
                        const float* __restrict__ v,
                        const float* __restrict__ gate,
                        float* __restrict__ out) {
    int lane = threadIdx.x;
    int bid = blockIdx.x;
    int j = bid % DCONST;
    int bh = bid / DCONST;
    int h = bh & 3;
    int b = bh >> 2;
    int base = (b * TCONST * HCONST + h) * DCONST;

    float st = 0.0f;
    unsigned mask = 0xffffffffu;

    #pragma unroll 1
    for (int t = 0; t < TCONST; ++t) {
        float g = 0.0f;
        float vv = 0.0f;
        if (lane == 0) {
            g = gate[base + j];
            vv = v[base + j];
        }
        g = __shfl_sync(mask, g, 0);
        vv = __shfl_sync(mask, vv, 0);

        float sum = 0.0f;
        if (lane < DCONST) {
            float kk = k[base + lane];
            float qq = q[base + lane];
            st = fmaf(g, st, kk * vv);
            sum = qq * st;
        }

        sum += __shfl_down_sync(mask, sum, 16);
        sum += __shfl_down_sync(mask, sum, 8);
        sum += __shfl_down_sync(mask, sum, 4);
        sum += __shfl_down_sync(mask, sum, 2);
        sum += __shfl_down_sync(mask, sum, 1);

        if (lane == 0) {
            out[base + j] = sum;
        }
        base += HCONST * DCONST;
    }
}

torch::Tensor gated_delta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate) {
    auto out = torch::empty_like(q);
    gated_delta_kernel<<<2 * HCONST * DCONST, 32>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), gate.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
"""

_ext = load_inline(
    name="gated_delta_ext_final",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["gated_delta_forward"],
    extra_cuda_cflags=["-O3"],
    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)