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_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t num_groups, double eps, bool affine);

torch::Tensor groupnorm_silu_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t num_groups, double eps, bool affine) {
    return groupnorm_silu_forward_cuda(x, weight, bias, num_groups, eps, affine);
}
"""

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

static inline int next_pow2_host(int x) {
    int p = 1;
    while (p < x) p <<= 1;
    return p;
}

__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 HW, int G, int cpg, int group_size,
    float eps, bool affine) {

    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng - n * G;
    int tid = threadIdx.x;

    int c0 = g * cpg;
    int base = (n * C + c0) * HW;

    float sum = 0.0f;
    float sumsq = 0.0f;

    for (int idx = tid; idx < group_size; idx += blockDim.x) {
        float v = x[base + idx];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }

    __shared__ float s_sum[1024];
    __shared__ float s_sumsq[1024];
    s_sum[tid] = sum;
    s_sumsq[tid] = sumsq;
    __syncthreads();

    // blockDim is power-of-two <= 1024
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_sum[tid] += s_sum[tid + stride];
            s_sumsq[tid] += s_sumsq[tid + stride];
        }
        __syncthreads();
    }

    float mean = s_sum[0] / (float)group_size;
    float var = s_sumsq[0] / (float)group_size - mean * mean;
    if (var < 0.0f) var = 0.0f;
    float inv_std = rsqrtf(var + eps);

    for (int idx = tid; idx < group_size; idx += blockDim.x) {
        int rel_c = idx / HW;
        int c = c0 + rel_c;
        float v = (x[base + idx] - mean) * inv_std;
        if (affine) {
            v = fmaf(v, weight[c], bias[c]);
        }
        float sig = 1.0f / (1.0f + expf(-v));
        y[base + idx] = v * sig;
    }
}

torch::Tensor groupnorm_silu_forward_cuda(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() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");

    auto xc = x.contiguous();
    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int H = (int)xc.size(2);
    int W = (int)xc.size(3);
    int G = (int)num_groups;
    int HW = H * W;
    int cpg = C / G;
    int group_size = cpg * HW;

    auto y = torch::empty_like(xc);

    int threads = next_pow2_host(group_size);
    if (threads > 1024) threads = 1024;
    if (threads < 32) threads = 32;

    const float* wptr = affine ? weight.data_ptr<float>() : nullptr;
    const float* bptr = affine ? bias.data_ptr<float>() : nullptr;

    dim3 grid(N * G);
    groupnorm_silu_kernel<<<grid, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        xc.data_ptr<float>(), wptr, bptr, y.data_ptr<float>(),
        N, C, HW, G, cpg, group_size, (float)eps, affine);

    return y;
}
"""

_ext = load_inline(
    name="groupnorm_silu_fused_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["groupnorm_silu_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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_channels = num_channels
        self.num_groups = num_groups
        self.eps = eps
        self.affine = affine
        # Keep the same public parameter names as nn.GroupNorm for harness weight copying.
        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)
        else:
            empty = torch.empty(0, device=x.device, dtype=x.dtype)
            return _ext.groupnorm_silu_forward(x, empty, empty, self.num_groups, self.eps, False)