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

torch.backends.cudnn.benchmark = True

# Fused ReLU + GroupNorm CUDA kernel.
# The transposed 3D convolution is left to cuDNN (highly optimized),
# while the elementwise ReLU and the GroupNorm (two full memory passes +
# a reduction pass in eager PyTorch) are fused into a single kernel that
# performs one persistent block per (batch, group) pair:
#   phase 1: block-parallel reduction of sum / sum-of-squares of relu(x)
#   phase 2: normalize + per-channel affine, writing the output.
# Vectorized float4 accesses are used when the spatial extent allows it.
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

#define BLOCK_SIZE 512

__global__ void fused_relu_groupnorm_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ out,
    int G,
    long C_per_G,
    long spatial,
    long group_size,
    float eps)
{
    int g = blockIdx.x % G;
    // Groups are contiguous in memory: block b handles (n, g) with b = n*G + g
    const float* x_base   = x   + (long)blockIdx.x * group_size;
    float*       out_base = out + (long)blockIdx.x * group_size;

    float sum   = 0.f;
    float sumsq = 0.f;

    bool vec4 = (spatial & 3) == 0;

    // ---- Phase 1: statistics over relu(x) ----
    if (vec4) {
        long sp4 = spatial >> 2;
        for (long c = 0; c < C_per_G; ++c) {
            const float4* xc = reinterpret_cast<const float4*>(x_base + c * spatial);
            for (long i = threadIdx.x; i < sp4; i += blockDim.x) {
                float4 v = xc[i];
                v.x = fmaxf(v.x, 0.f);
                v.y = fmaxf(v.y, 0.f);
                v.z = fmaxf(v.z, 0.f);
                v.w = fmaxf(v.w, 0.f);
                sum   += v.x + v.y + v.z + v.w;
                sumsq += v.x*v.x + v.y*v.y + v.z*v.z + v.w*v.w;
            }
        }
    } else {
        for (long i = threadIdx.x; i < group_size; i += blockDim.x) {
            float v = fmaxf(x_base[i], 0.f);
            sum   += v;
            sumsq += v * v;
        }
    }

    // ---- Block reduction ----
    __shared__ float s_sum[BLOCK_SIZE];
    __shared__ float s_sumsq[BLOCK_SIZE];
    s_sum[threadIdx.x]   = sum;
    s_sumsq[threadIdx.x] = sumsq;
    __syncthreads();

    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            s_sum[threadIdx.x]   += s_sum[threadIdx.x + stride];
            s_sumsq[threadIdx.x] += s_sumsq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    float mean = s_sum[0] / (float)group_size;
    float var  = s_sumsq[0] / (float)group_size - mean * mean;
    float rstd = rsqrtf(var + eps);

    // ---- Phase 2: normalize + affine, write output ----
    if (vec4) {
        long sp4 = spatial >> 2;
        for (long c = 0; c < C_per_G; ++c) {
            const float4* xc = reinterpret_cast<const float4*>(x_base + c * spatial);
            float4*       oc = reinterpret_cast<float4*>(out_base + c * spatial);
            float gm = gamma[g * C_per_G + c];
            float bt = beta[g * C_per_G + c];
            float scale = rstd * gm;
            for (long i = threadIdx.x; i < sp4; i += blockDim.x) {
                float4 v = xc[i];
                v.x = (fmaxf(v.x, 0.f) - mean) * scale + bt;
                v.y = (fmaxf(v.y, 0.f) - mean) * scale + bt;
                v.z = (fmaxf(v.z, 0.f) - mean) * scale + bt;
                v.w = (fmaxf(v.w, 0.f) - mean) * scale + bt;
                oc[i] = v;
            }
        }
    } else {
        for (long c = 0; c < C_per_G; ++c) {
            const float* xc = x_base + c * spatial;
            float*       oc = out_base + c * spatial;
            float gm = gamma[g * C_per_G + c];
            float bt = beta[g * C_per_G + c];
            float scale = rstd * gm;
            for (long i = threadIdx.x; i < spatial; i += blockDim.x) {
                float v = fmaxf(xc[i], 0.f);
                oc[i] = (v - mean) * scale + bt;
            }
        }
    }
}

torch::Tensor fused_relu_groupnorm_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
                                        int64_t N, int64_t C, int64_t G, int64_t spatial, double eps)
{
    auto out = torch::empty_like(x);
    long C_per_G    = (long)(C / G);
    long group_size = C_per_G * (long)spatial;
    long num_blocks = (long)N * (long)G;

    dim3 grid((unsigned int)num_blocks);
    dim3 block(BLOCK_SIZE);
    auto stream = at::cuda::getCurrentCUDAStream();

    fused_relu_groupnorm_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)G,
        C_per_G,
        (long)spatial,
        group_size,
        (float)eps);

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_relu_groupnorm_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, "
    "int64_t N, int64_t C, int64_t G, int64_t spatial, double eps);"
)

fused_relu_groupnorm = load_inline(
    name="fused_relu_groupnorm",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_relu_groupnorm_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: transposed 3D convolution (cuDNN) followed by a single
    fused CUDA kernel performing ReLU + GroupNorm.
    """
    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.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)
        self.groups = groups

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x.contiguous()
        N, C, D, H, W = x.shape
        spatial = D * H * W
        return fused_relu_groupnorm.fused_relu_groupnorm_cuda(
            x,
            self.group_norm.weight,
            self.group_norm.bias,
            N, C, self.groups, spatial, self.group_norm.eps)
```