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 parallel_scan_kernel(const float* __restrict__ a,
                                      const float* __restrict__ b,
                                      float* __restrict__ out,
                                      const int B, const int L, const int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * D) 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 aa = a[offset];
        float bb = b[offset];
        h = aa * h + bb;
        out[offset] = h;
    }
}

torch::Tensor parallel_scan_cuda(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;
    parallel_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_cuda(torch::Tensor a, torch::Tensor b);
'''

scan_module = load_inline(
    name='parallel_scan',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['parallel_scan_cuda'],
    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 scan_module.parallel_scan_cuda(a, b)