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

# Custom fused ReLU + GroupNorm CUDA kernels.
# Pass 1: compute per-group sum / sum-of-squares of relu(x) with block
#         reductions + double-precision atomics.
# Pass 2: recompute relu(x) on the fly and write the normalized output.
# This avoids materializing the intermediate ReLU tensor and removes a full
# read+write pass of the activation compared to eager PyTorch.
relu_gn_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <climits>
#include <cstdint>

#define GN_BLOCK 256
#define GN_WARPS (GN_BLOCK / 32)

__global__ void relu_gn_stats_vec4_kernel(const float4* __restrict__ x,
                                          double* __restrict__ acc,
                                          const int group_vec,
                                          const int blocks_per_group,
                                          const int chunk_vec) {
    const int bid = blockIdx.x;
    const int g = bid / blocks_per_group;
    const int t = bid - g * blocks_per_group;
    const int start = g * group_vec + t * chunk_vec;
    int end = start + chunk_vec;
    const int gend = (g + 1) * group_vec;
    if (end > gend) end = gend;

    float s = 0.f, q = 0.f;
    for (int i = start + threadIdx.x; i < end; i += GN_BLOCK) {
        const float4 v = x[i];
        const float a = fmaxf(v.x, 0.f);
        const float b = fmaxf(v.y, 0.f);
        const float c = fmaxf(v.z, 0.f);
        const float d = fmaxf(v.w, 0.f);
        s += (a + b) + (c + d);
        q += a * a + b * b + c * c + d * d;
    }

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, off);
        q += __shfl_down_sync(0xffffffffu, q, off);
    }
    __shared__ float sm_s[GN_WARPS];
    __shared__ float sm_q[GN_WARPS];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        sm_s[wid] = s;
        sm_q[wid] = q;
    }
    __syncthreads();
    if (wid == 0) {
        s = (lane < GN_WARPS) ? sm_s[lane] : 0.f;
        q = (lane < GN_WARPS) ? sm_q[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, off);
            q += __shfl_down_sync(0xffffffffu, q, off);
        }
        if (lane == 0) {
            atomicAdd(&acc[2 * g],     (double)s);
            atomicAdd(&acc[2 * g + 1], (double)q);
        }
    }
}

__global__ void relu_gn_stats_scalar_kernel(const float* __restrict__ x,
                                            double* __restrict__ acc,
                                            const long long group_elems,
                                            const int blocks_per_group,
                                            const long long chunk) {
    const long long bid = blockIdx.x;
    const long long g = bid / blocks_per_group;
    const long long t = bid - g * blocks_per_group;
    const long long start = g * group_elems + t * chunk;
    long long end = start + chunk;
    const long long gend = (g + 1) * group_elems;
    if (end > gend) end = gend;

    float s = 0.f, q = 0.f;
    for (long long i = start + threadIdx.x; i < end; i += GN_BLOCK) {
        const float a = fmaxf(x[i], 0.f);
        s += a;
        q += a * a;
    }

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, off);
        q += __shfl_down_sync(0xffffffffu, q, off);
    }
    __shared__ float sm_s[GN_WARPS];
    __shared__ float sm_q[GN_WARPS];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        sm_s[wid] = s;
        sm_q[wid] = q;
    }
    __syncthreads();
    if (wid == 0) {
        s = (lane < GN_WARPS) ? sm_s[lane] : 0.f;
        q = (lane < GN_WARPS) ? sm_q[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, off);
            q += __shfl_down_sync(0xffffffffu, q, off);
        }
        if (lane == 0) {
            atomicAdd(&acc[2 * g],     (double)s);
            atomicAdd(&acc[2 * g + 1], (double)q);
        }
    }
}

