```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.
# Strategy:
#   Kernel 1: computes per-(batch, group) sum and sum-of-squares of relu(x)
#             using float4 vectorized loads, double-precision accumulators,
#             warp-shuffle block reduction and atomics into a small workspace.
#   Kernel 2: reads the stats, then applies relu + group normalization
#             (with affine) in a single fused pass.
_fused_gn_cuda = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <algorithm>

#define GN_THREADS 256

__device__ __forceinline__ void gnBlockReduceD(double &sum, double &sumsq) {
    const unsigned int mask = 0xffffffffu;
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum   += __shfl_down_sync(mask, sum, offset);
        sumsq += __shfl_down_sync(mask, sumsq, offset);
    }
    __shared__ double sSum[GN_THREADS / 32];
    __shared__ double sSq[GN_THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) { sSum[wid] = sum; sSq[wid] = sumsq; }
    __syncthreads();
    if (wid == 0) {
        sum   = (lane < GN_THREADS / 32) ? sSum[lane] : 0.0;
        sumsq = (lane < GN_THREADS / 32) ? sSq[lane]  : 0.0;
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum   += __shfl_down_sync(mask, sum, offset);
            sumsq += __shfl_down_sync(mask, sumsq, offset);
        }
        if (lane == 0) { sSum[0] = sum; sSq[0] = sumsq; }
    }
    __syncthreads();
    sum   = sSum[0];
    sumsq = sSq[0];
}

// ------------------------- stats kernels -------------------------
__global__ __launch_bounds__(GN_THREADS)
void reluGnStatsVec4(const float* __restrict__ x, double* __restrict__ ws,
                     long long n4, long long chunk4, int chunksPerGroup) {
    const int gid   = blockIdx.x / chunksPerGroup;
    const int chunk = blockIdx.x - gid * chunksPerGroup;
    long long start = (long long)chunk * chunk4;
    long long end   = start + chunk4;
    if (end > n4) end = n4;
    const float4* xp = reinterpret_cast<const float4*>(x) + (long long)gid * n4;

    double sum = 0.0, sumsq = 0.0;
    for (long long i = start + threadIdx.x; i < end; i += GN_THREADS) {
        const float4 v = xp[i];
        const double a = (double)fmaxf(v.x, 0.0f);
        const double b = (double)fmaxf(v.y, 0.0f);
        const double c = (double)fmaxf(v.z, 0.0f);
        const double d = (double)fmaxf(v.w, 0.0f);
        sum   += (a + b) + (c + d);
        sumsq += (a * a + b * b) + (c * c + d * d);
    }
    gnBlockReduceD(sum, sumsq);
    if (threadIdx.x == 0) {
        atomicAdd(&ws[gid * 2],     sum);
        atomicAdd(&ws[gid * 2 + 1], sumsq);
    }
}

__global__ __launch_bounds__(GN_THREADS)
void reluGnStatsScalar(const float* __restrict__ x, double* __restrict__ ws,
                       long long groupSize, long long chunk, int chunksPerGroup) {
    const int gid = blockIdx.x / chunksPerGroup;
    const int ch  = blockIdx.x - gid * chunksPerGroup;
    long long start = (long long)ch * chunk;
    long long end   = start + chunk;
    if (end > groupSize) end = groupSize;
    const float* xp = x + (long long)gid * groupSize;

    double sum = 0.0, sumsq = 0.0;
    for (long long i = start + threadIdx.x; i < end; i += GN_THREADS) {
        const double v = (double)fmaxf(xp[i], 0.0f);
        sum   += v;
        sumsq += v * v;
    }
    gnBlockReduceD(sum, sumsq);
    if (threadIdx.x == 0) {
        atomicAdd(&ws[gid * 2],     sum);
        atomicAdd(&ws[gid * 2 + 1], sumsq);
    }
}

// ------------------------- apply kernels -------------------------
__global__ __launch_bounds__(GN_THREADS)
void reluGnApplyVec4(const float* __restrict__ x, const float* __restrict__ gamma,
                     const float* __restrict__ beta, const double* __restrict__ ws,
                     float* __restrict__ y, long long n4, long long chunk4,
                     int chunksPerGroup, unsigned int S, int G, int Cpg, float eps) {
    const int gid   = blockIdx.x / chunksPerGroup;
    const int chunk = blockIdx.x - gid * chunksPerGroup;
    long long start = (long long)chunk * chunk4;
    long long end   = start + chunk4;
    if (end > n4) end = n4;

    const double invN = 1.0 / (double)(n4 * 4);
    const double mean = ws[gid * 2] * invN;
    double var = ws[gid * 2 + 1] * invN - mean * mean;
    if (var < 0.0) var = 0.0;
    const float rstd  = (float)rsqrt(var + (double)eps);
    const float meanf = (float)mean;

    const int g  = gid % G;
    const int c0 = g * Cpg;

    const float4* xp = reinterpret_cast<const float4*>(x) + (long long)gid * n4;
    float4* yp = reinterpret_cast<float4*>(y) + (long long)gid * n4;

    for (long long i = start + threadIdx.x; i < end; i += GN_THREADS) {
        const float4 v = xp[i];
        // S % 4 == 0 in this path, so a float4 never crosses a channel boundary.
        const unsigned int baseIdx = (unsigned int)(i * 4);
        const int c = c0 + (int)(baseIdx / S);
        const float scale = rstd * gamma[c];
        const float shift = beta[c] - meanf * scale;
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), scale, shift);
        o.y = fmaf(fmaxf(v.y, 0.0f), scale, shift);
        o.z = fmaf(fmaxf(v.z, 0.0f), scale, shift);
        o.w = fmaf(fmaxf(v.w, 0.0f), scale, shift);
        yp[i] = o;
    }
}

__global__ __launch_bounds__(GN_THREADS)
void reluGnApplyScalar(const float* __restrict__ x, const float* __restrict__ gamma,
                       const float* __restrict__ beta, const double* __restrict__ ws,
                       float* __restrict__ y, long long groupSize, long long chunk,
                       int chunksPerGroup, long long S, int G, int Cpg, float eps) {
    const int gid = blockIdx.x / chunksPerGroup;
    const int ch  = blockIdx.x - gid * chunksPerGroup;
    long long start = (long long)ch * chunk;
    long long end   = start + chunk;
    if (end > groupSize) end = groupSize;

    const double invN = 1.0 / (double)groupSize;
    const double mean = ws[gid * 2] * invN;
    double var = ws[gid * 2 + 1] * invN - mean * mean;
    if (var < 0.0) var = 0.0;
    const float rstd  = (float)rsqrt(var + (double)eps);
    const float meanf = (float)mean;

    const int g  = gid % G;
    const int c0 = g * Cpg;

    const float* xp = x + (long long)gid * groupSize;
    float* yp = y + (long long)gid * groupSize;

    for (long long i = start + threadIdx.x; i < end; i += GN_THREADS) {
        const int c = c0 + (int)(i / S);
        const float scale = rstd * gamma[c];
        const float shift = beta[c] - meanf * scale;
        yp[i] = fmaf(fmaxf(xp[i], 0.0f), scale, shift);
    }
}

// ------------------------- host wrapper -------------------------
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");
    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();
    auto y  = torch::empty_like(xc);

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t S = xc.numel() / std::max<int64_t>(1, N * C);
    const int64_t G = groups;
    const int64_t Cpg = C / G;
    const int64_t groupSize = Cpg * S;
    if (groupSize == 0 || N == 0) return y;

    auto stream = at::cuda::getCurrentCUDAStream();
    auto ws = torch::zeros({N * G * 2}, xc.options().dtype(at::kDouble));

    const int numSM = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const long long groupsTotal = (long long)N * G;
    const long long target = (long long)numSM * 8;
    long long chunks = target / groupsTotal;
    if (chunks < 1) chunks = 1;
    if (chunks > 512) chunks = 512;

    const float* xp  = xc.data_ptr<float>();
    float* yp        = y.data_ptr<float>();
    const float* gp  = gc.data_ptr<float>();
    const float* bp  = bc.data_ptr<float>();
    double* wsp      = ws.data_ptr<double>();

    const bool aligned =
        ((reinterpret_cast<uintptr_t>(xp) | reinterpret_cast<uintptr_t>(yp)) & 0xF) == 0;
    const bool vec = aligned && (S % 4 == 0) && (groupSize <= 2147483647LL);

    if (vec) {
        long long n4 = groupSize / 4;
        long long chunk4 = (n4 + chunks - 1) / chunks;
        const int cpgChunks = (int)((n4 + chunk4 - 1) / chunk4);
        const dim3 grid((unsigned int)(groupsTotal * cpgChunks));
        reluGnStatsVec4<<<grid, GN_THREADS, 0, stream>>>(xp, wsp, n4, chunk4, cpgChunks);
        reluGnApplyVec4<<<grid, GN_THREADS, 0, stream>>>(xp, gp, bp, wsp, yp, n4, chunk4,
                                                         cpgChunks, (unsigned int)S,
                                                         (int)G, (int)Cpg, (float)eps);
    } else {
        long long chunk = (groupSize + chunks - 1) / chunks;
        const int cpgChunks = (int)((groupSize + chunk - 1) / chunk);
        const dim3 grid((unsigned int)(groupsTotal * cpgChunks));
        reluGnStatsScalar<<<grid, GN_THREADS, 0, stream>>>(xp, wsp, groupSize, chunk, cpgChunks);
        reluGnApplyScalar<<<grid, GN_THREADS, 0, stream>>>(xp, gp, bp, wsp, yp, groupSize, chunk,
                                                           cpgChunks, S, (int)G, (int)Cpg, (float)eps);
    }
    return y;
}
"""

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

_fused_gn = load_inline(
    name="fused_relu_groupnorm_v1",
    cpp_sources=_fused_gn_cpp,
    cuda_sources=_fused_gn_cuda,
    functions=["fused_relu_groupnorm"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model:
      - ConvTranspose3d (stride 1, pad 0, dilation 1) is rewritten as an
        equivalent cuDNN forward Conv3d with spatially-flipped / channel-swapped
        kernels and padding = k - 1, which uses the faster forward-conv path.
      - ReLU + GroupNorm (stats + normalize + affine) are fused into two custom
        CUDA kernels, eliminating multiple full-tensor memory passes.
    Parameter names are identical to Model, so state_dicts are interchangeable.
    """
    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)

        ct = self.conv_transpose
        self._rewrite = (
            all(s == 1 for s in ct.stride)
            and all(p == 0 for p in ct.padding)
            and all(d == 1 for d in ct.dilation)
            and all(o == 0 for o in ct.output_padding)
            and ct.groups == 1
        )
        self._kpad = tuple(k - 1 for k in ct.kernel_size)

    def forward(self, x):
        if self._rewrite:
            # ConvTranspose3d(s=1, p=0, d=1) == Conv3d(padding=k-1) with the
            # kernel flipped spatially and in/out channels swapped.
            w = self.conv_transpose.weight
            w = w.flip(dims=(2, 3, 4)).permute(1, 0, 2, 3, 4).contiguous()
            x = F.conv3d(x, w, self.conv_transpose.bias, stride=1,
                         padding=self._kpad)
        else:
            x = self.conv_transpose(x)

        if x.is_cuda and x.dtype == torch.float32:
            return _fused_gn.fused_relu_groupnorm(
                x, self.group_norm.weight, self.group_norm.bias,
                self.group_norm.num_groups, self.group_norm.eps)
        return F.group_norm(F.relu(x), self.group_norm.num_groups,
                            self.group_norm.weight, self.group_norm.bias,
                            self.group_norm.eps)
```