import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, batch normalization, tanh activation, max pooling, and group normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups, num_groups):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm2d(out_channels)
        self.tanh = nn.Tanh()
        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.group_norm = nn.GroupNorm(num_groups=num_groups, num_channels=out_channels)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.tanh(x)
        x = self.max_pool(x)
        x = self.group_norm(x)
        return x

batch_size = 512
in_channels  = 64  
out_channels = 128  
height = width = 2048  
kernel_size  = 5
stride       = 1  
padding      = 1
groups       = 8
num_groups   = 8
height, width = 32, 32

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, 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>

torch::Tensor fused_post(torch::Tensor conv,
                         torch::Tensor bn_weight,
                         torch::Tensor bn_bias,
                         torch::Tensor running_mean,
                         torch::Tensor running_var,
                         torch::Tensor gn_weight,
                         torch::Tensor gn_bias,
                         bool bn_training,
                         double bn_eps,
                         int64_t num_groups,
                         double gn_eps);
"""

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

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

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

__global__ void bn_stats_kernel(const float* __restrict__ x, float* __restrict__ mean, float* __restrict__ invstd,
                                int N, int C, int H, int W, float eps) {
    int c = blockIdx.x;
    int HW = H * W;
    int M = N * HW;
    float sum = 0.0f;
    float sum2 = 0.0f;
    for (int i = threadIdx.x; i < M; i += blockDim.x) {
        int n = i / HW;
        int hw = i - n * HW;
        float v = x[((long long)n * C + c) * HW + hw];
        sum += v;
        sum2 = fmaf(v, v, sum2);
    }
    sum = block_reduce_sum(sum);
    sum2 = block_reduce_sum(sum2);
    if (threadIdx.x == 0) {
        float m = sum / (float)M;
        float var = sum2 / (float)M - m * m;
        var = var < 0.0f ? 0.0f : var;
        mean[c] = m;
        invstd[c] = rsqrtf(var + eps);
    }
}

__global__ void pool_bn_tanh_kernel(const float* __restrict__ x,
                                    const float* __restrict__ bn_weight,
                                    const float* __restrict__ bn_bias,
                                    const float* __restrict__ mean,
                                    const float* __restrict__ invstd,
                                    float* __restrict__ pooled,
                                    int total, int C, int H, int W, int PH, int PW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    int PHW = PH * PW;
    int HW = H * W;
    for (; idx < total; idx += stride) {
        int t = idx;
        int pw = t % PW; t /= PW;
        int ph = t % PH; t /= PH;
        int c = t % C;
        int n = t / C;
        int h0 = ph * 2;
        int w0 = pw * 2;
        float a = bn_weight[c] * invstd[c];
        float b = bn_bias[c] - mean[c] * a;
        long long base = ((long long)n * C + c) * HW + h0 * W + w0;
        float v0 = tanhf(fmaf(x[base], a, b));
        float v1 = tanhf(fmaf(x[base + 1], a, b));
        float v2 = tanhf(fmaf(x[base + W], a, b));
        float v3 = tanhf(fmaf(x[base + W + 1], a, b));
        float m = fmaxf(fmaxf(v0, v1), fmaxf(v2, v3));
        pooled[idx] = m;
    }
}

__global__ void group_stats_kernel(const float* __restrict__ pooled, float* __restrict__ gmean, float* __restrict__ grstd,
                                   int N, int C, int PH, int PW, int G, float eps) {
    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng - n * G;
    int cpg = C / G;
    int PHW = PH * PW;
    int M = cpg * PHW;
    int c0 = g * cpg;
    float sum = 0.0f;
    float sum2 = 0.0f;
    long long nbase = (long long)n * C * PHW;
    for (int i = threadIdx.x; i < M; i += blockDim.x) {
        int lc = i / PHW;
        int hw = i - lc * PHW;
        float v = pooled[nbase + (long long)(c0 + lc) * PHW + hw];
        sum += v;
        sum2 = fmaf(v, v, sum2);
    }
    sum = block_reduce_sum(sum);
    sum2 = block_reduce_sum(sum2);
    if (threadIdx.x == 0) {
        float m = sum / (float)M;
        float var = sum2 / (float)M - m * m;
        var = var < 0.0f ? 0.0f : var;
        gmean[ng] = m;
        grstd[ng] = rsqrtf(var + eps);
    }
}

__global__ void group_apply_kernel(const float* __restrict__ pooled,
                                   const float* __restrict__ gmean,
                                   const float* __restrict__ grstd,
                                   const float* __restrict__ gn_weight,
                                   const float* __restrict__ gn_bias,
                                   float* __restrict__ out,
                                   int total, int C, int PH, int PW, int G) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    int PHW = PH * PW;
    int cpg = C / G;
    for (; idx < total; idx += stride) {
        int t = idx;
        int hw = t % PHW; (void)hw;
        t /= PHW;
        int c = t % C;
        int n = t / C;
        int g = c / cpg;
        int ng = n * G + g;
        float v = (pooled[idx] - gmean[ng]) * grstd[ng];
        out[idx] = fmaf(v, gn_weight[c], gn_bias[c]);
    }
}

torch::Tensor fused_post(torch::Tensor conv,
                         torch::Tensor bn_weight,
                         torch::Tensor bn_bias,
                         torch::Tensor running_mean,
                         torch::Tensor running_var,
                         torch::Tensor gn_weight,
                         torch::Tensor gn_bias,
                         bool bn_training,
                         double bn_eps_d,
                         int64_t num_groups,
                         double gn_eps_d) {
    TORCH_CHECK(conv.is_cuda(), "conv must be CUDA");
    TORCH_CHECK(conv.scalar_type() == at::kFloat, "only float32 supported");
    conv = conv.contiguous();
    int N = (int)conv.size(0);
    int C = (int)conv.size(1);
    int H = (int)conv.size(2);
    int W = (int)conv.size(3);
    int PH = (H - 2) / 2 + 1;
    int PW = (W - 2) / 2 + 1;
    float bn_eps = (float)bn_eps_d;
    float gn_eps = (float)gn_eps_d;
    auto opts = conv.options();
    auto mean = torch::empty({C}, opts);
    auto invstd = torch::empty({C}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (bn_training) {
        bn_stats_kernel<<<C, 256, 0, stream>>>(conv.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), N, C, H, W, bn_eps);
    } else {
        // invstd from running variance; mean is running mean. Small one-block-per-channel kernel via bn_stats_kernel not suitable,
        // so use a simple lambda-like kernel implemented below by reusing group_apply? no: launch dedicated local kernel unavailable.
        // Fill using a tiny generic transform kernel defined inline as a separate launch below.
        int threads = 128;
        int blocks = (C + threads - 1) / threads;
        // use cuda kernel declared in this translation unit
        // placeholder replaced by actual kernel call
        
        // Direct copy/compute with at operations would add overhead; kernel below.
        extern void launch_bn_eval_params(const float*, const float*, float*, float*, int, float, cudaStream_t);
        launch_bn_eval_params(running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), C, bn_eps, stream);
    }
    auto pooled = torch::empty({N, C, PH, PW}, opts);
    int total = N * C * PH * PW;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 65535) blocks = 65535;
    pool_bn_tanh_kernel<<<blocks, threads, 0, stream>>>(conv.data_ptr<float>(), bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), pooled.data_ptr<float>(), total, C, H, W, PH, PW);
    int G = (int)num_groups;
    auto gmean = torch::empty({N, G}, opts);
    auto grstd = torch::empty({N, G}, opts);
    group_stats_kernel<<<N * G, 256, 0, stream>>>(pooled.data_ptr<float>(), gmean.data_ptr<float>(), grstd.data_ptr<float>(), N, C, PH, PW, G, gn_eps);
    auto out = torch::empty({N, C, PH, PW}, opts);
    group_apply_kernel<<<blocks, threads, 0, stream>>>(pooled.data_ptr<float>(), gmean.data_ptr<float>(), grstd.data_ptr<float>(), gn_weight.data_ptr<float>(), gn_bias.data_ptr<float>(), out.data_ptr<float>(), total, C, PH, PW, G);
    return out;
}

__global__ void bn_eval_params_kernel(const float* __restrict__ running_mean,
                                      const float* __restrict__ running_var,
                                      float* __restrict__ mean,
                                      float* __restrict__ invstd,
                                      int C, float eps) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c < C) {
        mean[c] = running_mean[c];
        invstd[c] = rsqrtf(running_var[c] + eps);
    }
}

void launch_bn_eval_params(const float* rm, const float* rv, float* mean, float* invstd, int C, float eps, cudaStream_t stream) {
    int threads = 128;
    int blocks = (C + threads - 1) / threads;
    bn_eval_params_kernel<<<blocks, threads, 0, stream>>>(rm, rv, mean, invstd, C, eps);
}
"""

_ext = load_inline(
    name="fused_deconv_bn_tanh_pool_gn_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_post"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups, num_groups):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm2d(out_channels)
        self.tanh = nn.Tanh()
        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.group_norm = nn.GroupNorm(num_groups=num_groups, num_channels=out_channels)
        self.num_groups = num_groups

    def forward(self, x):
        y = F.conv_transpose2d(x, self.conv_transpose.weight, self.conv_transpose.bias,
                               stride=self.conv_transpose.stride, padding=self.conv_transpose.padding,
                               output_padding=self.conv_transpose.output_padding, groups=self.conv_transpose.groups,
                               dilation=self.conv_transpose.dilation)
        return _ext.fused_post(y,
                               self.batch_norm.weight, self.batch_norm.bias,
                               self.batch_norm.running_mean, self.batch_norm.running_var,
                               self.group_norm.weight, self.group_norm.bias,
                               bool(self.batch_norm.training), float(self.batch_norm.eps),
                               int(self.group_norm.num_groups), float(self.group_norm.eps))