```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

torch.backends.cudnn.benchmark = True

# CUDA source: fused ReLU + GroupNorm kernel.
# One thread block per (batch, group). The block first computes the mean and
# inverse-std of relu(x) over its group with a vectorized block reduction,
# then applies the normalization in a second sweep (data comes back from L2).
fused_relu_groupnorm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_THREADS 512
#define NWARPS (BLOCK_THREADS / 32)

template <bool VEC>
__global__ void fused_relu_groupnorm_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ y,
    const int C,
    const int S,
    const int groups,
    const float eps)
{
    const int g = blockIdx.x % groups;
    const int cpg = C / groups;
    // blockIdx.x == n * groups + g, and groups * cpg == C, so the element
    // offset of this (n, g) pair is simply blockIdx.x * cpg * S.
    const int64_t base = (int64_t)blockIdx.x * (int64_t)cpg * (int64_t)S;
    const int64_t M = (int64_t)cpg * (int64_t)S;

    // ---------------- Phase 1: compute sum / sum-of-squares of relu(x) -----
    float s0 = 0.f, s1 = 0.f, q0 = 0.f, q1 = 0.f;

    if (VEC) {
        const float4* xv = reinterpret_cast<const float4*>(x + base);
        const int64_t n4 = M >> 2;
        int64_t i = threadIdx.x;
        for (; i + BLOCK_THREADS < n4; i += 2 * BLOCK_THREADS) {
            float4 a = xv[i];
            float4 b = xv[i + BLOCK_THREADS];
            a.x = fmaxf(a.x, 0.f); a.y = fmaxf(a.y, 0.f);
            a.z = fmaxf(a.z, 0.f); a.w = fmaxf(a.w, 0.f);
            b.x = fmaxf(b.x, 0.f); b.y = fmaxf(b.y, 0.f);
            b.z = fmaxf(b.z, 0.f); b.w = fmaxf(b.w, 0.f);
            s0 += (a.x + a.y) + (a.z + a.w);
            q0 += (a.x * a.x + a.y * a.y) + (a.z * a.z + a.w * a.w);
            s1 += (b.x + b.y) + (b.z + b.w);
            q1 += (b.x * b.x + b.y * b.y) + (b.z * b.z + b.w * b.w);
        }
        if (i < n4) {
            float4 a = xv[i];
            a.x = fmaxf(a.x, 0.f); a.y = fmaxf(a.y, 0.f);
            a.z = fmaxf(a.z, 0.f); a.w = fmaxf(a.w, 0.f);
            s0 += (a.x + a.y) + (a.z + a.w);
            q0 += (a.x * a.x + a.y * a.y) + (a.z * a.z + a.w * a.w);
        }
    } else {
        const float* xs = x + base;
        int64_t i = threadIdx.x;
        for (; i + BLOCK_THREADS < M; i += 2 * BLOCK_THREADS) {
            const float a = fmaxf(xs[i], 0.f);
            const float b = fmaxf(xs[i + BLOCK_THREADS], 0.f);
            s0 += a; q0 += a * a;
            s1 += b; q1 += b * b;
        }
        if (i < M) {
            const float a = fmaxf(xs[i], 0.f);
            s0 += a; q0 += a * a;
        }
    }

    float sum = s0 + s1;
    float ssq = q0 + q1;

    __shared__ float sh_sum[NWARPS];
    __shared__ float sh_ssq[NWARPS];
    __shared__ float sh_mean;
    __shared__ float sh_rstd;

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, off);
        ssq += __shfl_down_sync(0xffffffffu, ssq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        sh_sum[wid] = sum;
        sh_ssq[wid] = ssq;
    }
    __syncthreads();
    if (wid == 0) {
        sum = (lane < NWARPS) ? sh_sum[lane] : 0.f;
        ssq = (lane < NWARPS) ? sh_ssq[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, off);
            ssq += __shfl_down_sync(0xffffffffu, ssq, off);
        }
        if (lane == 0) {
            const float mean = sum / (float)M;
            float var = ssq / (float)M - mean * mean;
            var = fmaxf(var, 0.f);
            sh_mean = mean;
            sh_rstd = rsqrtf(var + eps);
        }
    }
    __syncthreads();
    const float mean = sh_mean;
    const float rstd = sh_rstd;

    // ---------------- Phase 2: normalize -----------------------------------
    const int cbase = g * cpg;
    for (int c = 0; c < cpg; ++c) {
        const float scale = rstd * gamma[cbase + c];
        const float shift = beta[cbase + c] - mean * scale;
        const int64_t off = base + (int64_t)c * (int64_t)S;
        if (VEC) {
            const int n4 = S >> 2;
            const float4* xc = reinterpret_cast<const float4*>(x + off);
            float4* yc = reinterpret_cast<float4*>(y + off);
            int j = threadIdx.x;
            for (; j + BLOCK_THREADS < n4; j += 2 * BLOCK_THREADS) {
                const float4 a = xc[j];
                const float4 b = xc[j + BLOCK_THREADS];
                float4 oa, ob;
                oa.x = fmaxf(a.x, 0.f) * scale + shift;
                oa.y = fmaxf(a.y, 0.f) * scale + shift;
                oa.z = fmaxf(a.z, 0.f) * scale + shift;
                oa.w = fmaxf(a.w, 0.f) * scale + shift;
                ob.x = fmaxf(b.x, 0.f) * scale + shift;
                ob.y = fmaxf(b.y, 0.f) * scale + shift;
                ob.z = fmaxf(b.z, 0.f) * scale + shift;
                ob.w = fmaxf(b.w, 0.f) * scale + shift;
                yc[j] = oa;
                yc[j + BLOCK_THREADS] = ob;
            }
            if (j < n4) {
                const float4 a = xc[j];
                float4 oa;
                oa.x = fmaxf(a.x, 0.f) * scale + shift;
                oa.y = fmaxf(a.y, 0.f) * scale + shift;
                oa.z = fmaxf(a.z, 0.f) * scale + shift;
                oa.w = fmaxf(a.w, 0.f) * scale + shift;
                yc[j] = oa;
            }
        } else {
            const float* xc = x + off;
            float* yc = y + off;
            int j = threadIdx.x;
            for (; j + BLOCK_THREADS < S; j += 2 * BLOCK_THREADS) {
                yc[j] = fmaxf(xc[j], 0.f) * scale + shift;
                yc[j + BLOCK_THREADS] = fmaxf(xc[j + BLOCK_THREADS], 0.f) * scale + shift;
            }
            if (j < S) {
                yc[j] = fmaxf(xc[j], 0.f) * scale + shift;
            }
        }
    }
}

torch::Tensor fused_relu_groupnorm(torch::Tensor x, torch::Tensor gamma,
                                   torch::Tensor beta, int64_t groups,
                                   double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() >= 2, "x must have at least 2 dims");
    TORCH_CHECK(gamma.scalar_type() == at::kFloat && beta.scalar_type() == at::kFloat,
                "gamma/beta must be float32");

    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();
    auto y = at::empty_like(xc);

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    int64_t S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    TORCH_CHECK(C % groups == 0, "channels must be divisible by groups");

    const dim3 grid((unsigned)(N * groups));
    const dim3 block(BLOCK_THREADS);
    auto stream = at::cuda::getCurrentCUDAStream();

    if ((S & 3) == 0) {
        fused_relu_groupnorm_kernel<true><<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), gc.data_ptr<float>(), bc.data_ptr<float>(),
            y.data_ptr<float>(), (int)C, (int)S, (int)groups, (float)eps);
    } else {
        fused_relu_groupnorm_kernel<false><<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), gc.data_ptr<float>(), bc.data_ptr<float>(),
            y.data_ptr<float>(), (int)C, (int)S, (int)groups, (float)eps);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

fused_relu_groupnorm_cpp_src = (
    "torch::Tensor fused_relu_groupnorm(torch::Tensor x, torch::Tensor gamma, "
    "torch::Tensor beta, int64_t groups, double eps);"
)

fused_relu_groupnorm_ext = load_inline(
    name="fused_relu_groupnorm_ext",
    cpp_sources=fused_relu_groupnorm_cpp_src,
    cuda_sources=fused_relu_groupnorm_cuda_src,
    functions=["fused_relu_groupnorm"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      1. ConvTranspose3d (stride=1, pad=0) is evaluated as an equivalent
         forward Conv3d with padding k-1 and a channel-permuted, spatially
         flipped kernel. This is mathematically identical but runs on the
         faster cuDNN forward-convolution path.
      2. ReLU + GroupNorm are fused into a single custom CUDA kernel
         (one read of the conv output for statistics, one fused
         read/normalize/write pass), removing the separate ReLU kernel and
         the extra global-memory round trips of the unfused ops.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels,
                                                 kernel_size, bias=bias)
        self.relu = nn.ReLU()
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)

        ct = self.conv_transpose
        self._fast_path_ok = (
            ct.groups == 1
            and ct.stride == (1, 1, 1)
            and ct.dilation == (1, 1, 1)
            and ct.output_padding == (0, 0, 0)
        )
        # Equivalent forward-conv padding for transposed conv with stride 1.
        self._conv_padding = tuple(k - 1 - p for k, p in zip(ct.kernel_size, ct.padding))

    def forward(self, x):
        if (self._fast_path_ok
                and x.is_cuda
                and x.dtype == torch.float32
                and self.group_norm.affine):
            # ConvTranspose3d(stride=1) == Conv3d with padding k-1-p and the
            # weight permuted (in<->out channels) and flipped spatially.
            w = self.conv_transpose.weight
            w_conv = w.permute(1, 0, 2, 3, 4).flip([2, 3, 4]).contiguous()
            y = F.conv3d(x, w_conv, self.conv_transpose.bias,
                         stride=1, padding=self._conv_padding)
            return fused_relu_groupnorm_ext.fused_relu_groupnorm(
                y, self.group_norm.weight, self.group_norm.bias,
                self.group_norm.num_groups, self.group_norm.eps)

        # Reference fallback (CPU tensors, non-fp32, or unsupported configs).
        x = self.conv_transpose(x)
        x = self.relu(x)
        x = self.group_norm(x)
        return x
```