import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Selective scan (Mamba's S6 layer): projects input to derive input-dependent SSM parameters
    (delta, B, C), discretizes a learned state matrix A, runs the linear recurrence, and
    applies a skip connection via D. The fusion opportunity is keeping all intermediates
    (discretized parameters, hidden states) in SRAM rather than materializing them to global memory.

    Args:
        d_model (int): Input feature dimension.
        d_state (int): Hidden state dimension per channel (N in Mamba, typically 16).
    """
    def __init__(self, d_model: int, d_state: int):
        super(Model, self).__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        self.proj_B = nn.Linear(d_model, d_state, bias=False)
        self.proj_C = nn.Linear(d_model, d_state, bias=False)
        self.proj_delta = nn.Linear(d_model, d_model, bias=True)
        self.D = nn.Parameter(torch.ones(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        N = self.d_state

        delta = torch.nn.functional.softplus(self.proj_delta(x))
        B_t = self.proj_B(x)
        C_t = self.proj_C(x)

        A = -torch.exp(self.A_log)
        dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
        dB = delta.unsqueeze(-1) * B_t.unsqueeze(2)
        x_db = x.unsqueeze(-1) * dB

        h = torch.zeros(B, D, N, device=x.device, dtype=x.dtype)
        outputs = []
        for t in range(L):
            h = dA[:, t] * h + x_db[:, t]
            y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)
            outputs.append(y_t)

        y = torch.stack(outputs, dim=1)
        y = y + x * self.D.unsqueeze(0).unsqueeze(0)
        return y

batch_size = 16
d_model = 256
d_state = 16
seq_len = 2048

def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), popcorn_pri.jitter_int(seq_len), d_model)]

def get_init_inputs():
    return [d_model, d_state]

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


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

torch::Tensor selective_scan_forward(torch::Tensor x,
                                     torch::Tensor delta,
                                     torch::Tensor B_t,
                                     torch::Tensor C_t,
                                     torch::Tensor A_log,
                                     torch::Tensor D_param);
'''

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

#ifndef TILE_D
#define TILE_D 8
#endif
#ifndef STATE_W
#define STATE_W 16
#endif

__global__ __launch_bounds__(TILE_D * STATE_W, 2)
void selective_scan_kernel(const float* __restrict__ x,
                           const float* __restrict__ delta,
                           const float* __restrict__ B_t,
                           const float* __restrict__ C_t,
                           const float* __restrict__ A_log,
                           const float* __restrict__ D_param,
                           float* __restrict__ y,
                           int B, int L, int D, int N) {
    const int b = (int)blockIdx.y;
    const int d_base = (int)blockIdx.x * TILE_D;
    const int tid = (int)threadIdx.x;
    const int kd = tid >> 4;        // channel within the tile
    const int n  = tid & 15;        // state lane
    const int d = d_base + kd;

    const int lane_in_warp = tid & 31;
    const unsigned mask = 0xFFFFu << (lane_in_warp & 16);

    float h = 0.0f;
    float a = 0.0f;
    if (d < D && n < N) {
        a = -expf(A_log[d * N + n]);
    }

    float d_skip = 0.0f;
    if (d < D && n == 0) d_skip = D_param[d];

    const int bd_x_base = (b * L) * D;
    const int bd_s_base = (b * L) * N;

    for (int t = 0; t < L; ++t) {
        float xv = 0.0f;
        float dt = 0.0f;
        const int x_idx = bd_x_base + t * D + d;
        if (d < D && n == 0) {
            xv = x[x_idx];
            dt = delta[x_idx];
        }
        xv = __shfl_sync(mask, xv, 0, STATE_W);
        dt = __shfl_sync(mask, dt, 0, STATE_W);

        float acc = 0.0f;
        if (d < D && n < N) {
            const int s_idx = bd_s_base + t * N + n;
            const float bn = B_t[s_idx];
            const float cn = C_t[s_idx];
            const float ea = expf(dt * a);
            h = fmaf(xv * dt, bn, ea * h);
            acc = h * cn;
        }

        acc += __shfl_down_sync(mask, acc, 8, STATE_W);
        acc += __shfl_down_sync(mask, acc, 4, STATE_W);
        acc += __shfl_down_sync(mask, acc, 2, STATE_W);
        acc += __shfl_down_sync(mask, acc, 1, STATE_W);

        if (d < D && n == 0) {
            y[x_idx] = acc + xv * d_skip;
        }
    }
}

torch::Tensor selective_scan_forward(torch::Tensor x,
                                     torch::Tensor delta,
                                     torch::Tensor B_t,
                                     torch::Tensor C_t,
                                     torch::Tensor A_log,
                                     torch::Tensor D_param) {
    const int B = (int)x.size(0);
    const int L = (int)x.size(1);
    const int D = (int)x.size(2);
    const int N = (int)A_log.size(1);

    auto y = torch::empty_like(x);
    dim3 block(TILE_D * STATE_W);
    dim3 grid((D + TILE_D - 1) / TILE_D, B);

    selective_scan_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), delta.data_ptr<float>(), B_t.data_ptr<float>(), C_t.data_ptr<float>(),
        A_log.data_ptr<float>(), D_param.data_ptr<float>(), y.data_ptr<float>(), B, L, D, N);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='selective_scan_mamba_s6_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['selective_scan_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, d_model: int, d_state: int):
        super(ModelNew, self).__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        self.proj_B = nn.Linear(d_model, d_state, bias=False)
        self.proj_C = nn.Linear(d_model, d_state, bias=False)
        self.proj_delta = nn.Linear(d_model, d_model, bias=True)
        self.D = nn.Parameter(torch.ones(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        delta = F.softplus(self.proj_delta(x))
        B_t = self.proj_B(x)
        C_t = self.proj_C(x)
        return _ext.selective_scan_forward(x, delta, B_t, C_t, self.A_log, self.D)