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

# Fused ReLU + GroupNorm CUDA implementation.
#
# The standard PyTorch implementation performs:
#   1) ReLU:          read N, write N
#   2) GroupNorm:     read N (stats), read N + write N (normalize)
# This fused implementation performs:
#   1) stats kernel:  read N  (ReLU applied on-the-fly, per-group mean/rstd)
#   2) affine kernel: fold mean/rstd/gamma/beta into per-(batch,channel) a,b
#   3) apply kernel:  read N, write N (y = relu(x) * a + b, vectorized float4)
# which eliminates one full read and one full write of the activation tensor.
fused_relu_groupnorm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <climits>

constexpr int STATS_BLOCK = 512;

// One block per (batch, group): computes mean and rstd of relu(x).
__global__ void relu_groupnorm_stats_kernel(
        const float* __restrict__ x,
        float* __restrict__ mean,
        float* __restrict__ rstd,
        int group_size,
        double eps) {
    const int ng = blockIdx.x;
    const float* base = x + (size_t)ng * (size_t)group_size;
    float s = 0.0f, sq = 0.0f;
    if ((group_size & 3) == 0) {
        // base is 16B aligned because group_size is a multiple of 4
        const float4* base4 = reinterpret_cast<const float4*>(base);
        const int vec = group_size >> 2;
        #pragma unroll 4
        for (int i = threadIdx.x; i < vec; i += blockDim.x) {
            float4 v = base4[i];
            float r0 = fmaxf(v.x, 0.0f);
            float r1 = fmaxf(v.y, 0.0f);
            float r2 = fmaxf(v.z, 0.0f);
            float r3 = fmaxf(v.w, 0.0f);
            s  += (r0 + r1) + (r2 + r3);
            sq += (r0 * r0 + r1 * r1) + (r2 * r2 + r3 * r3);
        }
    } else {
        for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
            float r = fmaxf(base[i], 0.0f);
            s += r;
            sq += r * r;
        }
    }
    // Block reduction in double precision for accuracy.
    double ds = (double)s;
    double dq = (double)sq;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        ds += __shfl_down_sync(0xffffffffu, ds, off);
        dq += __shfl_down_sync(0xffffffffu, dq, off);
    }
    __shared__ double sh_s[STATS_BLOCK / 32];
    __shared__ double sh_q[STATS_BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) { sh_s[wid] = ds; sh_q[wid] = dq; }
    __syncthreads();
    if (wid == 0) {
        const int nw = STATS_BLOCK / 32;
        ds = (lane < nw) ? sh_s[lane] : 0.0;
        dq = (lane < nw) ? sh_q[lane] : 0.0;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            ds += __shfl_down_sync(0xffffffffu, ds, off);
            dq += __shfl_down_sync(0xffffffffu, dq, off);
        }
        if (lane == 0) {
            double m = ds / (double)group_size;
            double var = dq / (double)group_size - m * m;
            if (var < 0.0) var = 0.0;
            mean[ng] = (float)m;
            rstd[ng] = (float)(1.0 / sqrt(var + eps));
        }
    }
}

// Fold stats and affine params into per-(batch, channel) scale/shift:
//   y = relu(x) * a[n, c] + b[n, c]
__global__ void relu_groupnorm_affine_kernel(
        const float* __restrict__ mean,
        const float* __restrict__ rstd,
        const float* __restrict__ gamma,
        const float* __restrict__ beta,
        float* __restrict__ a,
        float* __restrict__ b,
        int NC, int C, int Cg, int G) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= NC) return;
    int n = idx / C;
    int c = idx - n * C;
    int g = c / Cg;
    float m = mean[n * G + g];
    float r = rstd[n * G + g];
    float gm = gamma[c];
    float av = r * gm;
    a[idx] = av;
    b[idx] = beta[c] - m * av;
}

