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 depthwise_conv1d_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 <c10/cuda/CUDAGuard.h>

constexpr int BLOCK = 256;

__global__ void dwconv1d_k3_512_8192_vec4_kernel(const float* __restrict__ x,
                                             const float* __restrict__ w,
                                             const float* __restrict__ bias,
                                             float* __restrict__ y,
                                             int has_bias) {
    int c = blockIdx.x;
    int base = ((blockIdx.y << 9) + c) << 13; // (b*512+c)*8192
    const float w0 = w[c * 3 + 0];
    const float w1 = w[c * 3 + 1];
    const float w2 = w[c * 3 + 2];
    const float bv = has_bias ? bias[c] : 0.0f;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x + base);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y + base);
    #pragma unroll 8
    for (int vg = threadIdx.x; vg < 2048; vg += BLOCK) {
        int i = vg << 2;
        float4 xv = x4[vg];
        float left = (i > 0) ? x[base + i - 1] : 0.0f;
        float right = (i < 8188) ? x[base + i + 4] : 0.0f;
        float4 out;
        out.x = left * w0 + xv.x * w1 + xv.y * w2 + bv;
        out.y = xv.x * w0 + xv.y * w1 + xv.z * w2 + bv;
        out.z = xv.y * w0 + xv.z * w1 + xv.w * w2 + bv;
        out.w = xv.z * w0 + xv.w * w1 + right * w2 + bv;
        y4[vg] = out;
    }
}

__global__ void dwconv1d_k3s1p1_vec4_kernel(const float* __restrict__ x,
                                             const float* __restrict__ w,
                                             const float* __restrict__ bias,
                                             float* __restrict__ y,
                                             int C, int L, int has_bias) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    int base = (b * C + c) * L;
    const float w0 = w[c * 3 + 0];
    const float w1 = w[c * 3 + 1];
    const float w2 = w[c * 3 + 2];
    const float bv = has_bias ? bias[c] : 0.0f;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x + base);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y + base);
    int vecL = L >> 2;
    for (int vg = threadIdx.x; vg < vecL; vg += BLOCK) {
        int i = vg << 2;
        float4 xv = x4[vg];
        float left = (i > 0) ? x[base + i - 1] : 0.0f;
        float right = (i + 4 < L) ? x[base + i + 4] : 0.0f;
        float4 out;
        out.x = left * w0 + xv.x * w1 + xv.y * w2 + bv;
        out.y = xv.x * w0 + xv.y * w1 + xv.z * w2 + bv;
        out.z = xv.y * w0 + xv.z * w1 + xv.w * w2 + bv;
        out.w = xv.z * w0 + xv.w * w1 + right * w2 + bv;
        y4[vg] = out;
    }
}

__global__ void dwconv1d_k3s1p1_scalar_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ bias,
                                        float* __restrict__ y,
                                        int C, int L, int has_bias) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    int base = (b * C + c) * L;
    const float w0 = w[c * 3 + 0];
    const float w1 = w[c * 3 + 1];
    const float w2 = w[c * 3 + 2];
    const float bv = has_bias ? bias[c] : 0.0f;
    for (int idx = threadIdx.x; idx < L; idx += BLOCK) {
        float l = (idx > 0) ? x[base + idx - 1] : 0.0f;
        float m = x[base + idx];
        float r = (idx + 1 < L) ? x[base + idx + 1] : 0.0f;
        y[base + idx] = l * w0 + m * w1 + r * w2 + bv;
    }
}

__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 outL,
                                        int stride, int padding, int has_bias) {
    int linear = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C * outL;
    if (linear >= total) return;
    int pos = linear % outL;
    int tmp = linear / outL;
    int c = tmp % C;
    int b = tmp / C;
    int xbase = (b * C + c) * L;
    float acc = has_bias ? bias[c] : 0.0f;
    int in0 = pos * stride - padding;
    int wbase = c * K;
    for (int k = 0; k < K; ++k) {
        int ix = in0 + k;
        if ((unsigned)ix < (unsigned)L) acc += x[xbase + ix] * w[wbase + k];
    }
    y[linear] = acc;
}

torch::Tensor depthwise_conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t stride64, int64_t padding64) {
    c10::cuda::CUDAGuard device_guard(x.device());
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat, "float32 required");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(), "contiguous required");
    int B = (int)x.size(0), C = (int)x.size(1), L = (int)x.size(2);
    int K = (int)weight.size(2), stride = (int)stride64, padding = (int)padding64;
    int outL = (L + 2 * padding - K) / stride + 1;
    auto y = torch::empty({B, C, outL}, x.options());
    int has_bias = bias.numel() > 0;
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (K == 3 && stride == 1 && padding == 1 && outL == L) {
        dim3 grid(C, B);
        if (C == 512 && L == 8192) dwconv1d_k3_512_8192_vec4_kernel<<<grid, BLOCK, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(), has_bias);
        else if ((L & 3) == 0) dwconv1d_k3s1p1_vec4_kernel<<<grid, BLOCK, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(), C, L, has_bias);
        else dwconv1d_k3s1p1_scalar_kernel<<<grid, BLOCK, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(), C, L, has_bias);
    } else {
        int total = B * C * outL;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        dwconv1d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(), B, C, L, K, outL, stride, padding, has_bias);
    }
    return y;
}
"""

_ext = load_inline(name="dwconv1d_ext_final", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["depthwise_conv1d_forward"], extra_cflags=["-O3"],
                   extra_cuda_cflags=["-O3", "--use_fast_math"], 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.register_buffer("_empty_bias", torch.empty(0), persistent=False)
        self.stride = stride
        self.padding = padding
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv1d.bias if self.conv1d.bias is not None else self._empty_bias
        return _ext.depthwise_conv1d_forward(x, self.conv1d.weight, b, self.stride, self.padding)