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 fused_groupnorm_silu_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    int N, int C, int H, int W, int G, int C_per_G, float eps) {
    
    // Each block handles one (n, g) pair
    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng % G;
    
    int HW = H * W;
    int group_size = C_per_G * HW;
    
    extern __shared__ float sdata[];
    
    // Compute mean
    float sum = 0.0f;
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c_in_g = i / HW;
        int hw = i % HW;
        int c = g * C_per_G + c_in_g;
        int idx = ((n * C + c) * HW) + hw;
        sum += input[idx];
    }
    
    sdata[threadIdx.x] = sum;
    __syncthreads();
    
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    float mean = sdata[0] / group_size;
    
    // Compute variance
    float var_sum = 0.0f;
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c_in_g = i / HW;
        int hw = i % HW;
        int c = g * C_per_G + c_in_g;
        int idx = ((n * C + c) * HW) + hw;
        float val = input[idx] - mean;
        var_sum += val * val;
    }
    
    sdata[threadIdx.x] = var_sum;
    __syncthreads();
    
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    float var = sdata[0] / group_size;
    float inv_std = rsqrtf(var + eps);
    
    // Normalize, apply affine and SiLU
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c_in_g = i / HW;
        int hw = i % HW;
        int c = g * C_per_G + c_in_g;
        int idx = ((n * C + c) * HW) + hw;
        float val = input[idx];
        float norm = (val - mean) * inv_std;
        if (gamma != nullptr && beta != nullptr) {
            norm = norm * gamma[c] + beta[c];
        }
        float silu = norm * (1.0f / (1.0f + expf(-norm)));
        output[idx] = silu;
    }
}

torch::Tensor fused_groupnorm_silu_cuda(
    torch::Tensor input,
    torch::Tensor gamma,
    torch::Tensor beta,
    int num_groups,
    float eps) {
    
    auto N = input.size(0);
    auto C = input.size(1);
    auto H = input.size(2);
    auto W = input.size(3);
    int G = num_groups;
    int C_per_G = C / G;
    
    auto output = torch::empty_like(input);
    
    int num_blocks = N * G;
    int threads = 256;
    
    fused_groupnorm_silu_kernel<<<num_blocks, threads, threads * sizeof(float)>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        gamma.defined() ? gamma.data_ptr<float>() : nullptr,
        beta.defined() ? beta.data_ptr<float>() : nullptr,
        N, C, H, W, G, C_per_G, eps
    );
    
    return output;
}
'''

cpp_source = r'''
torch::Tensor fused_groupnorm_silu_cuda(
    torch::Tensor input,
    torch::Tensor gamma,
    torch::Tensor beta,
    int num_groups,
    float eps);
'''

module = load_inline(
    name='fused_gn_silu',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_groupnorm_silu_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, num_channels: int, num_groups: int, eps: float = 1e-5, affine: bool = True):
        super().__init__()
        self.num_groups = num_groups
        self.eps = eps
        self.affine = affine
        self.weight = nn.Parameter(torch.ones(num_channels)) if affine else None
        self.bias = nn.Parameter(torch.zeros(num_channels)) if affine else None
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gamma = self.weight if self.weight is not None else torch.empty(0, device=x.device)
        beta = self.bias if self.bias is not None else torch.empty(0, device=x.device)
        return module.fused_groupnorm_silu_cuda(x, gamma, beta, self.num_groups, self.eps)