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__ input,
    const float* __restrict__ dw_weight,
    const float* __restrict__ pw_weight,
    float* __restrict__ output,
    int B, int C, int T, int K, int D, int P, int Cout
) {
    int b = blockIdx.x;
    int t = blockIdx.y * blockDim.x + threadIdx.x;
    if (t >= T) return;

    extern __shared__ float shared_dw[]; // C * K floats

    // Load dw weights into shared (simplified, assume small)
    for (int c = 0; c < C; c++) {
        for (int k = 0; k < K; k++) {
            shared_dw[c * K + k] = dw_weight[c * K + k];
        }
    }
    __syncthreads();

    float dw_out[128]; // assume C<=128
    for (int c = 0; c < C; c++) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            int t_idx = t + (k - K/2) * D;
            if (t_idx >= 0 && t_idx < T) {
                sum += input[(b * C + c) * T + t_idx] * shared_dw[c * K + k];
            }
        }
        dw_out[c] = sum / (1.0f + expf(-sum)); // SiLU
    }

    for (int co = 0; co < Cout; co++) {
        float sum = 0.0f;
        for (int ci = 0; ci < C; ci++) {
            sum += dw_out[ci] * pw_weight[co * C + ci];
        }
        output[(b * Cout + co) * T + t] = sum;
    }
}

torch::Tensor fused_forward(
    torch::Tensor x,
    torch::Tensor dw_w,
    torch::Tensor pw_w,
    int padding,
    int dilation,
    int kernel_size
) {
    int B = x.size(0);
    int C = x.size(1);
    int T = x.size(2);
    int Cout = pw_w.size(0);
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
    auto out = torch::empty({B, Cout, T}, options);

    dim3 grid(B, (T + 255) / 256);
    dim3 block(256);
    size_t shmem = C * kernel_size * sizeof(float);

    fused_dw_silu_pw_kernel<<<grid, block, shmem>>>(
        x.data_ptr<float>(),
        dw_w.data_ptr<float>(),
        pw_w.data_ptr<float>(),
        out.data_ptr<float>(),
        B, C, T, kernel_size, dilation, padding, Cout
    );
    return out;
}
'''

cpp_source = '''
torch::Tensor fused_forward(
    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_forward'],
    verbose=False,
    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.dilation = dilation
        self.kernel_size = kernel_size
        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, 1, bias=False)

    def forward(self, x):
        return module.fused_forward(
            x, self.dw.weight, self.pw.weight,
            self.padding, self.dilation, self.kernel_size
        )