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

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

__global__ void scan_kernel(const float* a, const float* b, float* out, int B, int L, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * D;
    if (idx >= total) return;
    
    int batch = idx / D;
    int dim = idx % D;
    
    float h = 0.0f;
    for (int t = 0; t < L; t++) {
        int offset = (batch * L + t) * D + dim;
        float at = a[offset];
        float bt = b[offset];
        h = at * h + bt;
        out[offset] = h;
    }
}

torch::Tensor parallel_scan(torch::Tensor a, torch::Tensor b) {
    auto B = a.size(0);
    auto L = a.size(1);
    auto D = a.size(2);
    
    auto out = torch::empty({B, L, D}, a.options());
    
    int total = B * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    scan_kernel<<<blocks, threads>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        out.data_ptr<float>(),
        B, L, D
    );
    
    return out;
}
"""

cpp_source = """
torch::Tensor parallel_scan(torch::Tensor a, torch::Tensor b);
"""

module = load_inline(
    name='scan_module',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['parallel_scan'],
    verbose=True
)

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

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return self.module.parallel_scan(a, b)