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

# Fused ReLU + GroupNorm CUDA kernels.
# The ConvTranspose3d is left to cuDNN (it already uses tensor cores / tuned
# implicit-GEMM kernels). The elementwise ReLU and the GroupNorm are memory
# bound, so we fuse them into two custom kernels:
#   1) stats kernel: computes per-(batch,group) sum / sum-of-squares of relu(x)
#   2) apply kernel: re-reads x, applies relu + normalization + affine
# This eliminates the separate ReLU read+write and one GroupNorm read.
relu_groupnorm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

#define TPB 256
#define NWARP (TPB / 32)

__device__ __forceinline__ double warpReduceSum(double v) {
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;
}

__device__ __forceinline__ void blockReduceAtomicAdd(double s, double q,
                                                     double* __restrict__ sums,
                                                     double* __restrict__ sumsqs,
                                                     int ng) {
    __shared__ double sh_s[NWARP];
    __shared__ double sh_q[NWARP];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    s = warpReduceSum(s);
    q = warpReduceSum(q);
    if (lane == 0) { sh_s[wid] = s; sh_q[wid] = q; }
    __syncthreads();
    if (wid == 0) {
        s = (lane < NWARP) ? sh_s[lane] : 0.0;
        q = (lane < NWARP) ? sh_q[lane] : 0.0;
        s = warpReduceSum(s);
        q = warpReduceSum(q);
        if (lane == 0) {
            atomicAdd(&sums[ng], s);
            atomicAdd(&sumsqs[ng], q);
        }
    }
}

// ------------------------- stats kernels -------------------------
__global__ void __launch_bounds__(TPB)
stats_vec4_kernel(const float4* __restrict__ x,
                  double* __restrict__ sums,
                  double* __restrict__ sumsqs,
                  int v4_per_group,
                  int chunk_len) {
    const int ng = blockIdx.y;
    const int start = blockIdx.x * chunk_len;
    int end = start + chunk_len;
    if (end > v4_per_group) end = v4_per_group;
    const float4* xg = x + (size_t)ng * (size_t)v4_per_group;

    double s = 0.0, q = 0.0;
    for (int i = start + threadIdx.x; i < end; i += TPB) {
        float4 v = xg[i];
        float a = fmaxf(v.x, 0.f);
        float b = fmaxf(v.y, 0.f);
        float c = fmaxf(v.z, 0.f);
        float d = fmaxf(v.w, 0.f);
        s += (double)(a + b + c + d);
        q += (double)(a * a + b * b + c * c + d * d);
    }
    blockReduceAtomicAdd(s, q, sums, sumsqs, ng);
}

__global__ void __launch_bounds__(TPB)
stats_scalar_kernel(const float* __restrict__ x,
                    double* __restrict__ sums,
                    double* __restrict__ sumsqs,
                    int64_t size_per_group,
                    int64_t chunk_len) {
    const int64_t ng = blockIdx.y;
    const int64_t start = (int64_t)blockIdx.x * chunk_len;
    int64_t end = start + chunk_len;
    if (end > size_per_group) end = size_per_group;
    const float* xg = x + (size_t)ng * (size_t)size_per_group;

    double s = 0.0, q = 0.0;
    for (int64_t i = start + threadIdx.x; i < end; i += TPB) {
        float v = fmaxf(xg[i], 0.f);
        s += (double)v;
        q += (double)v * (double)v;
    }
    blockReduceAtomicAdd(s, q, sums, sumsqs, (int)ng);
}

// ------------------------- apply kernels -------------------------
// Requires spatial % 4 == 0 so a float4 never spans two channels.
__global__ void __launch_bounds__(TPB)
apply_vec4_kernel(const float4* __restrict__ x,
                  const float* __restrict__ gamma,
                  const float* __restrict__ beta,
                  const double* __restrict__ sums,
                  const double* __restrict__ sumsqs,
                  float4* __restrict__ out,
                  int v4_per_group,
                  int chunk_len,
                  int channels_per_group,
                  int spatial,
                  int groups,
                  float eps) {
    const int ng = blockIdx.y;
    const int g = ng % groups;
    const double n = (double)v4_per_group * 4.0;
    double mean = sums[ng] / n;
    double var = sumsqs[ng] / n - mean * mean;
    if (var < 0.0) var = 0.0;
    const float rstd = rsqrtf((float)var + eps);
    const float fmean = (float)mean;

    const int start = blockIdx.x * chunk_len;
    int end = start + chunk_len;
    if (end > v4_per_group) end = v4_per_group;

    const float4* xg = x + (size_t)ng * (size_t)v4_per_group;
    float4* og = out + (size_t)ng * (size_t)v4_per_group;
    const float* gg = gamma + g * channels_per_group;
    const float* gb = beta + g * channels_per_group;

    int i = start + threadIdx.x;
    if (i >= end) return;

    // Incremental channel tracking (one division per thread only).
    int c = (int)(((int64_t)i * 4) / spatial);
    int64_t next = (int64_t)(c + 1) * (int64_t)spatial;
    for (; i < end; i += TPB) {
        int64_t elem = (int64_t)i * 4;
        while (elem >= next) { ++c; next += spatial; }
        const float w = gg[c];
        const float b = gb[c];
        float4 v = xg[i];
        float4 r;
        r.x = (fmaxf(v.x, 0.f) - fmean) * rstd * w + b;
        r.y = (fmaxf(v.y, 0.f) - fmean) * rstd * w + b;
        r.z = (fmaxf(v.z, 0.f) - fmean) * rstd * w + b;
        r.w = (fmaxf(v.w, 0.f) - fmean) * rstd * w + b;
        og[i] = r;
    }
}

