import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, adds a bias term, scales, applies sigmoid, and performs group normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape)) 
        self.scale = nn.Parameter(torch.randn(scale_shape))
        self.group_norm = nn.GroupNorm(num_groups, out_channels)

    def forward(self, x):
        x = self.conv(x)
        x = x + self.bias
        x = x * self.scale
        x = torch.sigmoid(x)
        x = self.group_norm(x)
        return x

batch_size = 128
in_channels = 8
out_channels = 32
height = width = 256
kernel_size = 3
num_groups = 8
bias_shape = (out_channels, 1, 1)
scale_shape = (out_channels, 1, 1)

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

static __device__ __forceinline__ float sigmoidf_fast(float x) {
    return 1.0f / (1.0f + expf(-x));
}

__global__ void sigmoid_stats_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    float* __restrict__ y,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    int N, int C, int H, int W, int G, float eps)
{
    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng - n * G;
    int cpg = C / G;
    int HW = H * W;
    int K = cpg * HW;
    int c0 = g * cpg;
    size_t base_n = ((size_t)n) * C * HW;

    float sum = 0.0f;
    float sum2 = 0.0f;

    for (int i = threadIdx.x; i < K; i += blockDim.x) {
        int lc = i / HW;
        int hw = i - lc * HW;
        int c = c0 + lc;
        size_t idx = base_n + (size_t)c * HW + hw;
        float v = sigmoidf_fast((x[idx] + bias[c]) * scale[c]);
        y[idx] = v;
        sum += v;
        sum2 += v * v;
    }

    __shared__ float ssum[256];
    __shared__ float ssq[256];
    ssum[threadIdx.x] = sum;
    ssq[threadIdx.x] = sum2;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            ssum[threadIdx.x] += ssum[threadIdx.x + stride];
            ssq[threadIdx.x] += ssq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        float m = ssum[0] / (float)K;
        float ex2 = ssq[0] / (float)K;
        float var = ex2 - m * m;
        if (var < 0.0f) var = 0.0f;
        mean[ng] = m;
        invstd[ng] = rsqrtf(var + eps);
    }
}

__global__ void norm_affine_kernel(
    float* __restrict__ y,
    const float* __restrict__ mean,
    const float* __restrict__ invstd,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    int total, int C, int H, int W, int G)
{
    int HW = H * W;
    int cpg = C / G;
    int CHW = C * HW;
    for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) {
        int t = idx;
        int n = t / CHW;
        int rem = t - n * CHW;
        int c = rem / HW;
        int g = c / cpg;
        int ng = n * G + g;
        float v = y[idx];
        y[idx] = (v - mean[ng]) * invstd[ng] * gamma[c] + beta[c];
    }
}

torch::Tensor fused_post_gn(torch::Tensor x, torch::Tensor bias, torch::Tensor scale,
                            torch::Tensor gamma, torch::Tensor beta, int64_t groups, double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    auto xc = x.contiguous();
    auto bc = bias.contiguous();
    auto sc = scale.contiguous();
    auto gc = gamma.contiguous();
    auto betac = beta.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)groups;
    auto y = torch::empty_like(xc);
    auto opts = torch::TensorOptions().device(xc.device()).dtype(torch::kFloat32);
    auto mean = torch::empty({N * G}, opts);
    auto invstd = torch::empty({N * G}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int threads = 256;
    sigmoid_stats_kernel<<<N * G, threads, 0, stream>>>(
        xc.data_ptr<float>(), bc.data_ptr<float>(), sc.data_ptr<float>(), y.data_ptr<float>(),
        mean.data_ptr<float>(), invstd.data_ptr<float>(), N, C, H, W, G, (float)eps);
    int total = N * C * H * W;
    int blocks = (total + 255) / 256;
    if (blocks > 65535) blocks = 65535;
    norm_affine_kernel<<<blocks, 256, 0, stream>>>(
        y.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        gc.data_ptr<float>(), betac.data_ptr<float>(), total, C, H, W, G);
    return y;
}
'''

_ext = load_inline(
    name='fused_conv_post_groupnorm_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['fused_post_gn'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.scale = nn.Parameter(torch.randn(scale_shape))
        self.group_norm = nn.GroupNorm(num_groups, out_channels)
        self.num_groups = num_groups

    def forward(self, x):
        z = self.conv(x)
        return _ext.fused_post_gn(z, self.bias, self.scale, self.group_norm.weight, self.group_norm.bias, self.num_groups, self.group_norm.eps)