import torch
import torch.nn as nn

# 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 1D convolution + GroupNorm for temporal sequences (N, C, T).

    Typical in time-series heads, audio / motor trajectories, and sequence models
    that use strided or padded Conv1d followed by per-channel normalization without
    collapsing the time dimension (GroupNorm over `out_channels` on (N, C_out, T')).

    Reference fusion target: `Conv1d` → `GroupNorm` (no activation), matching a
    common stem or residual branch in temporal backbones.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_groups: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        dilation: int = 1,
        eps: float = 1e-5,
        affine: bool = True,
    ):
        super().__init__()
        assert out_channels % num_groups == 0
        self.conv = nn.Conv1d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=False,
        )
        self.norm = nn.GroupNorm(
            num_groups=num_groups,
            num_channels=out_channels,
            eps=eps,
            affine=affine,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (N, C_in, T) temporal feature map
        Returns:
            (N, C_out, T') with T' determined by conv stride / padding / kernel
        """
        return self.norm(self.conv(x))


class ModelNew(nn.Module):
    """
    KernelBench candidate entry point. Replace with a fused CUDA/Triton kernel;
    default delegates to `Model` for harness smoke tests.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_groups: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        dilation: int = 1,
        eps: float = 1e-5,
        affine: bool = True,
    ):
        super().__init__()
        self._impl = Model(
            in_channels,
            out_channels,
            num_groups,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            eps=eps,
            affine=affine,
        )

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


batch_size = 16
in_channels = 128
out_channels = 256
num_groups = 32
seq_len = 1024
kernel_size = 3
stride = 1
padding = 1


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,
        num_groups,
        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 fused_conv_gn(torch::Tensor x,
                            torch::Tensor weight,
                            torch::Tensor gamma,
                            torch::Tensor beta,
                            int64_t num_groups,
                            double eps,
                            int64_t stride,
                            int64_t padding,
                            int64_t dilation);
'''

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

// Specialized fast path for the benchmark's Conv1d parameters:
// kernel=3, stride=1, padding=1, dilation=1, fp32, contiguous NCT layout.
// One block computes a 16 (time) x 16 (output-channel) tile for one batch item.
__global__ void conv1d_k3s1p1_tile_kernel(const float* __restrict__ x,
                                           const float* __restrict__ w,
                                           float* __restrict__ y,
                                           int N, int Cin, int Cout, int T, int Tout) {
    const int tx = threadIdx.x;          // 0..15 time within tile
    const int ty = threadIdx.y;          // 0..15 output channel within tile
    const int tid = ty * 16 + tx;
    const int t0 = blockIdx.x * 16;
    const int co = blockIdx.y * 16 + ty;
    const int n = blockIdx.z;
    const int t = t0 + tx;

    __shared__ float sx[18];             // x[n, ci, t0-1 : t0+16]
    __shared__ float sw[16][3];          // w[co_tile, ci, k]

    float acc = 0.0f;
    for (int ci = 0; ci < Cin; ++ci) {
        if (tid < 18) {
            const int tt = t0 + tid - 1;
            float v = 0.0f;
            if ((unsigned)tt < (unsigned)T) {
                v = x[(n * Cin + ci) * T + tt];
            }
            sx[tid] = v;
        }
        if (tx < 3) {
            const int c = co;
            if (c < Cout) {
                sw[ty][tx] = w[(c * Cin + ci) * 3 + tx];
            }
        }
        __syncthreads();
        if (co < Cout && t < Tout) {
            acc = fmaf(sw[ty][0], sx[tx], acc);
            acc = fmaf(sw[ty][1], sx[tx + 1], acc);
            acc = fmaf(sw[ty][2], sx[tx + 2], acc);
        }
        __syncthreads();
    }
    if (co < Cout && t < Tout) {
        y[(n * Cout + co) * Tout + t] = acc;
    }
}

// Generic direct Conv1d fallback for unusual constructor arguments.  It is mainly
// for correctness outside the benchmark shape; the specialized path above is used
// for the provided case.
__global__ void conv1d_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      float* __restrict__ y,
                                      int N, int Cin, int Cout, int T, int K,
                                      int Tout, int stride, int padding, int dilation) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * Cout * Tout;
    for (; idx < total; idx += blockDim.x * gridDim.x) {
        int t = idx % Tout;
        int co = (idx / Tout) % Cout;
        int n = idx / (Tout * Cout);
        float acc = 0.0f;
        int base = t * stride - padding;
        for (int ci = 0; ci < Cin; ++ci) {
            #pragma unroll
            for (int k = 0; k < 8; ++k) {
                if (k >= K) break;
                int ti = base + k * dilation;
                if ((unsigned)ti < (unsigned)T) {
                    acc = fmaf(w[(co * Cin + ci) * K + k], x[(n * Cin + ci) * T + ti], acc);
                }
            }
            for (int k = 8; k < K; ++k) {
                int ti = base + k * dilation;
                if ((unsigned)ti < (unsigned)T) {
                    acc = fmaf(w[(co * Cin + ci) * K + k], x[(n * Cin + ci) * T + ti], acc);
                }
            }
        }
        y[idx] = acc;
    }
}

// One block per (sample, group).  Reduces over C_per_group * T, then writes the
// normalized and affine-transformed output.
__global__ void groupnorm_forward_kernel(const float* __restrict__ inp,
                                         const float* __restrict__ gamma,
                                         const float* __restrict__ beta,
                                         float* __restrict__ out,
                                         int N, int C, int T, int G,
                                         float eps, bool affine) {
    extern __shared__ float smem[];
    float* ssum = smem;
    float* ssq = smem + blockDim.x;
    const int tid = threadIdx.x;
    const int ng = blockIdx.x;
    const int n = ng / G;
    const int g = ng - n * G;
    const int cpg = C / G;
    const int elems = cpg * T;
    const int cbase = g * cpg;

    float sum = 0.0f;
    float sumsq = 0.0f;
    for (int i = tid; i < elems; i += blockDim.x) {
        int c = cbase + i / T;
        int t = i - (i / T) * T;
        float v = inp[(n * C + c) * T + t];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }
    ssum[tid] = sum;
    ssq[tid] = sumsq;
    __syncthreads();

    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
        if (tid < offset) {
            ssum[tid] += ssum[tid + offset];
            ssq[tid] += ssq[tid + offset];
        }
        __syncthreads();
    }

    const float mean = ssum[0] / (float)elems;
    float var = ssq[0] / (float)elems - mean * mean;
    var = var < 0.0f ? 0.0f : var;
    const float rstd = rsqrtf(var + eps);

    for (int i = tid; i < elems; i += blockDim.x) {
        int c_local = i / T;
        int c = cbase + c_local;
        int t = i - c_local * T;
        float v = inp[(n * C + c) * T + t];
        float z = (v - mean) * rstd;
        if (affine) z = z * gamma[c] + beta[c];
        out[(n * C + c) * T + t] = z;
    }
}

torch::Tensor fused_conv_gn(torch::Tensor x,
                            torch::Tensor weight,
                            torch::Tensor gamma,
                            torch::Tensor beta,
                            int64_t num_groups,
                            double eps,
                            int64_t stride,
                            int64_t padding,
                            int64_t dilation) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only fp32 is supported");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "only fp32 weights are supported");

    x = x.contiguous();
    weight = weight.contiguous();
    if (gamma.defined()) gamma = gamma.contiguous();
    if (beta.defined()) beta = beta.contiguous();

    const int N = (int)x.size(0);
    const int Cin = (int)x.size(1);
    const int T = (int)x.size(2);
    const int Cout = (int)weight.size(0);
    const int K = (int)weight.size(2);
    const int Tout = (int)((T + 2 * padding - dilation * (K - 1) - 1) / stride + 1);

    auto conv = torch::empty({N, Cout, Tout}, x.options());
    auto out = torch::empty({N, Cout, Tout}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && padding == 1 && dilation == 1) {
        dim3 block(16, 16, 1);
        dim3 grid((Tout + 15) / 16, (Cout + 15) / 16, N);
        conv1d_k3s1p1_tile_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), conv.data_ptr<float>(),
            N, Cin, Cout, T, Tout);
    } else {
        int total = N * Cout * Tout;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 4096) blocks = 4096;
        conv1d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), conv.data_ptr<float>(),
            N, Cin, Cout, T, K, Tout, (int)stride, (int)padding, (int)dilation);
    }

    const bool affine = gamma.defined() && gamma.numel() == Cout && beta.defined() && beta.numel() == Cout;
    const float* gptr = affine ? gamma.data_ptr<float>() : nullptr;
    const float* bptr = affine ? beta.data_ptr<float>() : nullptr;
    int threads = 256;
    size_t shmem = 2 * threads * sizeof(float);
    groupnorm_forward_kernel<<<N * (int)num_groups, threads, shmem, stream>>>(
        conv.data_ptr<float>(), gptr, bptr, out.data_ptr<float>(),
        N, Cout, Tout, (int)num_groups, (float)eps, affine);

    return out;
}
'''

_ext = load_inline(
    name="fused_conv1d_groupnorm_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_conv_gn"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_groups: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        dilation: int = 1,
        eps: float = 1e-5,
        affine: bool = True,
    ):
        super().__init__()
        assert out_channels % num_groups == 0
        self.conv = nn.Conv1d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=False,
        )
        self.norm = nn.GroupNorm(
            num_groups=num_groups,
            num_channels=out_channels,
            eps=eps,
            affine=affine,
        )
        self.num_groups = num_groups
        self.eps = eps
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.affine = affine

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.affine:
            gamma = self.norm.weight
            beta = self.norm.bias
        else:
            gamma = torch.empty(0, device=x.device, dtype=x.dtype)
            beta = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.fused_conv_gn(
            x,
            self.conv.weight,
            gamma,
            beta,
            self.num_groups,
            float(self.eps),
            int(self.stride),
            int(self.padding),
            int(self.dilation),
        )