import torch
import torch.nn as nn
import torch.nn.functional as F

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


class Model(nn.Module):
    """
    **Fused depthwise–pointwise 1D temporal convolution** on ``(B, C, T)`` sequences.

    Fuses a **temporal depthwise** ``k``-tap conv (per-channel along time) with
    **SiLU** and a **pointwise** ``1×1`` mix—standard in temporal encoders for
    proprio, audio, or token-wise time series in robotics / VLA backbones.

    Padding uses dilated effective radius so length is unchanged for odd
    ``kernel_size``.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dilation: int = 1,
    ):
        super().__init__()
        assert kernel_size % 2 == 1, "odd kernel_size keeps T unchanged with symmetric padding"
        assert dilation >= 1
        self.padding = ((kernel_size - 1) * dilation) // 2
        self.dw = nn.Conv1d(
            in_channels,
            in_channels,
            kernel_size,
            padding=self.padding,
            dilation=dilation,
            groups=in_channels,
            bias=False,
        )
        self.pw = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C_in, T) temporal features.
        Returns:
            (B, C_out, T)
        """
        x = self.dw(x)
        x = F.silu(x)
        return self.pw(x)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dilation: int = 1,
    ):
        super().__init__()
        self._impl = Model(in_channels, out_channels, kernel_size, dilation=dilation)

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


batch_size = 16
in_channels = 128
out_channels = 128
kernel_size = 5
dilation = 2
seq_len = 2048


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), in_channels, p.jitter_int(seq_len))
    return [x]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, dilation]

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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void fused_kernel(
    const float* __restrict__ input,
    const float* __restrict__ dw_weight,
    const float* __restrict__ pw_weight,
    float* __restrict__ output,
    const int B, const int C_in, const int T,
    const int K, const int D, const int P, const int C_out
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C_out * T;
    if (idx >= total) return;

    int b = idx / (C_out * T);
    int oc = (idx / T) % C_out;
    int t = idx % T;

    double val = 0.0;
    for (int ic = 0; ic < C_in; ++ic) {
        double dw_sum = 0.0;
        for (int k = 0; k < K; ++k) {
            int t_in = t - P + k * D;
            float x_val = 0.0f;
            if (t_in >= 0 && t_in < T) {
                x_val = input[((b * C_in + ic) * T) + t_in];
            }
            dw_sum += (double)x_val * (double)dw_weight[(ic * K) + k];
        }
        double silu = dw_sum / (1.0 + exp(-dw_sum));
        val += silu * (double)pw_weight[(oc * C_in + ic)];
    }
    output[idx] = (float)val;
}

torch::Tensor fused_dw_silu_pw(torch::Tensor x, torch::Tensor dw_w, torch::Tensor pw_w, int padding, int dilation, int kernel_size) {
    const int B = x.size(0);
    const int C_in = x.size(1);
    const int T = x.size(2);
    const int C_out = pw_w.size(0);
    const int K = kernel_size;
    const int D = dilation;
    const int P = padding;

    auto output = torch::empty({B, C_out, T}, x.options());

    int total = B * C_out * T;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;

    fused_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        dw_w.data_ptr<float>(),
        pw_w.data_ptr<float>(),
        output.data_ptr<float>(),
        B, C_in, T, K, D, P, C_out
    );
    cudaDeviceSynchronize();

    return output;
}
'''

cpp_source = '''
torch::Tensor fused_dw_silu_pw(torch::Tensor x, torch::Tensor dw_w, torch::Tensor pw_w, int padding, int dilation, int kernel_size);
'''

module = load_inline(
    name='fused_conv',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_dw_silu_pw'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation=1):
        super().__init__()
        assert kernel_size % 2 == 1
        self.padding = ((kernel_size - 1) * dilation) // 2
        self.dw = nn.Conv1d(in_channels, in_channels, kernel_size, padding=self.padding, dilation=dilation, groups=in_channels, bias=False)
        self.pw = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
        self.module = module

    def forward(self, x):
        dw_w = self.dw.weight.squeeze(1)
        pw_w = self.pw.weight.squeeze(-1)
        return self.module.fused_dw_silu_pw(x.contiguous(), dw_w, pw_w, self.padding, self.dw.dilation[0], self.dw.kernel_size[0])