__global__ void relu_gn_norm_vec4_kernel(const float4* __restrict__ x,
                                         float4* __restrict__ y,
                                         const double* __restrict__ acc,
                                         const float* __restrict__ gamma,
                                         const float* __restrict__ beta,
                                         const int total_vec,
                                         const int dhw_vec,
                                         const int C,
                                         const int CPG,
                                         const double inv_count,
                                         const float eps) {
    const int i = blockIdx.x * GN_BLOCK + threadIdx.x;
    if (i >= total_vec) return;
    const int c_global = i / dhw_vec;          // n * C + c
    const int g = c_global / CPG;              // n * G + group
    const int c = c_global % C;                // channel in batch
    const float mean = (float)(acc[2 * g] * inv_count);
    float var = (float)(acc[2 * g + 1] * inv_count) - mean * mean;
    var = fmaxf(var, 0.f);
    const float rstd = rsqrtf(var + eps);
    const float gc = gamma ? gamma[c] : 1.f;
    const float bc = beta  ? beta[c]  : 0.f;
    const float4 v = x[i];
    float4 o;
    o.x = (fmaxf(v.x, 0.f) - mean) * rstd * gc + bc;
    o.y = (fmaxf(v.y, 0.f) - mean) * rstd * gc + bc;
    o.z = (fmaxf(v.z, 0.f) - mean) * rstd * gc + bc;
    o.w = (fmaxf(v.w, 0.f) - mean) * rstd * gc + bc;
    y[i] = o;
}

__global__ void relu_gn_norm_scalar_kernel(const float* __restrict__ x,
                                           float* __restrict__ y,
                                           const double* __restrict__ acc,
                                           const float* __restrict__ gamma,
                                           const float* __restrict__ beta,
                                           const long long total,
                                           const long long dhw,
                                           const int C,
                                           const int CPG,
                                           const double inv_count,
                                           const float eps) {
    const long long i = (long long)blockIdx.x * GN_BLOCK + threadIdx.x;
    if (i >= total) return;
    const long long c_global = i / dhw;
    const int g = (int)(c_global / CPG);
    const int c = (int)(c_global % C);
    const float mean = (float)(acc[2 * g] * inv_count);
    float var = (float)(acc[2 * g + 1] * inv_count) - mean * mean;
    var = fmaxf(var, 0.f);
    const float rstd = rsqrtf(var + eps);
    const float gc = gamma ? gamma[c] : 1.f;
    const float bc = beta  ? beta[c]  : 0.f;
    y[i] = (fmaxf(x[i], 0.f) - mean) * rstd * gc + bc;
}

