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.utils.cpp_extension as cpp_extension

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

// Fused GroupNorm + SiLU kernel
// Each thread block processes one (batch, group) pair
// Threads within block cooperatively compute mean, variance, then apply normalization + SiLU

template <int BLOCK_SIZE>
__global__ void fused_groupnorm_silu_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    int N, int C, int H, int W, int G, int channels_per_group, float eps
) {
    // Each block handles one (batch_item, group) pair
    int batch_idx = blockIdx.x / G;
    int group_idx = blockIdx.x % G;
    
    int group_start_channel = group_idx * channels_per_group;
    int group_size = channels_per_group * H * W;
    
    // Index for this thread's element within the group
    int tid = threadIdx.x;
    int total_elements = group_size;
    
    // Pointers to the start of this batch and group in input/output
    int batch_offset = batch_idx * C * H * W;
    int group_offset = group_start_channel * H * W;
    const float* __restrict__ input_ptr = input + batch_offset + group_offset;
    float* __restrict__ output_ptr = output + batch_offset + group_offset;
    
    // Shared memory for reduction
    __shared__ float s_mean[BLOCK_SIZE];
    __shared__ float s_var[BLOCK_SIZE];
    
    // Compute mean using parallel reduction
    float local_sum = 0.0f;
    float local_sq_sum = 0.0f;
    
    // Grid-stride loop to process all elements in the group
    for (int i = tid; i < total_elements; i += BLOCK_SIZE) {
        float val = input_ptr[i];
        local_sum += val;
        local_sq_sum += val * val;
    }
    
    s_mean[tid] = local_sum;
    s_var[tid] = local_sq_sum;
    __syncthreads();
    
    // Parallel reduction in shared memory
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_mean[tid] += s_mean[tid + stride];
            s_var[tid] += s_var[tid + stride];
        }
        __syncthreads();
    }
    
    // Compute mean and variance
    float mean = 0.0f;
    float inv_std = 0.0f;
    if (tid == 0) {
        mean = s_mean[0] / total_elements;
        float variance = s_var[0] / total_elements - mean * mean;
        inv_std = 1.0f / sqrtf(variance + eps);
        s_mean[0] = mean;
        s_var[0] = inv_std;
    }
    __syncthreads();
    
    mean = s_mean[0];
    inv_std = s_var[0];
    
    // Apply normalization, affine, and SiLU
    for (int i = tid; i < total_elements; i += BLOCK_SIZE) {
        int c_local = i / (H * W);
        int c_global = group_start_channel + c_local;
        
        float val = input_ptr[i];
        float norm_val = (val - mean) * inv_std;
        
        // Apply affine transformation if weight/bias are provided
        if (weight != nullptr && bias != nullptr) {
            norm_val = weight[c_global] * norm_val + bias[c_global];
        }
        
        // Apply SiLU: x * sigmoid(x) = x / (1 + exp(-x))
        float silu_val = norm_val / (1.0f + expf(-norm_val));
        
        output_ptr[i] = silu_val;
    }
}

at::Tensor fused_groupnorm_silu_cuda(
    at::Tensor input,
    at::Tensor weight,
    at::Tensor bias,
    int num_groups,
    double eps
) {
    int N = input.size(0);
    int C = input.size(1);
    int H = input.size(2);
    int W = input.size(3);
    int channels_per_group = C / num_groups;
    
    auto output = at::empty_like(input);
    
    int num_blocks = N * num_groups;
    int block_size = 256;
    
    const float* input_ptr = input.data_ptr<float>();
    float* output_ptr = output.data_ptr<float>();
    const float* weight_ptr = weight.defined() ? weight.data_ptr<float>() : nullptr;
    const float* bias_ptr = bias.defined() ? bias.data_ptr<float>() : nullptr;
    
    fused_groupnorm_silu_kernel<256><<<num_blocks, block_size>>>(
        input_ptr, output_ptr, weight_ptr, bias_ptr,
        N, C, H, W, num_groups, channels_per_group, static_cast<float>(eps)
    );
    
    return output;
}
"""

CPP_SOURCE = r"""
at::Tensor fused_groupnorm_silu_cuda(
    at::Tensor input,
    at::Tensor weight,
    at::Tensor bias,
    int num_groups,
    double eps
);
"""

# Compile the extension
fused_gn_silu = cpp_extension.load_inline(
    name='fused_gn_silu',
    cpp_sources=CPP_SOURCE,
    cuda_sources=CUDA_SOURCE,
    functions=['fused_groupnorm_silu_cuda'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
)

import torch
import torch.nn as nn

class ModelNew(nn.Module):
    """
    Fused GroupNorm + SiLU using custom CUDA kernel.
    """
    
    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
        
        if affine:
            self.weight = nn.Parameter(torch.ones(num_channels))
            self.bias = nn.Parameter(torch.zeros(num_channels))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (N, C, H, W) feature map
        Returns:
            Same shape as x with GroupNorm + SiLU applied
        """
        # Ensure contiguous input
        if not x.is_contiguous():
            x = x.contiguous()
        
        return fused_gn_silu.fused_groupnorm_silu_cuda(
            x,
            self.weight if self.weight is not None else torch.empty(0),
            self.bias if self.bias is not None else torch.empty(0),
            self.num_groups,
            self.eps
        )