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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

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

__global__ void compute_stats_kernel(
    const float* __restrict__ input,
    float* __restrict__ mean_out,
    float* __restrict__ rstd_out,
    int N, int C, int HW, int G, int CpG, float eps
) {
    int ng = blockIdx.x * blockDim.x + threadIdx.x;
    int total_ng = N * G;
    if (ng >= total_ng) return;
    int n = ng / G;
    int g = ng % G;

    float sum = 0.0f;
    float sq_sum = 0.0f;
    int base = n * C * HW + g * CpG * HW;
    int count = CpG * HW;
    for (int i = 0; i < count; ++i) {
        float v = input[base + i];
        sum += v;
        sq_sum += v * v;
    }
    float mean = sum / count;
    float var = sq_sum / count - mean * mean;
    float rstd = rsqrtf(var + eps);
    mean_out[ng] = mean;
    rstd_out[ng] = rstd;
}

__global__ void apply_norm_silu_kernel(
    const float* __restrict__ input,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mean,
    const float* __restrict__ rstd,
    float* __restrict__ output,
    int N, int C, int HW, int G, int CpG
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * HW;
    if (idx >= total) return;

    int c = (idx / HW) % C;
    int g = c / CpG;
    int n = idx / (C * HW);
    int ng = n * G + g;
    float m = mean[ng];
    float r = rstd[ng];
    float g_ = gamma[c];
    float b = beta[c];
    float val = input[idx];
    float normed = (val - m) * r * g_ + b;
    output[idx] = normed * (1.0f / (1.0f + __expf(-normed)));
}

torch::Tensor groupnorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int num_groups, float eps) {
    auto N = x.size(0);
    auto C = x.size(1);
    auto H = x.size(2);
    auto W = x.size(3);
    int HW = H * W;
    int G = num_groups;
    int CpG = C / G;
    auto output = torch::empty_like(x);
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
    auto mean = torch::empty({N, G}, options);
    auto rstd = torch::empty({N, G}, options);

    int total_ng = N * G;
    compute_stats_kernel<<<(total_ng + 255) / 255, 255>>>(x.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), N, C, HW, G, CpG, eps);

    int total = N * C * HW;
    apply_norm_silu_kernel<<<(total + 255) / 255, 255>>>(x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(), output.data_ptr<float>(), N, C, HW, G, CpG);

    return output;
}
"""

cpp_source = r"""
torch::Tensor groupnorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int num_groups, float eps);
"""

module = load_inline(
    name="groupnorm_silu",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["groupnorm_silu_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

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

    def forward(self, x):
        gamma = self.norm.weight
        beta = self.norm.bias
        return module.groupnorm_silu_cuda(x.contiguous(), gamma, beta, self.num_groups, self.eps)