torch::Tensor relu_groupnorm(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             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");
    auto xc = x.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    int64_t DHW = 1;
    for (int64_t i = 2; i < xc.dim(); ++i) DHW *= xc.size(i);
    TORCH_CHECK(groups > 0 && (C % groups) == 0, "invalid group count");
    const int64_t CPG = C / groups;
    const int64_t NG = N * groups;
    const int64_t numel = N * C * DHW;

    auto y = at::empty_like(xc);
    auto acc = at::zeros({NG * 2}, xc.options().dtype(at::kDouble));

    torch::Tensor wc, bc;
    const float* gamma = nullptr;
    const float* beta = nullptr;
    if (weight.defined() && weight.numel() > 0) {
        TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
        wc = weight.contiguous();
        gamma = wc.data_ptr<float>();
    }
    if (bias.defined() && bias.numel() > 0) {
        TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");
        bc = bias.contiguous();
        beta = bc.data_ptr<float>();
    }

    const double inv_count = 1.0 / (double)(CPG * DHW);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const bool aligned =
        ((reinterpret_cast<uintptr_t>(xc.data_ptr<float>()) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(y.data_ptr<float>())  & 15) == 0);

    if ((DHW % 4 == 0) && aligned && (numel / 4 <= (int64_t)INT_MAX)) {
        const int total_vec = (int)(numel >> 2);
        const int dhw_vec   = (int)(DHW >> 2);
        const int group_vec = (int)(CPG * dhw_vec);
        int blocks_per_group = (group_vec + 4095) / 4096;
        if (blocks_per_group < 1) blocks_per_group = 1;
        if (blocks_per_group > 512) blocks_per_group = 512;
        const int chunk_vec = (group_vec + blocks_per_group - 1) / blocks_per_group;
        const int64_t grid_stats = NG * (int64_t)blocks_per_group;
        relu_gn_stats_vec4_kernel<<<(unsigned)grid_stats, GN_BLOCK, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            acc.data_ptr<double>(), group_vec, blocks_per_group, chunk_vec);
        const int64_t grid_norm = ((int64_t)total_vec + GN_BLOCK - 1) / GN_BLOCK;
        relu_gn_norm_vec4_kernel<<<(unsigned)grid_norm, GN_BLOCK, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<float4*>(y.data_ptr<float>()),
            acc.data_ptr<double>(), gamma, beta,
            total_vec, dhw_vec, (int)C, (int)CPG, inv_count, (float)eps);
    } else {
        const long long group_elems = (long long)(CPG * DHW);
        int blocks_per_group = (int)((group_elems + 16383) / 16384);
        if (blocks_per_group < 1) blocks_per_group = 1;
        if (blocks_per_group > 512) blocks_per_group = 512;
        const long long chunk = (group_elems + blocks_per_group - 1) / blocks_per_group;
        const long long grid_stats = (long long)NG * blocks_per_group;
        relu_gn_stats_scalar_kernel<<<(unsigned)grid_stats, GN_BLOCK, 0, stream>>>(
            xc.data_ptr<float>(), acc.data_ptr<double>(),
            group_elems, blocks_per_group, chunk);
        const long long grid_norm = (numel + GN_BLOCK - 1) / GN_BLOCK;
        relu_gn_norm_scalar_kernel<<<(unsigned)grid_norm, GN_BLOCK, 0, stream>>>(
            xc.data_ptr<float>(), y.data_ptr<float>(),
            acc.data_ptr<double>(), gamma, beta,
            (long long)numel, (long long)DHW, (int)C, (int)CPG, inv_count, (float)eps);
    }
    return y;
}
"""

relu_gn_cpp_source = (
    "torch::Tensor relu_groupnorm(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, int64_t groups, double eps);"
)

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


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      1) ConvTranspose3d (stride=1, pad=0) is algebraically rewritten as an
         equivalent regular Conv3d with padding = K-1 and the kernel flipped
         spatially with in/out channels swapped. This lets cuDNN run its
         highly optimized forward-convolution kernels instead of the
         transposed-convolution path.
      2) ReLU + GroupNorm are replaced by a pair of fused custom CUDA kernels
         that never materialize the ReLU output (ReLU is recomputed on the
         fly), cutting global-memory traffic for the elementwise/normalization
         stages roughly in half.
    """
    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.fused = fused_relu_groupnorm
        # Let cuDNN autotune the best algorithm for the (fixed) conv shape.
        torch.backends.cudnn.benchmark = True

    def forward(self, x):
        ct = self.conv_transpose
        gn = self.group_norm

        # ConvTranspose3d -> equivalent Conv3d (valid for the default
        # stride=1, padding=0, dilation=1 configuration).
        if (ct.groups == 1
                and ct.stride == (1, 1, 1)
                and ct.padding == (0, 0, 0)
                and ct.dilation == (1, 1, 1)
                and ct.output_padding == (0, 0, 0)):
            w2 = ct.weight.permute(1, 0, 2, 3, 4).flip([2, 3, 4]).contiguous()
            pad = (ct.kernel_size[0] - 1, ct.kernel_size[1] - 1, ct.kernel_size[2] - 1)
            y = F.conv3d(x, w2, ct.bias, 1, pad)
        else:
            y = ct(x)

        # Fused ReLU + GroupNorm custom CUDA kernels.
        if y.is_cuda and y.dtype == torch.float32:
            w = gn.weight
            b = gn.bias
            if w is None:
                w = torch.empty(0, device=y.device, dtype=y.dtype)
            if b is None:
                b = torch.empty(0, device=y.device, dtype=y.dtype)
            return self.fused.relu_groupnorm(y, w, b, gn.num_groups, gn.eps)

        # Generic fallback for non-CUDA / non-float32 tensors.
        return F.group_norm(F.relu(y), gn.num_groups, gn.weight, gn.bias, gn.eps)
```