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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void fused_dw_silu_pw_kernel(
    const float* __restrict__ x,
    const float* __restrict__ dw_weight,
    const float* __restrict__ pw_weight,
    float* __restrict__ out,
    int B, int C, int T, int K, int D, int P, int OC
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * OC * T;
    if (idx >= total) return;

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

    float sum = 0.0f;
    for (int ic = 0; ic < C; ++ic) {
        double dw_sum = 0.0;
        for (int k = 0; k < K; ++k) {
            int offset = k * D - P;
            int pos = t + offset;
            float val = 0.0f;
            if (pos >= 0 && pos < T) {
                val = x[b * C * T + ic * T + pos];
            }
            dw_sum += (double)dw_weight[ic * K + k] * val;
        }
        float silu_val = (float)(dw_sum * (1.0 / (1.0 + exp(-dw_sum))));
        sum += pw_weight[oc * C + ic] * silu_val;
    }
    out[b * OC * T + oc * T + t] = sum;
}

torch::Tensor fused_dw_silu_pw(
    torch::Tensor x,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    int K, int D, int P, int OC
) {
    int B = x.size(0);
    int C = x.size(1);
    int T = x.size(2);
    auto out = torch::empty({B, OC, T}, x.options());

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

    fused_dw_silu_pw_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        dw_weight.data_ptr<float>(),
        pw_weight.data_ptr<float>(),
        out.data_ptr<float>(),
        B, C, T, K, D, P, OC
    );
    return out;
}
"""

cpp_source = """
torch::Tensor fused_dw_silu_pw(
    torch::Tensor x,
    torch::Tensor dw_weight,
    torch::Tensor pw_weight,
    int K, int D, int P, int OC
);
"""

module = load_inline(
    name='fused_conv',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_dw_silu_pw'],
    verbose=True,
    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.K = kernel_size
        self.D = dilation
        self.P = self.padding
        self.OC = out_channels
        self.C = in_channels

    def forward(self, x):
        dw_w = self.dw.weight.view(self.C, self.K)
        pw_w = self.pw.weight.view(self.OC, self.C)
        return module.fused_dw_silu_pw(x.contiguous(), dw_w, pw_w, self.K, self.D, self.P, self.OC)