import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a depthwise 1D convolution where each input channel is convolved with its own filter.
    Fundamental primitive in Hyena, Mamba, and genomic sequence models (Enformer, HyenaDNA).

    Args:
        channels (int): Number of input/output channels.
        kernel_size (int): Size of the convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        bias (bool, optional): If True, adds a learnable bias. Defaults to False.
    """
    def __init__(self, channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d = nn.Conv1d(channels, channels, kernel_size, stride=stride, padding=padding, groups=channels, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv1d(x)

batch_size = 32
channels = 512
kernel_size = 3
seq_len = 8192
stride = 1
padding = 1

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

def get_init_inputs():
    return [channels, kernel_size, stride, padding]

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

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

torch::Tensor dwconv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride, int64_t padding);
"""

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

__global__ void dwconv1d_k3s1p1_vec4_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    int rows, int C, int L, int has_bias)
{
    int vecs_per_row = L >> 2;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t total = (int64_t)rows * vecs_per_row;
    if (idx >= total) return;

    int row = (int)(idx / vecs_per_row);
    int vec = (int)(idx - (int64_t)row * vecs_per_row);
    int t = vec << 2;
    int c = row % C;
    int base = row * L + t;

    float w0 = w[c * 3 + 0];
    float w1 = w[c * 3 + 1];
    float w2 = w[c * 3 + 2];
    float bv = has_bias ? bias[c] : 0.0f;

    const float* xp = x + base;
    float xm1 = (t == 0) ? 0.0f : xp[-1];
    float x0 = xp[0];
    float x1 = xp[1];
    float x2 = xp[2];
    float x3 = xp[3];
    float x4 = (t + 4 >= L) ? 0.0f : xp[4];

    float o0 = fmaf(w2, x1, fmaf(w1, x0, w0 * xm1)) + bv;
    float o1 = fmaf(w2, x2, fmaf(w1, x1, w0 * x0)) + bv;
    float o2 = fmaf(w2, x3, fmaf(w1, x2, w0 * x1)) + bv;
    float o3 = fmaf(w2, x4, fmaf(w1, x3, w0 * x2)) + bv;

    *reinterpret_cast<float4*>(y + base) = make_float4(o0, o1, o2, o3);
}

__global__ void dwconv1d_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    int B, int C, int L, int K, int Lout, int stride, int padding, int has_bias)
{
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t total = (int64_t)B * C * Lout;
    if (idx >= total) return;

    int t = (int)(idx % Lout);
    int tmp = (int)(idx / Lout);
    int c = tmp % C;
    int b = tmp / C;

    float acc = has_bias ? bias[c] : 0.0f;
    int in0 = t * stride - padding;
    const float* xrow = x + ((int64_t)b * C + c) * L;
    const float* wrow = w + (int64_t)c * K;
    #pragma unroll
    for (int k = 0; k < 16; ++k) {
        if (k >= K) break;
        int in_t = in0 + k;
        if ((unsigned)in_t < (unsigned)L) {
            acc = fmaf(wrow[k], xrow[in_t], acc);
        }
    }
    for (int k = 16; k < K; ++k) {
        int in_t = in0 + k;
        if ((unsigned)in_t < (unsigned)L) {
            acc = fmaf(wrow[k], xrow[in_t], acc);
        }
    }
    y[idx] = acc;
}

torch::Tensor dwconv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride64, int64_t padding64) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only fp32 supported");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "only fp32 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");

    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int L = (int)x.size(2);
    int K = (int)weight.size(2);
    int stride = (int)stride64;
    int padding = (int)padding64;
    int Lout = (L + 2 * padding - K) / stride + 1;
    auto y = torch::empty({B, C, Lout}, x.options());
    int has_bias = (bias.numel() != 0) ? 1 : 0;

    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && padding == 1 && Lout == L && ((L & 3) == 0)) {
        int rows = B * C;
        int64_t total_vec = (int64_t)rows * (L >> 2);
        int blocks = (int)((total_vec + threads - 1) / threads);
        dwconv1d_k3s1p1_vec4_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), rows, C, L, has_bias);
    } else {
        int64_t total = (int64_t)B * C * Lout;
        int blocks = (int)((total + threads - 1) / threads);
        dwconv1d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), B, C, L, K, Lout, stride, padding, has_bias);
    }
    return y;
}
"""

_dwconv1d_ext = load_inline(
    name="dwconv1d_depthwise_k3_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["dwconv1d_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv1d = nn.Conv1d(channels, channels, kernel_size, stride=stride, padding=padding, groups=channels, bias=bias)
        self._stride_val = stride
        self._padding_val = padding
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        b = self.conv1d.bias if self.conv1d.bias is not None else self._empty_bias
        return _dwconv1d_ext.dwconv1d_forward(x, self.conv1d.weight, b, int(self._stride_val), int(self._padding_val))