```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

# Custom CUDA kernel that fuses ReLU + GroupNorm into a single kernel.
# The transposed 3D convolution is left to cuDNN (already highly tuned);
# the two following elementwise/reduction ops are fused to eliminate one
# full read + write pass over the activation tensor.
_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

#define FULL_MASK 0xffffffffu

template <int THREADS>
__device__ __forceinline__ void blockReduceSumSq(float& s, float& q,
                                                 float* sh_s, float* sh_q) {
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
#pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        s += __shfl_down_sync(FULL_MASK, s, o);
        q += __shfl_down_sync(FULL_MASK, q, o);
    }
    if (lane == 0) { sh_s[wid] = s; sh_q[wid] = q; }
    __syncthreads();
    constexpr int NW = THREADS / 32;
    if (wid == 0) {
        s = (lane < NW) ? sh_s[lane] : 0.0f;
        q = (lane < NW) ? sh_q[lane] : 0.0f;
#pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            s += __shfl_down_sync(FULL_MASK, s, o);
            q += __shfl_down_sync(FULL_MASK, q, o);
        }
        if (lane == 0) { sh_s[0] = s; sh_q[0] = q; }
    }
    __syncthreads();
}

// One block handles one (batch, group) slice.
template <int THREADS>
__global__ void reluGroupNormKernelScalar(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ y,
    const int Cg, const int DHW, const int groups, const float eps)
{
    const int64_t ng = blockIdx.x;
    const int g = (int)(ng % groups);
    const int inner = Cg * DHW;
    const float* xp = x + ng * (int64_t)inner;
    float* yp = y + ng * (int64_t)inner;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;
    const int tid = threadIdx.x;
    const int stride = THREADS;
    int i = tid;
    for (; i + 3 * stride < inner; i += 4 * stride) {
        float v0 = fmaxf(xp[i], 0.f);
        float v1 = fmaxf(xp[i + stride], 0.f);
        float v2 = fmaxf(xp[i + 2 * stride], 0.f);
        float v3 = fmaxf(xp[i + 3 * stride], 0.f);
        s0 += v0; s1 += v1; s2 += v2; s3 += v3;
        q0 = fmaf(v0, v0, q0); q1 = fmaf(v1, v1, q1);
        q2 = fmaf(v2, v2, q2); q3 = fmaf(v3, v3, q3);
    }
    for (; i < inner; i += stride) {
        float v = fmaxf(xp[i], 0.f);
        s0 += v; q0 = fmaf(v, v, q0);
    }
    float s = (s0 + s1) + (s2 + s3);
    float q = (q0 + q1) + (q2 + q3);
    __shared__ float sh_s[THREADS / 32];
    __shared__ float sh_q[THREADS / 32];
    blockReduceSumSq<THREADS>(s, q, sh_s, sh_q);

    const float mean = sh_s[0] / (float)inner;
    const float var  = sh_q[0] / (float)inner - mean * mean;
    const float rstd = rsqrtf(var + eps);

    const float* gp = gamma + g * Cg;
    const float* bp = beta + g * Cg;
    for (int c = 0; c < Cg; ++c) {
        const float scale = rstd * gp[c];
        const float shift = bp[c] - mean * scale;
        const float* xc = xp + (int64_t)c * DHW;
        float* yc = yp + (int64_t)c * DHW;
        int j = tid;
        for (; j + 3 * stride < DHW; j += 4 * stride) {
            yc[j]            = fmaf(fmaxf(xc[j], 0.f), scale, shift);
            yc[j + stride]   = fmaf(fmaxf(xc[j + stride], 0.f), scale, shift);
            yc[j + 2*stride] = fmaf(fmaxf(xc[j + 2*stride], 0.f), scale, shift);
            yc[j + 3*stride] = fmaf(fmaxf(xc[j + 3*stride], 0.f), scale, shift);
        }
        for (; j < DHW; j += stride) {
            yc[j] = fmaf(fmaxf(xc[j], 0.f), scale, shift);
        }
    }
}

// Vectorized (float4) version used when D*H*W is divisible by 4,
// which guarantees 16-byte alignment of every (batch, group, channel) slice.
template <int THREADS>
__global__ void reluGroupNormKernelVec4(
    const float4* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float4* __restrict__ y,
    const int Cg, const int DHW4, const int groups, const float eps)
{
    const int64_t ng = blockIdx.x;
    const int g = (int)(ng % groups);
    const int inner4 = Cg * DHW4;
    const float4* xp = x + ng * (int64_t)inner4;
    float4* yp = y + ng * (int64_t)inner4;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;
    const int tid = threadIdx.x;
    const int stride = THREADS;
    int i = tid;
    for (; i + 3 * stride < inner4; i += 4 * stride) {
        float4 a = xp[i];
        float4 b = xp[i + stride];
        float4 c = xp[i + 2 * stride];
        float4 d = xp[i + 3 * stride];
        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);
        c.x = fmaxf(c.x, 0.f); c.y = fmaxf(c.y, 0.f); c.z = fmaxf(c.z, 0.f); c.w = fmaxf(c.w, 0.f);
        d.x = fmaxf(d.x, 0.f); d.y = fmaxf(d.y, 0.f); d.z = fmaxf(d.z, 0.f); d.w = fmaxf(d.w, 0.f);
        s0 += (a.x + a.y) + (a.z + a.w);
        s1 += (b.x + b.y) + (b.z + b.w);
        s2 += (c.x + c.y) + (c.z + c.w);
        s3 += (d.x + d.y) + (d.z + d.w);
        q0 += (a.x * a.x + a.y * a.y) + (a.z * a.z + a.w * a.w);
        q1 += (b.x * b.x + b.y * b.y) + (b.z * b.z + b.w * b.w);
        q2 += (c.x * c.x + c.y * c.y) + (c.z * c.z + c.w * c.w);
        q3 += (d.x * d.x + d.y * d.y) + (d.z * d.z + d.w * d.w);
    }
    for (; i < inner4; i += stride) {
        float4 a = xp[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);
    }
    float s = (s0 + s1) + (s2 + s3);
    float q = (q0 + q1) + (q2 + q3);
    __shared__ float sh_s[THREADS / 32];
    __shared__ float sh_q[THREADS / 32];
    blockReduceSumSq<THREADS>(s, q, sh_s, sh_q);

    const float count = (float)inner4 * 4.0f;
    const float mean = sh_s[0] / count;
    const float var  = sh_q[0] / count - mean * mean;
    const float rstd = rsqrtf(var + eps);

    const float* gp = gamma + g * Cg;
    const float* bp = beta + g * Cg;
    for (int ch = 0; ch < Cg; ++ch) {
        const float scale = rstd * gp[ch];
        const float shift = bp[ch] - mean * scale;
        const float4* xc = xp + (int64_t)ch * DHW4;
        float4* yc = yp + (int64_t)ch * DHW4;
        int j = tid;
        for (; j + 3 * stride < DHW4; j += 4 * stride) {
            float4 a = xc[j];
            float4 b = xc[j + stride];
            float4 c = xc[j + 2 * stride];
            float4 d = xc[j + 3 * stride];
            a.x = fmaf(fmaxf(a.x, 0.f), scale, shift);
            a.y = fmaf(fmaxf(a.y, 0.f), scale, shift);
            a.z = fmaf(fmaxf(a.z, 0.f), scale, shift);
            a.w = fmaf(fmaxf(a.w, 0.f), scale, shift);
            b.x = fmaf(fmaxf(b.x, 0.f), scale, shift);
            b.y = fmaf(fmaxf(b.y, 0.f), scale, shift);
            b.z = fmaf(fmaxf(b.z, 0.f), scale, shift);
            b.w = fmaf(fmaxf(b.w, 0.f), scale, shift);
            c.x = fmaf(fmaxf(c.x, 0.f), scale, shift);
            c.y = fmaf(fmaxf(c.y, 0.f), scale, shift);
            c.z = fmaf(fmaxf(c.z, 0.f), scale, shift);
            c.w = fmaf(fmaxf(c.w, 0.f), scale, shift);
            d.x = fmaf(fmaxf(d.x, 0.f), scale, shift);
            d.y = fmaf(fmaxf(d.y, 0.f), scale, shift);
            d.z = fmaf(fmaxf(d.z, 0.f), scale, shift);
            d.w = fmaf(fmaxf(d.w, 0.f), scale, shift);
            yc[j]            = a;
            yc[j + stride]   = b;
            yc[j + 2*stride] = c;
            yc[j + 3*stride] = d;
        }
        for (; j < DHW4; j += stride) {
            float4 a = xc[j];
            a.x = fmaf(fmaxf(a.x, 0.f), scale, shift);
            a.y = fmaf(fmaxf(a.y, 0.f), scale, shift);
            a.z = fmaf(fmaxf(a.z, 0.f), scale, shift);
            a.w = fmaf(fmaxf(a.w, 0.f), scale, shift);
            yc[j] = a;
        }
    }
}

torch::Tensor fused_relu_groupnorm(torch::Tensor x, torch::Tensor weight,
                                   torch::Tensor bias, int64_t num_groups,
                                   double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 5, "x must have shape (N, C, D, H, W)");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(num_groups > 0, "num_groups must be positive");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    const int64_t N   = xc.size(0);
    const int64_t C   = xc.size(1);
    TORCH_CHECK(C % num_groups == 0, "C must be divisible by num_groups");
    const int64_t Cg  = C / num_groups;
    const int64_t DHW = xc.size(2) * xc.size(3) * xc.size(4);

    auto y = at::empty_like(xc);
    if (xc.numel() == 0) return y;

    constexpr int THREADS = 1024;
    const int blocks = (int)(N * num_groups);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((DHW & 3) == 0 && DHW * Cg <= INT32_MAX) {
        reluGroupNormKernelVec4<THREADS><<<blocks, THREADS, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            wc.data_ptr<float>(),
            bc.data_ptr<float>(),
            reinterpret_cast<float4*>(y.data_ptr<float>()),
            (int)Cg, (int)(DHW >> 2), (int)num_groups, (float)eps);
    } else {
        TORCH_CHECK(DHW * Cg <= INT32_MAX, "group size too large");
        reluGroupNormKernelScalar<THREADS><<<blocks, THREADS, 0, stream>>>(
            xc.data_ptr<float>(),
            wc.data_ptr<float>(),
            bc.data_ptr<float>(),
            y.data_ptr<float>(),
            (int)Cg, (int)DHW, (int)num_groups, (float)eps);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_relu_groupnorm(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, int64_t num_groups, double eps);"
)

_fused_ops = load_inline(
    name="fused_relu_groupnorm_ext",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_relu_groupnorm"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model: ConvTranspose3d (cuDNN) followed by a custom fused
    ReLU + GroupNorm CUDA kernel (single kernel, one pass over the stats
    and one pass for normalization, ReLU applied on-the-fly in both).
    """
    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)

    def forward(self, x):
        x = self.conv_transpose(x)
        if x.is_cuda and x.dtype == torch.float32:
            return _fused_ops.fused_relu_groupnorm(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.group_norm.num_groups,
                self.group_norm.eps,
            )
        return self.group_norm(self.relu(x))
```