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 = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void fused_groupnorm_silu_kernel(
    const float* __restrict__ input,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ output,
    int N, int C, int H, int W,
    int num_groups, float eps) {
    
    int nc = blockIdx.x; // n * num_groups + g
    int n = nc / num_groups;
    int g = nc % num_groups;
    
    int group_size = C / num_groups * H * W;
    int c_start = g * (C / num_groups);
    
    // Compute mean and var for this group
    float sum = 0.0f;
    float sum_sq = 0.0f;
    
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c = c_start + (i / (H*W));
        int hw = i % (H*W);
        int idx = ((n * C + c) * H * W) + hw;
        float val = input[idx];
        sum += val;
        sum_sq += val * val;
    }
    
    // Block reduce
    __shared__ float shared_sum[256];
    __shared__ float shared_sq[256];
    
    shared_sum[threadIdx.x] = sum;
    shared_sq[threadIdx.x] = sum_sq;
    __syncthreads();
    
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            shared_sum[threadIdx.x] += shared_sum[threadIdx.x + s];
            shared_sq[threadIdx.x] += shared_sq[threadIdx.x + s];
        }
        __syncthreads();
    }
    
    float mean = shared_sum[0] / group_size;
    float var = shared_sq[0] / group_size - mean * mean;
    float inv_std = rsqrtf(var + eps);
    
    // Normalize, apply gamma/beta, SiLU
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c = c_start + (i / (H*W));
        int hw = i % (H*W);
        int idx = ((n * C + c) * H * W) + hw;
        float val = input[idx];
        float normed = (val - mean) * inv_std;
        float out_val = normed * gamma[c] + beta[c];
        // SiLU
        out_val = out_val / (1.0f + expf(-out_val));
        output[idx] = out_val;
    }
}

torch::Tensor fused_groupnorm_silu(
    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);
    
    auto output = torch::empty_like(input);
    
    int num_blocks = N * num_groups;
    fused_groupnorm_silu_kernel<<<num_blocks, 256>>>(
        input.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, H, W,
        num_groups, eps);
    
    return output;
}
'''

cpp_source = '''
torch::Tensor fused_groupnorm_silu(
    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'],
    verbose=False
)

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_groups, num_channels=num_channels, eps=eps, affine=affine)
        self.num_groups = num_groups
        self.eps = eps
        self.fused = module
    
    def forward(self, x):
        gamma = self.norm.weight
        beta = self.norm.bias
        return self.fused.fused_groupnorm_silu(x, gamma, beta, self.num_groups, self.eps)