// Generic fallback: block handles a chunk of a single channel (no division).
__global__ void __launch_bounds__(TPB)
apply_scalar_kernel(const float* __restrict__ x,
                    const float* __restrict__ gamma,
                    const float* __restrict__ beta,
                    const double* __restrict__ sums,
                    const double* __restrict__ sumsqs,
                    float* __restrict__ out,
                    int64_t size_per_group,
                    int channels_per_group,
                    int64_t spatial,
                    int groups,
                    float eps,
                    int64_t chunk_len) {
    const int64_t ng = blockIdx.z;
    const int c = blockIdx.y;
    const int g = (int)(ng % groups);
    double mean = sums[ng] / (double)size_per_group;
    double var = sumsqs[ng] / (double)size_per_group - mean * mean;
    if (var < 0.0) var = 0.0;
    const float rstd = rsqrtf((float)var + eps);
    const float fmean = (float)mean;
    const float w = gamma[g * channels_per_group + c];
    const float b = beta[g * channels_per_group + c];

    const size_t base = ((size_t)ng * (size_t)channels_per_group + (size_t)c) * (size_t)spatial;
    const float* xb = x + base;
    float* ob = out + base;

    const int64_t start = (int64_t)blockIdx.x * chunk_len;
    int64_t end = start + chunk_len;
    if (end > spatial) end = spatial;
    for (int64_t i = start + threadIdx.x; i < end; i += TPB) {
        float v = fmaxf(xb[i], 0.f);
        ob[i] = (v - fmean) * rstd * w + b;
    }
}

torch::Tensor relu_groupnorm_cuda(torch::Tensor x,
                                  torch::Tensor weight,
                                  torch::Tensor bias,
                                  int64_t groups,
                                  double eps) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);
    const int64_t numel = xc.numel();
    if (numel == 0) return out;

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t spatial = numel / (N * C);
    const int64_t cpg = C / groups;
    const int64_t size_per_group = cpg * spatial;
    const int64_t ng = N * groups;

    auto dopts = torch::TensorOptions().dtype(torch::kDouble).device(xc.device());
    auto sums = torch::zeros({ng}, dopts);
    auto sumsqs = torch::zeros({ng}, dopts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    float* op = out.data_ptr<float>();
    const float* gp = weight.contiguous().data_ptr<float>();
    const float* bp = bias.contiguous().data_ptr<float>();
    double* sp = sums.data_ptr<double>();
    double* qp = sumsqs.data_ptr<double>();

    const bool vec_ok = (size_per_group % 4 == 0) && (spatial % 4 == 0) &&
                        (size_per_group < (int64_t)INT32_MAX) &&
                        (ng <= 65535);
    if (vec_ok) {
        const int v4 = (int)(size_per_group / 4);
        const int chunk = 8192; // in float4 units
        const int chunks = (v4 + chunk - 1) / chunk;
        dim3 grid(chunks, (unsigned)ng);
        stats_vec4_kernel<<<grid, TPB, 0, stream>>>(
            reinterpret_cast<const float4*>(xp), sp, qp, v4, chunk);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        apply_vec4_kernel<<<grid, TPB, 0, stream>>>(
            reinterpret_cast<const float4*>(xp), gp, bp, sp, qp,
            reinterpret_cast<float4*>(op), v4, chunk, (int)cpg,
            (int)spatial, (int)groups, (float)eps);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        const int64_t chunk = 8192;
        const int schunks = (int)((size_per_group + chunk - 1) / chunk);
        dim3 sgrid(schunks, (unsigned)ng);
        stats_scalar_kernel<<<sgrid, TPB, 0, stream>>>(
            xp, sp, qp, size_per_group, chunk);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        const int achunks = (int)((spatial + chunk - 1) / chunk);
        dim3 agrid(achunks, (unsigned)cpg, (unsigned)ng);
        apply_scalar_kernel<<<agrid, TPB, 0, stream>>>(
            xp, gp, bp, sp, qp, op, size_per_group, (int)cpg,
            spatial, (int)groups, (float)eps, chunk);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return out;
}
"""

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

fused_relu_groupnorm = load_inline(
    name="fused_relu_groupnorm",
    cpp_sources=relu_groupnorm_cpp_source,
    cuda_sources=relu_groupnorm_cuda_source,
    functions=["relu_groupnorm_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: ConvTranspose3d stays on cuDNN (tensor-core tuned), while
    ReLU + GroupNorm are replaced by a fused custom CUDA kernel.
    """
    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
        # Let cuDNN autotune the transposed convolution algorithm.
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True

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