// Elementwise: y = relu(x) * a[n, c] + b[n, c], vectorized over float4.
template <typename index_t>
__global__ void relu_groupnorm_apply_kernel(
        const float* __restrict__ x,
        const float* __restrict__ a,
        const float* __restrict__ b,
        float* __restrict__ y,
        index_t total, int S, int C) {
    const index_t total4 = total >> 2;
    const index_t idx = (index_t)blockIdx.x * blockDim.x + threadIdx.x;
    const index_t CS = (index_t)C * (index_t)S;
    if (idx < total4) {
        const float4* x4 = reinterpret_cast<const float4*>(x);
        float4* y4 = reinterpret_cast<float4*>(y);
        float4 v = x4[idx];
        float r0 = fmaxf(v.x, 0.0f);
        float r1 = fmaxf(v.y, 0.0f);
        float r2 = fmaxf(v.z, 0.0f);
        float r3 = fmaxf(v.w, 0.0f);
        float4 o;
        if ((S & 3) == 0) {
            // All 4 lanes are guaranteed to be in the same channel.
            const index_t p = idx << 2;
            const index_t n = p / CS;
            const int c = (int)((p - n * CS) / S);
            const float av = a[n * C + c];
            const float bv = b[n * C + c];
            o.x = fmaf(r0, av, bv);
            o.y = fmaf(r1, av, bv);
            o.z = fmaf(r2, av, bv);
            o.w = fmaf(r3, av, bv);
        } else {
            const index_t p0 = idx << 2;
            const index_t p1 = p0 + 1, p2 = p0 + 2, p3 = p0 + 3;
            const index_t n0 = p0 / CS; const int c0 = (int)((p0 - n0 * CS) / S);
            const index_t n1 = p1 / CS; const int c1 = (int)((p1 - n1 * CS) / S);
            const index_t n2 = p2 / CS; const int c2 = (int)((p2 - n2 * CS) / S);
            const index_t n3 = p3 / CS; const int c3 = (int)((p3 - n3 * CS) / S);
            o.x = fmaf(r0, a[n0 * C + c0], b[n0 * C + c0]);
            o.y = fmaf(r1, a[n1 * C + c1], b[n1 * C + c1]);
            o.z = fmaf(r2, a[n2 * C + c2], b[n2 * C + c2]);
            o.w = fmaf(r3, a[n3 * C + c3], b[n3 * C + c3]);
        }
        y4[idx] = o;
    } else if (idx == total4) {
        // Handle the (up to 3) tail elements.
        for (index_t p = total4 << 2; p < total; ++p) {
            const index_t n = p / CS;
            const int c = (int)((p - n * CS) / S);
            const float r = fmaxf(x[p], 0.0f);
            y[p] = fmaf(r, a[n * C + c], b[n * C + c]);
        }
    }
}

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");
    TORCH_CHECK(x.dim() == 5, "x must have shape (N, C, D, H, W)");

    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    const int64_t G = groups;
    TORCH_CHECK(G > 0 && C % G == 0, "C must be divisible by groups");
    const int64_t Cg = C / G;
    const int64_t group_size = Cg * S;
    const int64_t total = N * C * S;
    TORCH_CHECK(group_size <= INT_MAX && S <= INT_MAX && C <= INT_MAX,
                "tensor dimensions too large");

    auto y    = at::empty_like(xc);
    auto mean = at::empty({N * G}, xc.options());
    auto rstd = at::empty({N * G}, xc.options());
    auto a    = at::empty({N * C}, xc.options());
    auto b    = at::empty({N * C}, xc.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    relu_groupnorm_stats_kernel<<<(int)(N * G), STATS_BLOCK, 0, stream>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        (int)group_size, eps);

    const int NC = (int)(N * C);
    relu_groupnorm_affine_kernel<<<(NC + 255) / 256, 256, 0, stream>>>(
        mean.data_ptr<float>(), rstd.data_ptr<float>(),
        gc.data_ptr<float>(), bc.data_ptr<float>(),
        a.data_ptr<float>(), b.data_ptr<float>(), NC, (int)C, (int)Cg, (int)G);

    const float* xp = xc.data_ptr<float>();
    const float* ap = a.data_ptr<float>();
    const float* bp = b.data_ptr<float>();
    float* yp = y.data_ptr<float>();

    if (total <= INT_MAX) {
        const int t = (int)total;
        const int64_t threads = (int64_t)(t >> 2) + 1;
        const int blocks = (int)((threads + 255) / 256);
        relu_groupnorm_apply_kernel<int><<<blocks, 256, 0, stream>>>(
            xp, ap, bp, yp, t, (int)S, (int)C);
    } else {
        const int64_t threads = (total >> 2) + 1;
        const int64_t blocks = (threads + 255) / 256;
        relu_groupnorm_apply_kernel<int64_t><<<(int)blocks, 256, 0, stream>>>(
            xp, ap, bp, yp, total, (int)S, (int)C);
    }

    return y;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN ConvTranspose3d and replaces the
    ReLU + GroupNorm chain with a single fused custom CUDA operator
    (ReLU is applied on-the-fly inside the GroupNorm kernels, so the
    intermediate ReLU tensor is never materialized).
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        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
        self.eps = self.group_norm.eps
        self.fused = fused_relu_groupnorm

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.fused.fused_relu_groupnorm(
            x, self.group_norm.weight, self.group_norm.bias, self.groups, self.eps
        )
        return x
```