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
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <vector>

torch::Tensor gelu_gn_cuda(torch::Tensor y, torch::Tensor weight, torch::Tensor bias, int64_t num_groups);

torch::Tensor model_forward(torch::Tensor x,
                            torch::Tensor conv_w,
                            torch::Tensor conv_b,
                            torch::Tensor gn_w,
                            torch::Tensor gn_b,
                            int64_t num_groups,
                            int64_t stride) {
    // ConvTranspose2d parameters match the reference model: padding=0,
    // output_padding=0, dilation=1, groups=1.
    auto y = at::conv_transpose2d(x, conv_w, conv_b,
                                  {stride, stride}, {0, 0}, {0, 0},
                                  1, {1, 1});
    return gelu_gn_cuda(y, gn_w, gn_b, num_groups);
}
"""

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

#define BLOCK 256

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

__global__ void gelu_stats_kernel(float* __restrict__ y,
                                  float* __restrict__ sums,
                                  float* __restrict__ sqsums,
                                  int N, int C, int H, int W, int G) {
    int gid = blockIdx.x;              // one block per (n, group)
    int g = gid % G;
    int n = gid / G;
    int cpg = C / G;
    int HW = H * W;
    int M = cpg * HW;
    long long base = ((long long)n * C + (long long)g * cpg) * HW;

    float local_sum = 0.0f;
    float local_sq = 0.0f;

    for (int i = threadIdx.x; i < M; i += BLOCK) {
        float v = y[base + i];
        float gv = gelu_exact(v);
        y[base + i] = gv;
        local_sum += gv;
        local_sq += gv * gv;
    }

    __shared__ float sh_sum[BLOCK];
    __shared__ float sh_sq[BLOCK];
    sh_sum[threadIdx.x] = local_sum;
    sh_sq[threadIdx.x] = local_sq;
    __syncthreads();

    for (int offset = BLOCK / 2; offset > 0; offset >>= 1) {
        if (threadIdx.x < offset) {
            sh_sum[threadIdx.x] += sh_sum[threadIdx.x + offset];
            sh_sq[threadIdx.x] += sh_sq[threadIdx.x + offset];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        sums[gid] = sh_sum[0];
        sqsums[gid] = sh_sq[0];
    }
}

__global__ void normalize_kernel(float* __restrict__ y,
                                 const float* __restrict__ sums,
                                 const float* __restrict__ sqsums,
                                 const float* __restrict__ gamma,
                                 const float* __restrict__ beta,
                                 int N, int C, int H, int W, int G) {
    int gid = blockIdx.x;              // one block per (n, group)
    int g = gid % G;
    int n = gid / G;
    int cpg = C / G;
    int HW = H * W;
    int M = cpg * HW;
    long long base = ((long long)n * C + (long long)g * cpg) * HW;

    float invM = 1.0f / (float)M;
    float mean = sums[gid] * invM;
    float var = sqsums[gid] * invM - mean * mean;
    var = var < 0.0f ? 0.0f : var;
    float rstd = rsqrtf(var + 1.0e-5f);

    for (int i = threadIdx.x; i < M; i += BLOCK) {
        int local_c = i / HW;
        int c = g * cpg + local_c;
        float v = y[base + i];
        y[base + i] = (v - mean) * rstd * gamma[c] + beta[c];
    }
}

torch::Tensor gelu_gn_cuda(torch::Tensor y, torch::Tensor weight, torch::Tensor bias, int64_t num_groups) {
    const int N = (int)y.size(0);
    const int C = (int)y.size(1);
    const int H = (int)y.size(2);
    const int W = (int)y.size(3);
    const int G = (int)num_groups;
    auto stats = torch::empty({2, N * G}, y.options());
    float* sums = stats.data_ptr<float>();
    float* sqsums = sums + N * G;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 grid(N * G);
    dim3 block(BLOCK);
    gelu_stats_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), sums, sqsums, N, C, H, W, G);
    normalize_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), sums, sqsums,
                                                 weight.data_ptr<float>(), bias.data_ptr<float>(),
                                                 N, C, H, W, G);
    return y;
}
"""

_ext = load_inline(
    name="convtranspose_gelu_groupnorm_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["model_forward"],
    extra_cflags=["-O3"],
    extra_cuda_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.stride = int(stride)
        self.num_groups = int(num_groups)

    def forward(self, x):
        return _ext.model_forward(x, self.conv_transpose.weight, self.conv_transpose.bias,
                                  self.group_norm.weight, self.group_norm.bias,
                                  self.num_groups, self.stride)