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

# ---------------------------------------------------------------------------
# Fused ReLU + GroupNorm CUDA kernels.
#   Kernel 1: per (batch, group) reduction of relu(x) -> sum / sum of squares
#   Kernel 2: tiny finalize that turns sums into mean / rstd
#   Kernel 3: elementwise y = (relu(x) - mean) * rstd * gamma + beta
# This performs the ReLU and the full GroupNorm in only 2 reads of the
# activations and 1 write (instead of 3 reads + 2 writes for the unfused
# PyTorch implementation).
# ---------------------------------------------------------------------------
_fused_gn_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <stdint.h>

#define FULL_MASK 0xffffffffu

__global__ void relu_gn_stats_kernel(
    const float* __restrict__ x,
    double* __restrict__ sums,
    double* __restrict__ sumsqs,
    const int64_t group_size,
    const int ng_total)
{
    for (int ng = blockIdx.y; ng < ng_total; ng += gridDim.y) {
        const int64_t base  = (int64_t)ng * group_size;
        const int64_t chunk = (group_size + (int64_t)gridDim.x - 1) / (int64_t)gridDim.x;
        const int64_t start = base + (int64_t)blockIdx.x * chunk;
        int64_t end = start + chunk;
        const int64_t limit = base + group_size;
        if (end > limit) end = limit;

        float s = 0.0f;
        float q = 0.0f;
        for (int64_t i = start + (int64_t)threadIdx.x; i < end; i += (int64_t)blockDim.x) {
            const float v = fmaxf(x[i], 0.0f);
            s += v;
            q += v * v;
        }

        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(FULL_MASK, s, off);
            q += __shfl_down_sync(FULL_MASK, q, off);
        }

        __shared__ float sh_s[32];
        __shared__ float sh_q[32];
        const int lane = threadIdx.x & 31;
        const int wid  = threadIdx.x >> 5;
        if (lane == 0) {
            sh_s[wid] = s;
            sh_q[wid] = q;
        }
        __syncthreads();

        if (wid == 0) {
            const int nw = (int)(blockDim.x + 31) >> 5;
            float a = (lane < nw) ? sh_s[lane] : 0.0f;
            float b = (lane < nw) ? sh_q[lane] : 0.0f;
            #pragma unroll
            for (int off = 16; off > 0; off >>= 1) {
                a += __shfl_down_sync(FULL_MASK, a, off);
                b += __shfl_down_sync(FULL_MASK, b, off);
            }
            if (lane == 0) {
                atomicAdd(&sums[ng], (double)a);
                atomicAdd(&sumsqs[ng], (double)b);
            }
        }
        __syncthreads();
    }
}

__global__ void gn_finalize_kernel(
    const double* __restrict__ sums,
    const double* __restrict__ sumsqs,
    float* __restrict__ means,
    float* __restrict__ rstds,
    const int64_t group_size,
    const double eps,
    const int ng_total)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < ng_total) {
        const double inv_m = 1.0 / (double)group_size;
        const double mean  = sums[i] * inv_m;
        double var = sumsqs[i] * inv_m - mean * mean;
        if (var < 0.0) var = 0.0;
        means[i] = (float)mean;
        rstds[i] = (float)rsqrt(var + eps);
    }
}

__global__ void relu_gn_apply_vec4_kernel(
    const float4* __restrict__ x4,
    float4* __restrict__ y4,
    const float* __restrict__ means,
    const float* __restrict__ rstds,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const int C,
    const int G,
    const int Cg,
    const int NC,
    const int64_t S4)
{
    for (int nc = blockIdx.y; nc < NC; nc += gridDim.y) {
        const int c  = nc % C;
        const int n  = nc / C;
        const int ng = n * G + c / Cg;
        const float mean = means[ng];
        const float rstd = rstds[ng];
        const float g = gamma[c];
        const float b = beta[c];
        const int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < S4) {
            const int64_t off = (int64_t)nc * S4 + idx;
            const float4 v = x4[off];
            float4 o;
            o.x = (fmaxf(v.x, 0.0f) - mean) * rstd * g + b;
            o.y = (fmaxf(v.y, 0.0f) - mean) * rstd * g + b;
            o.z = (fmaxf(v.z, 0.0f) - mean) * rstd * g + b;
            o.w = (fmaxf(v.w, 0.0f) - mean) * rstd * g + b;
            y4[off] = o;
        }
    }
}

__global__ void relu_gn_apply_scalar_kernel(
    const float* __restrict__ x,
    float* __restrict__ y,
    const float* __restrict__ means,
    const float* __restrict__ rstds,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const int C,
    const int G,
    const int Cg,
    const int NC,
    const int64_t S)
{
    for (int nc = blockIdx.y; nc < NC; nc += gridDim.y) {
        const int c  = nc % C;
        const int n  = nc / C;
        const int ng = n * G + c / Cg;
        const float mean = means[ng];
        const float rstd = rstds[ng];
        const float g = gamma[c];
        const float b = beta[c];
        const int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < S) {
            const int64_t off = (int64_t)nc * S + idx;
            const float v = fmaxf(x[off], 0.0f);
            y[off] = (v - mean) * rstd * g + b;
        }
    }
}

