import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a parallel associative scan (linear recurrence) over a 1D sequence.
    Given coefficient a and input b at each timestep, computes:
        h_0 = b_0
        h_t = a_t * h_{t-1} + b_t  for t > 0
    This is the core parallelizable primitive underlying Mamba's selective scan and
    other SSM architectures. Sequential in naive PyTorch but can be parallelized
    via the associative scan algorithm in O(L log L) span.

    Args:
        d_model (int): Number of independent recurrence channels.
    """
    def __init__(self, d_model: int):
        super(Model, self).__init__()
        self.d_model = d_model

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        B, L, D = a.shape
        h = torch.zeros(B, D, device=a.device, dtype=a.dtype)
        outputs = []
        for t in range(L):
            h = a[:, t] * h + b[:, t]
            outputs.append(h)
        return torch.stack(outputs, dim=1)

batch_size = 32
d_model = 256
seq_len = 4096

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

def get_init_inputs():
    return [d_model]

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor scan_forward(torch::Tensor a, torch::Tensor b);
"""

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

#define CHUNK 128

__global__ void chunk_summary_kernel(const float* __restrict__ a, const float* __restrict__ b,
                                     float* __restrict__ sA, float* __restrict__ sB,
                                     int B, int L, int D, int C) {
    int batch = blockIdx.x;
    int chunk = blockIdx.y;
    int d = threadIdx.x;
    if (d >= D) return;
    int start = chunk * CHUNK;
    int end = start + CHUNK;
    if (end > L) end = L;
    float A = 1.0f;
    float h = 0.0f;
    #pragma unroll 4
    for (int t = start; t < end; ++t) {
        int idx = (batch * L + t) * D + d;
        float av = a[idx];
        h = av * h + b[idx];
        A = av * A;
    }
    int sidx = (batch * C + chunk) * D + d;
    sA[sidx] = A;
    sB[sidx] = h;
}

__global__ void apply_kernel(const float* __restrict__ a, const float* __restrict__ b,
                             const float* __restrict__ sA, const float* __restrict__ sB,
                             float* __restrict__ out,
                             int B, int L, int D, int C) {
    int batch = blockIdx.x;
    int chunk = blockIdx.y;
    int d = threadIdx.x;
    if (d >= D) return;

    float h = 0.0f;
    for (int c = 0; c < chunk; ++c) {
        int idx = (batch * C + c) * D + d;
        h = sA[idx] * h + sB[idx];
    }

    int start = chunk * CHUNK;
    int end = start + CHUNK;
    if (end > L) end = L;
    #pragma unroll 4
    for (int t = start; t < end; ++t) {
        int idx = (batch * L + t) * D + d;
        h = a[idx] * h + b[idx];
        out[idx] = h;
    }
}

torch::Tensor scan_forward(torch::Tensor a, torch::Tensor b) {
    const int B = (int)a.size(0);
    const int L = (int)a.size(1);
    const int D = (int)a.size(2);
    const int C = (L + CHUNK - 1) / CHUNK;
    auto out = torch::empty_like(a);
    auto opts = a.options();
    auto sA = torch::empty({B, C, D}, opts);
    auto sB = torch::empty({B, C, D}, opts);
    dim3 block(D);
    dim3 grid(B, C);
    chunk_summary_kernel<<<grid, block>>>(a.data_ptr<float>(), b.data_ptr<float>(), sA.data_ptr<float>(), sB.data_ptr<float>(), B, L, D, C);
    apply_kernel<<<grid, block>>>(a.data_ptr<float>(), b.data_ptr<float>(), sA.data_ptr<float>(), sB.data_ptr<float>(), out.data_ptr<float>(), B, L, D, C);
    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, d_model: int):
        super(ModelNew, self).__init__()
        self.d_model = d_model

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return _ext.scan_forward(a, b)