import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, applies GELU, and normalizes with GroupNorm.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, groups, num_groups):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride)
        self.group_norm = nn.GroupNorm(num_groups=num_groups, num_channels=out_channels)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.nn.functional.gelu(x)
        x = self.group_norm(x)
        return x

batch_size   = 128  
in_channels  = 64  
out_channels = 64  
height = width = 256  
kernel_size  = 3
stride       = 1
groups = 8
num_groups = 8

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, groups, num_groups]

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

cpp_src = r'''
#include <torch/extension.h>
#include <vector>

void gelu_gn_launcher(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, torch::Tensor output, int num_groups, double eps);

torch::Tensor fused_gelu_groupnorm(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, int64_t num_groups, double eps) {
    auto output = torch::empty_like(input);
    gelu_gn_launcher(input, gamma, beta, output, (int)num_groups, eps);
    return output;
}
'''

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

#define THREADS_RED 1024
#define THREADS_NORM 256

__device__ __forceinline__ float gelu_exact(float x) {
    return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f));
}

__inline__ __device__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__inline__ __device__ float block_reduce_sum(float val) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    val = warp_reduce_sum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) val = warp_reduce_sum(val);
    return val;
}

__global__ void gelu_reduce_kernel(float* __restrict__ x, float* __restrict__ mean, float* __restrict__ invstd,
                                   int N, int C, int S, int G, float eps) {
    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng - n * G;
    int group_c = C / G;
    int64_t base = ((int64_t)n * C + g * group_c) * (int64_t)S;
    int M = group_c * S;
    float sum = 0.0f;
    float sumsq = 0.0f;

    for (int i = threadIdx.x; i < M; i += blockDim.x) {
        float v = x[base + i];
        float gv = gelu_exact(v);
        x[base + i] = gv;
        sum += gv;
        sumsq = fmaf(gv, gv, sumsq);
    }

    float bsum = block_reduce_sum(sum);
    float bsq = block_reduce_sum(sumsq);
    if (threadIdx.x == 0) {
        float m = bsum / (float)M;
        float var = bsq / (float)M - m * m;
        mean[ng] = m;
        invstd[ng] = rsqrtf((var < 0.0f ? 0.0f : var) + eps);
    }
}

__global__ void norm_nc_kernel(const float* __restrict__ x, const float* __restrict__ gamma, const float* __restrict__ beta,
                               const float* __restrict__ mean, const float* __restrict__ invstd, float* __restrict__ out,
                               int C, int S, int G) {
    int nc = blockIdx.x;
    int c = nc % C;
    int n = nc / C;
    int group_c = C / G;
    int ng = n * G + c / group_c;
    float m = mean[ng];
    float r = invstd[ng];
    float w = gamma[c];
    float b = beta[c];
    int64_t base = (int64_t)nc * (int64_t)S;

    for (int s = threadIdx.x; s < S; s += blockDim.x) {
        float y = (x[base + s] - m) * r;
        out[base + s] = y * w + b;
    }
}

void gelu_gn_launcher(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, torch::Tensor output, int num_groups, double eps) {
    const int N = input.size(0);
    const int C = input.size(1);
    const int H = input.size(2);
    const int W = input.size(3);
    const int S = H * W;

    auto mean = torch::empty({N * num_groups}, input.options());
    auto invstd = torch::empty({N * num_groups}, input.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    gelu_reduce_kernel<<<N * num_groups, THREADS_RED, 0, stream>>>(
        input.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), N, C, S, num_groups, (float)eps);
    norm_nc_kernel<<<N * C, THREADS_NORM, 0, stream>>>(
        input.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), output.data_ptr<float>(), C, S, num_groups);
}
'''

_ext = load_inline(
    name='fused_gelu_gn_ext_final_a',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['fused_gelu_groupnorm'],
    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, stride, groups, num_groups):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride)
        self.group_norm = nn.GroupNorm(num_groups=num_groups, num_channels=out_channels)
        self.num_groups = num_groups

    def forward(self, x):
        w = torch.flip(self.conv_transpose.weight.permute(1, 0, 2, 3), [2, 3]).contiguous()
        y = F.conv2d(x, w, self.conv_transpose.bias, stride=1, padding=2)
        return _ext.fused_gelu_groupnorm(y, self.group_norm.weight, self.group_norm.bias, self.num_groups, self.group_norm.eps)