torch::Tensor fused_relu_groupnorm_cuda(
    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(weight.defined() && weight.scalar_type() == at::kFloat, "weight must be float32");
    TORCH_CHECK(bias.defined() && bias.scalar_type() == at::kFloat, "bias must be float32");

    auto xc = x.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t S = (N * C > 0) ? (xc.numel() / (N * C)) : 0;
    const int64_t G = groups;
    const int64_t Cg = C / G;
    const int64_t group_size = Cg * S;
    const int64_t NG = N * G;
    const int64_t NC = N * C;

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(xc.device());
    auto opts_f = torch::TensorOptions().dtype(torch::kFloat32).device(xc.device());

    auto sums   = torch::zeros({NG}, opts_d);
    auto sumsqs = torch::zeros({NG}, opts_d);
    auto means  = torch::empty({NG}, opts_f);
    auto rstds  = torch::empty({NG}, opts_f);
    auto y      = torch::empty_like(xc);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;

    // 1) per-(batch, group) sums of relu(x)
    int64_t splits = group_size / (int64_t)(threads * 64);
    if (splits < 1) splits = 1;
    if (splits > 32) splits = 32;
    const unsigned gy = (unsigned)(NG < 65535 ? NG : 65535);
    dim3 stat_grid((unsigned)splits, gy, 1);
    relu_gn_stats_kernel<<<stat_grid, threads, 0, stream>>>(
        xc.data_ptr<float>(),
        sums.data_ptr<double>(),
        sumsqs.data_ptr<double>(),
        group_size,
        (int)NG);

    // 2) mean / rstd
    const int fin_threads = 128;
    const unsigned fin_blocks = (unsigned)((NG + fin_threads - 1) / fin_threads);
    gn_finalize_kernel<<<fin_blocks, fin_threads, 0, stream>>>(
        sums.data_ptr<double>(),
        sumsqs.data_ptr<double>(),
        means.data_ptr<float>(),
        rstds.data_ptr<float>(),
        group_size,
        eps,
        (int)NG);

    // 3) normalize
    const unsigned gy2 = (unsigned)(NC < 65535 ? NC : 65535);
    const bool aligned = ((((uintptr_t)xc.data_ptr<float>()) & 0xF) == 0) &&
                         ((((uintptr_t)y.data_ptr<float>())  & 0xF) == 0);
    if ((S & 3) == 0 && aligned) {
        const int64_t S4 = S >> 2;
        const unsigned bx = (unsigned)((S4 + threads - 1) / threads);
        dim3 apply_grid(bx, gy2, 1);
        relu_gn_apply_vec4_kernel<<<apply_grid, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<float4*>(y.data_ptr<float>()),
            means.data_ptr<float>(),
            rstds.data_ptr<float>(),
            weight.data_ptr<float>(),
            bias.data_ptr<float>(),
            (int)C, (int)G, (int)Cg, (int)NC, S4);
    } else {
        const unsigned bx = (unsigned)((S + threads - 1) / threads);
        dim3 apply_grid(bx, gy2, 1);
        relu_gn_apply_scalar_kernel<<<apply_grid, threads, 0, stream>>>(
            xc.data_ptr<float>(),
            y.data_ptr<float>(),
            means.data_ptr<float>(),
            rstds.data_ptr<float>(),
            weight.data_ptr<float>(),
            bias.data_ptr<float>(),
            (int)C, (int)G, (int)Cg, (int)NC, S);
    }

    return y;
}
"""

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

_fused_gn = load_inline(
    name="fused_relu_groupnorm_ext",
    cpp_sources=_fused_gn_cpp_src,
    cuda_sources=_fused_gn_cuda_src,
    functions=["fused_relu_groupnorm_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Two optimizations are applied:
      1. The ConvTranspose3d (stride=1, padding=0) is mathematically rewritten
         as an equivalent regular Conv3d with padding = kernel_size - 1 and a
         spatially flipped / channel-transposed kernel.  This lets cuDNN use
         its highly optimized forward-convolution kernels.
      2. ReLU + GroupNorm are fused into custom CUDA kernels that perform the
         statistics reduction and the normalization in a single pass pair,
         eliminating an entire read+write pass over the activations.
    """

    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
        if isinstance(kernel_size, int):
            k = (kernel_size, kernel_size, kernel_size)
        else:
            k = tuple(kernel_size)
        # Equivalent "full" convolution padding for stride-1 transposed conv.
        self._pad = (k[0] - 1, k[1] - 1, k[2] - 1)

    def forward(self, x):
        # ConvTranspose3d(stride=1, pad=0) weight: (in, out, kd, kh, kw)
        # Equivalent Conv3d weight: (out, in, kd, kh, kw) flipped spatially.
        w = self.conv_transpose.weight
        w_conv = w.permute(1, 0, 2, 3, 4).flip([2, 3, 4]).contiguous()
        out = F.conv3d(x, w_conv, self.conv_transpose.bias, 1, self._pad)
        out = _fused_gn.fused_relu_groupnorm_cuda(
            out,
            self.group_norm.weight,
            self.group_norm.bias,
            self.groups,
            float(self.group_norm.eps),
        )
        return out
```