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"""
torch::Tensor recurrence_forward(torch::Tensor a, torch::Tensor b);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

__global__ __launch_bounds__(32, 16)
void recurrence_kernel(const float* __restrict__ a,
                       const float* __restrict__ b,
                       float* __restrict__ out,
                       int B, int L, int D) {
    // One warp processes one contiguous tile of 32 channels for one batch item.
    // This preserves fully coalesced memory accesses for the [B,L,D] layout while
    // each lane performs the dependent recurrence for its channel.
    int lane = threadIdx.x;  // blockDim.x == 32
    int tiles_per_batch = (D + 31) >> 5;
    int tile = blockIdx.x;
    int batch = tile / tiles_per_batch;
    int d0 = (tile - batch * tiles_per_batch) << 5;
    int d = d0 + lane;
    if (batch >= B || d >= D) return;

    int idx = (batch * L * D) + d;
    float h = 0.0f;

    int t = 0;
    // small manual unroll; L is 4096 in the benchmark
    for (; t + 3 < L; t += 4) {
        float av0 = a[idx];
        float bv0 = b[idx];
        h = __fadd_rn(__fmul_rn(av0, h), bv0);
        out[idx] = h;
        idx += D;

        float av1 = a[idx];
        float bv1 = b[idx];
        h = __fadd_rn(__fmul_rn(av1, h), bv1);
        out[idx] = h;
        idx += D;

        float av2 = a[idx];
        float bv2 = b[idx];
        h = __fadd_rn(__fmul_rn(av2, h), bv2);
        out[idx] = h;
        idx += D;

        float av3 = a[idx];
        float bv3 = b[idx];
        h = __fadd_rn(__fmul_rn(av3, h), bv3);
        out[idx] = h;
        idx += D;
    }
    for (; t < L; ++t) {
        float av = a[idx];
        float bv = b[idx];
        h = __fadd_rn(__fmul_rn(av, h), bv);
        out[idx] = h;
        idx += D;
    }
}

torch::Tensor recurrence_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);
    auto out = torch::empty_like(a);
    c10::cuda::CUDAGuard device_guard(a.device());
    int tiles_per_batch = (D + 31) >> 5;
    dim3 grid(B * tiles_per_batch);
    dim3 block(32);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    recurrence_kernel<<<grid, block, 0, stream>>>(
        a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), B, L, D);
    return out;
}
"""

_ext = load_inline(
    name="recurrence_scan_warp_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["recurrence_forward"],
    extra_cflags=["-O3"],
    extra_cuda_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.recurrence_forward(a, b)