```
# 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 CUDA kernels: fused ReLU + GroupNorm.
#
# Strategy:
#   1) The ConvTranspose3d (stride=1, pad=0) is mathematically identical to a
#      regular Conv3d with the kernel spatially flipped, in/out channels
#      swapped, and padding = K-1.  Regular forward convolutions hit better
#      tuned cuDNN kernels than transposed convolutions.
#   2) ReLU and GroupNorm are fused into two custom kernels:
#        - gn_stats_kernel: computes per-(batch,group) sum & sum-of-squares of
#          relu(x) using vectorized float4 loads, warp/block reductions and
#          double-precision atomics.  The ReLU output is never materialized.
#        - gn_apply_kernel: recomputes relu(x) on the fly and applies the
#          normalization with the per-channel affine parameters folded into a
#          single scale (a) and shift (b): y = relu(x)*a + b.
#      This eliminates one full read+write pass over the activation tensor
#      compared to the unfused PyTorch implementation.
# ---------------------------------------------------------------------------

fused_gn_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

__global__ void gn_stats_kernel(const float4* __restrict__ x,
                                double* __restrict__ accum,
                                int ch4,
                                int channels_per_group) {
    const int global_chan = blockIdx.x;              // in [0, B*C)
    const int g = global_chan / channels_per_group;  // in [0, B*groups)
    const float4* base = x + (size_t)global_chan * (size_t)ch4;

    float s = 0.f, q = 0.f;
    for (int i = threadIdx.x; i < ch4; i += blockDim.x) {
        float4 v = base[i];
        v.x = fmaxf(v.x, 0.f);
        v.y = fmaxf(v.y, 0.f);
        v.z = fmaxf(v.z, 0.f);
        v.w = fmaxf(v.w, 0.f);
        s += (v.x + v.y) + (v.z + v.w);
        q += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
    }

    __shared__ float sh_s[32];
    __shared__ float sh_q[32];
    for (int o = 16; o > 0; o >>= 1) {
        s += __shfl_down_sync(0xffffffffu, s, o);
        q += __shfl_down_sync(0xffffffffu, q, o);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) { sh_s[warp] = s; sh_q[warp] = q; }
    __syncthreads();
    const int nwarp = (blockDim.x + 31) >> 5;
    if (warp == 0) {
        s = (lane < nwarp) ? sh_s[lane] : 0.f;
        q = (lane < nwarp) ? sh_q[lane] : 0.f;
        for (int o = 16; o > 0; o >>= 1) {
            s += __shfl_down_sync(0xffffffffu, s, o);
            q += __shfl_down_sync(0xffffffffu, q, o);
        }
        if (lane == 0) {
            atomicAdd(accum + 2 * g,     (double)s);
            atomicAdd(accum + 2 * g + 1, (double)q);
        }
    }
}

__global__ void gn_apply_kernel(const float4* __restrict__ x,
                                float4* __restrict__ y,
                                const double* __restrict__ accum,
                                const float* __restrict__ gamma,
                                const float* __restrict__ beta,
                                int ch4,
                                int channels_per_group,
                                int channels,
                                double inv_n,
                                double eps) {
    const int global_chan = blockIdx.x;
    const int g = global_chan / channels_per_group;
    const int c = global_chan % channels;

    const double sum = accum[2 * g];
    const double ssq = accum[2 * g + 1];
    const double mean = sum * inv_n;
    double var = ssq * inv_n - mean * mean;
    if (var < 0.0) var = 0.0;
    const float rstd = (float)rsqrt(var + eps);
    const float gm = gamma[c];
    const float a = rstd * gm;
    const float b = beta[c] - (float)mean * rstd * gm;

    const float4* in  = x + (size_t)global_chan * (size_t)ch4;
    float4* out       = y + (size_t)global_chan * (size_t)ch4;
    for (int i = threadIdx.x; i < ch4; i += blockDim.x) {
        float4 v = in[i];
        v.x = fmaxf(v.x, 0.f) * a + b;
        v.y = fmaxf(v.y, 0.f) * a + b;
        v.z = fmaxf(v.z, 0.f) * a + b;
        v.w = fmaxf(v.w, 0.f) * a + b;
        out[i] = v;
    }
}

torch::Tensor relu_groupnorm(torch::Tensor x,
                             torch::Tensor gamma,
                             torch::Tensor beta,
                             int64_t groups,
                             double eps) {
    TORCH_CHECK(x.is_cuda() && x.is_contiguous() &&
                x.scalar_type() == torch::kFloat);
    const int B = (int)x.size(0);
    const int C = (int)x.size(1);
    long inner = 1;
    for (int i = 2; i < x.dim(); ++i) inner *= x.size(i);
    TORCH_CHECK(inner % 4 == 0);
    const int ch4 = (int)(inner / 4);
    const int cpg = (int)(C / groups);
    const int nblocks = B * C;

    auto y = torch::empty_like(x);
    auto accum = torch::zeros({(long)(2 * B * groups)},
                              x.options().dtype(torch::kDouble));

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

    gn_stats_kernel<<<nblocks, threads, 0, stream>>>(
        reinterpret_cast<const float4*>(x.data_ptr<float>()),
        accum.data_ptr<double>(),
        ch4, cpg);

    const double inv_n = 1.0 / (double)(inner * (long)cpg);

    gn_apply_kernel<<<nblocks, threads, 0, stream>>>(
        reinterpret_cast<const float4*>(x.data_ptr<float>()),
        reinterpret_cast<float4*>(y.data_ptr<float>()),
        accum.data_ptr<double>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        ch4, cpg, C, inv_n, eps);

    return y;
}
"""

fused_gn_cpp_source = (
    "torch::Tensor 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",
    cpp_sources=fused_gn_cpp_source,
    cuda_sources=fused_gn_cuda_source,
    functions=["relu_groupnorm"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - ConvTranspose3d (stride=1, pad=0) is rewritten as an equivalent regular
        Conv3d with flipped/swapped kernels (better cuDNN kernel selection).
      - ReLU + GroupNorm are fused into custom CUDA kernels that never
        materialize the ReLU output.
    """

    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, (tuple, list)):
            self._pad = tuple(k - 1 for k in kernel_size)
        else:
            self._pad = kernel_size - 1

    def forward(self, x):
        # ConvTranspose3d(stride=1, pad=0) == Conv3d with flipped kernel,
        # swapped in/out channels and padding = K - 1.
        w = self.conv_transpose.weight
        w_conv = w.permute(1, 0, 2, 3, 4).flip([2, 3, 4]).contiguous()
        x = F.conv3d(x, w_conv, self.conv_transpose.bias, 1, self._pad)

        inner = x.shape[2] * x.shape[3] * x.shape[4]
        if (
            inner % 4 == 0
            and x.is_contiguous()
            and x.dtype == torch.float32
            and x.is_cuda
        ):
            return fused_relu_groupnorm.relu_groupnorm(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.groups,
                self.group_norm.eps,
            )

        # Fallback for shapes/dtypes the custom kernel does not support.
        x = F.relu(x)
        return self.group_norm(x)
```