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 <cuda.h>
#include <cuda_runtime.h>

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

template<int BLOCK>
__inline__ __device__ float block_reduce_sum(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    float val = (threadIdx.x < (BLOCK >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) val = warp_reduce_sum(val);
    if (threadIdx.x == 0) shared[0] = val;
    __syncthreads();
    val = shared[0];
    __syncthreads();
    return val;
}
__device__ __forceinline__ float silu_f(float v) { return v / (1.0f + expf(-v)); }

template<int BLOCK, bool AFFINE>
__global__ void gn_silu_fixed_kernel(const float* __restrict__ x, const float* __restrict__ weight, const float* __restrict__ bias, float* __restrict__ y, float eps) {
    int gid = blockIdx.x;
    int n = gid >> 5;
    int g = gid & 31;
    int c_start = g << 3;
    int base4 = n * 36864 + g * 1152;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
    float sum = 0.0f, sumsq = 0.0f;
    #pragma unroll
    for (int i = threadIdx.x; i < 1152; i += BLOCK) {
        float4 r = x4[base4 + i];
        sum += r.x + r.y + r.z + r.w;
        sumsq = fmaf(r.x, r.x, sumsq);
        sumsq = fmaf(r.y, r.y, sumsq);
        sumsq = fmaf(r.z, r.z, sumsq);
        sumsq = fmaf(r.w, r.w, sumsq);
    }
    float total = block_reduce_sum<BLOCK>(sum);
    float total2 = block_reduce_sum<BLOCK>(sumsq);
    float mean = total * (1.0f / 4608.0f);
    float invstd = rsqrtf(fmaxf(total2 * (1.0f / 4608.0f) - mean * mean, 0.0f) + eps);

    if (threadIdx.x < 144) {
        #pragma unroll
        for (int cl = 0; cl < 8; ++cl) {
            float ww = 1.0f, bb = 0.0f;
            if (AFFINE) { ww = weight[c_start + cl]; bb = bias[c_start + cl]; }
            int idx = base4 + cl * 144 + threadIdx.x;
            float4 r = x4[idx];
            r.x = fmaf((r.x - mean) * invstd, ww, bb);
            r.y = fmaf((r.y - mean) * invstd, ww, bb);
            r.z = fmaf((r.z - mean) * invstd, ww, bb);
            r.w = fmaf((r.w - mean) * invstd, ww, bb);
            r.x = silu_f(r.x); r.y = silu_f(r.y); r.z = silu_f(r.z); r.w = silu_f(r.w);
            y4[idx] = r;
        }
    }
}

template<int BLOCK, bool AFFINE>
__global__ void gn_silu_vec4_kernel(const float* __restrict__ x, const float* __restrict__ weight, const float* __restrict__ bias,
                                    float* __restrict__ y, int C, int H, int W, int G,
                                    int elems_per_group, int channels_per_group, float eps) {
    int group_id = blockIdx.x;
    int n = group_id / G;
    int g = group_id - n * G;
    int HW = H * W;
    int c_start = g * channels_per_group;
    int base = ((n * C + c_start) * H) * W;
    int base4 = base >> 2;
    int e4 = elems_per_group >> 2;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(x);
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y);
    float sum = 0.0f, sumsq = 0.0f;
    for (int i = threadIdx.x; i < e4; i += BLOCK) { float4 r = x4[base4 + i]; sum += r.x + r.y + r.z + r.w; sumsq = fmaf(r.x, r.x, sumsq); sumsq = fmaf(r.y, r.y, sumsq); sumsq = fmaf(r.z, r.z, sumsq); sumsq = fmaf(r.w, r.w, sumsq); }
    float total = block_reduce_sum<BLOCK>(sum); float total2 = block_reduce_sum<BLOCK>(sumsq); float mean = total / (float)elems_per_group; float invstd = rsqrtf(fmaxf(total2 / (float)elems_per_group - mean * mean, 0.0f) + eps);
    int HW4 = HW >> 2;
    for (int cl = 0; cl < channels_per_group; ++cl) { float ww = 1.0f, bb = 0.0f; if (AFFINE) { ww = weight[c_start + cl]; bb = bias[c_start + cl]; } int off4 = base4 + cl * HW4; for (int j = threadIdx.x; j < HW4; j += BLOCK) { float4 r = x4[off4 + j]; r.x = fmaf((r.x - mean) * invstd, ww, bb); r.y = fmaf((r.y - mean) * invstd, ww, bb); r.z = fmaf((r.z - mean) * invstd, ww, bb); r.w = fmaf((r.w - mean) * invstd, ww, bb); r.x = silu_f(r.x); r.y = silu_f(r.y); r.z = silu_f(r.z); r.w = silu_f(r.w); y4[off4 + j] = r; } }
}

torch::Tensor groupnorm_silu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t num_groups, double eps, bool affine) {
    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;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int BLOCK = 256;
    const float* w = affine ? weight.data_ptr<float>() : nullptr;
    const float* b = affine ? bias.data_ptr<float>() : nullptr;
    if (C == 256 && H == 24 && W == 24 && G == 32) {
        if (affine) gn_silu_fixed_kernel<BLOCK, true><<<N * 32, BLOCK, 0, stream>>>(x.data_ptr<float>(), w, b, y.data_ptr<float>(), (float)eps);
        else gn_silu_fixed_kernel<BLOCK, false><<<N * 32, BLOCK, 0, stream>>>(x.data_ptr<float>(), w, b, y.data_ptr<float>(), (float)eps);
    } else {
        int cpg = C / G; int elems = cpg * H * W;
        if (affine) gn_silu_vec4_kernel<BLOCK, true><<<N * G, BLOCK, 0, stream>>>(x.data_ptr<float>(), w, b, y.data_ptr<float>(), C,H,W,G,elems,cpg,(float)eps);
        else gn_silu_vec4_kernel<BLOCK, false><<<N * G, BLOCK, 0, stream>>>(x.data_ptr<float>(), w, b, y.data_ptr<float>(), C,H,W,G,elems,cpg,(float)eps);
    }
    return y;
}
'''
_ext = load_inline(name='gn_silu_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['groupnorm_silu_forward'], verbose=False, extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'])
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)
        d=torch.empty(0,device=x.device,dtype=x.dtype); return _ext.groupnorm_silu_forward(x,d,d,self.num_groups,self.eps,False)