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


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

torch::Tensor fused_dw_silu_pw_cuda(torch::Tensor x,
                                    torch::Tensor dw_w,
                                    torch::Tensor pw_w,
                                    int64_t dilation,
                                    int64_t padding);
"""

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

__global__ void dw_silu_kernel(const float* __restrict__ x,
                               const float* __restrict__ dw,
                               float* __restrict__ tmp,
                               int B, int C, int T, int K, int dilation, int padding) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int n = B * T * C;
    if (idx >= n) return;

    int c = idx % C;
    int m = idx / C;       // m = b*T + t
    int t = m % T;
    int b = m / T;

    float acc = 0.0f;
    const float* xbc = x + ((b * C + c) * T);
    const float* dwc = dw + c * K;

    // K is 5 for the benchmark, but keep the loop generic.
    #pragma unroll
    for (int k = 0; k < 16; ++k) {
        if (k >= K) break;
        int tt = t + k * dilation - padding;
        if ((unsigned)tt < (unsigned)T) {
            acc += xbc[tt] * dwc[k];
        }
    }

    // SiLU: x * sigmoid(x)
    float sig = 1.0f / (1.0f + expf(-acc));
    tmp[idx] = acc * sig;
}

torch::Tensor fused_dw_silu_pw_cuda(torch::Tensor x,
                                    torch::Tensor dw_w,
                                    torch::Tensor pw_w,
                                    int64_t dilation,
                                    int64_t padding) {
    const int B = (int)x.size(0);
    const int C = (int)x.size(1);
    const int T = (int)x.size(2);
    const int K = (int)dw_w.size(2);
    const int O = (int)pw_w.size(0);

    // Intermediate is laid out as (B*T, C), directly consumable by GEMM.
    auto tmp = torch::empty({(int64_t)B * T, C}, x.options());

    const int threads = 256;
    const int total = B * T * C;
    const int blocks = (total + threads - 1) / threads;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dw_silu_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), dw_w.data_ptr<float>(), tmp.data_ptr<float>(),
        B, C, T, K, (int)dilation, (int)padding);

    // Pointwise conv is a matrix multiply: (B*T,C) x (C,O).
    auto pw2 = pw_w.view({O, C}).transpose(0, 1);
    auto out2 = at::mm(tmp, pw2);

    // Return logical (B,O,T) without a physical transpose/copy.
    return out2.view({B, T, O}).permute({0, 2, 1});
}
"""

_ext = load_inline(
    name="fused_dw_silu_pw1d_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_dw_silu_pw_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    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.dilation = dilation
        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:
        return _ext.fused_dw_silu_pw_cuda(x.contiguous(), self.dw.weight, self.pw.weight, self.dilation, self.padding)