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

# Default ModelNew is PyTorch-only. run_and_check with check_kernel=True and backend=cuda
# requires real CUDA (__global__ + load_inline/cpp_extension) in the kernel file—use
# check_kernel=False until you add a custom kernel implementation.

class Model(nn.Module):
    """
    Fused GroupNorm + SiLU (Swish), a common post-conv / residual-branch pattern in
    ResNet-family backbones (SiLU instead of ReLU) and in spatial blocks of diffusion
    transformers (DiT) where features are (N, C, H, W).

    Reference pipeline for a single fused kernel: normalize channels with
    `GroupNorm`, then apply `SiLU` elementwise (`x * sigmoid(x)`).
    """

    def __init__(
        self,
        num_channels: int,
        num_groups: int,
        eps: float = 1e-5,
        affine: bool = True,
    ):
        super().__init__()
        assert num_channels % num_groups == 0
        self.norm = nn.GroupNorm(
            num_groups=num_groups,
            num_channels=num_channels,
            eps=eps,
            affine=affine,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (N, C, H, W) feature map
        Returns:
            Same shape as x
        """
        return F.silu(self.norm(x))


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

    def __init__(
        self,
        num_channels: int,
        num_groups: int,
        eps: float = 1e-5,
        affine: bool = True,
    ):
        super().__init__()
        self._impl = Model(num_channels, num_groups, eps=eps, affine=affine)

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


batch_size = 8
num_channels = 256
num_groups = 32
height = 24
width = 24


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), num_channels, p.jitter_int(height), p.jitter_int(width))
    return [x]


def get_init_inputs():
    return [num_channels, num_groups]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor groupnorm_silu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int64_t num_groups, double eps, bool affine);
'''

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

__inline__ __device__ float warp_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);
    return v;
}

__inline__ __device__ void block_reduce_pair(float &s, float &q) {
    __shared__ float shs[32];
    __shared__ float shq[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    s = warp_sum(s);
    q = warp_sum(q);
    if (lane == 0) { shs[wid] = s; shq[wid] = q; }
    __syncthreads();
    int nw = blockDim.x >> 5;
    s = (threadIdx.x < nw) ? shs[lane] : 0.0f;
    q = (threadIdx.x < nw) ? shq[lane] : 0.0f;
    if (wid == 0) { s = warp_sum(s); q = warp_sum(q); }
}

__device__ __forceinline__ float silu_f(float v) {
    return v / (1.0f + expf(-v));
}

__global__ void gn_silu_256_32_24_affine_kernel(const float* __restrict__ x,
                                                const float* __restrict__ weight,
                                                const float* __restrict__ bias,
                                                float* __restrict__ y,
                                                float eps) {
    int bg = blockIdx.x;
    int g = bg & 31;
    int base = bg * 4608;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x + base);
    float s = 0.0f, q = 0.0f;
    for (int vi = threadIdx.x; vi < 1152; vi += blockDim.x) {
        float4 a = x4[vi];
        s += a.x + a.y + a.z + a.w;
        q += a.x*a.x + a.y*a.y + a.z*a.z + a.w*a.w;
    }
    block_reduce_pair(s, q);
    __shared__ float sh_mean, sh_rstd;
    if (threadIdx.x == 0) {
        float mean = s * (1.0f / 4608.0f);
        float var = q * (1.0f / 4608.0f) - mean * mean;
        if (var < 0.0f) var = 0.0f;
        sh_mean = mean;
        sh_rstd = rsqrtf(var + eps);
    }
    __syncthreads();
    float mean = sh_mean;
    float rstd = sh_rstd;
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y + base);
    for (int vi = threadIdx.x; vi < 1152; vi += blockDim.x) {
        float4 a = x4[vi];
        int c = g * 8 + (vi / 144);
        float w = weight[c];
        float b = bias[c];
        float4 o;
        float v;
        v = (a.x - mean) * rstd; v = v * w + b; o.x = silu_f(v);
        v = (a.y - mean) * rstd; v = v * w + b; o.y = silu_f(v);
        v = (a.z - mean) * rstd; v = v * w + b; o.z = silu_f(v);
        v = (a.w - mean) * rstd; v = v * w + b; o.w = silu_f(v);
        y4[vi] = o;
    }
}

__global__ void groupnorm_silu_kernel(const float* __restrict__ x,
                                      const float* __restrict__ weight,
                                      const float* __restrict__ bias,
                                      float* __restrict__ y,
                                      int N, int C, int H, int W, int G,
                                      float eps, int affine) {
    int bg = blockIdx.x;
    int n = bg / G;
    int g = bg - n * G;
    int HW = H * W;
    int cpg = C / G;
    int m = cpg * HW;
    int base = ((n * C) + g * cpg) * HW;
    float s = 0.0f, q = 0.0f;
    for (int i = threadIdx.x; i < m; i += blockDim.x) { float v = x[base + i]; s += v; q += v * v; }
    block_reduce_pair(s, q);
    __shared__ float sh_mean, sh_rstd;
    if (threadIdx.x == 0) {
        float inv_m = 1.0f / (float)m;
        float mean = s * inv_m;
        float var = q * inv_m - mean * mean;
        if (var < 0.0f) var = 0.0f;
        sh_mean = mean; sh_rstd = rsqrtf(var + eps);
    }
    __syncthreads();
    float mean = sh_mean, rstd = sh_rstd;
    for (int i = threadIdx.x; i < m; i += blockDim.x) {
        int cl = i / HW;
        int c = g * cpg + cl;
        float v = (x[base + i] - mean) * rstd;
        if (affine) v = v * weight[c] + bias[c];
        y[base + i] = silu_f(v);
    }
}

torch::Tensor groupnorm_silu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int64_t num_groups, double eps, bool affine) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
    auto y = torch::empty_like(x);
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int G = (int)num_groups;
    const float* wptr = affine ? weight.data_ptr<float>() : nullptr;
    const float* bptr = affine ? bias.data_ptr<float>() : nullptr;
    if (affine && C == 256 && G == 32 && H == 24 && W == 24) {
        gn_silu_256_32_24_affine_kernel<<<N * 32, 1024, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), wptr, bptr, y.data_ptr<float>(), (float)eps);
    } else {
        groupnorm_silu_kernel<<<N * G, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), wptr, bptr, y.data_ptr<float>(), N, C, H, W, G, (float)eps, affine ? 1 : 0);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='gn_silu_ext_submit',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['groupnorm_silu_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_channels: int, num_groups: int, eps: float = 1e-5, affine: bool = True):
        super().__init__()
        assert num_channels % num_groups == 0
        self.num_groups = num_groups
        self.eps = eps
        self.affine = affine
        self.norm = nn.GroupNorm(num_groups=num_groups, num_channels=num_channels, eps=eps, affine=affine)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.affine:
            return _ext.groupnorm_silu_forward(x, self.norm.weight, self.norm.bias, self.num_groups, self.eps, True)
        empty = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.groupnorm_silu_forward(x, empty, empty, self.num_groups, self.eps, False)