```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 that operates directly on channels_last_3d (NHWC) data.
#
# Profiler analysis of the previous version showed the ConvTranspose3d internally
# transposing activations NCHW->NHWC (x2) and the output NHWC->NCHW (~670 us total).
# By running the whole model in channels_last_3d and making the fused norm kernels
# NHWC-aware, those transposes are eliminated entirely:
#   * weight kept in channels_last_3d  -> no per-call weight transpose
#   * conv output stays in NHWC        -> no output transpose
#   * norm consumes/produces NHWC      -> ReLU still fused (never materialized)
#
# Pipeline per forward:
#   1) stats kernel:   grid (splits, G, N); each block reduces a slice of one
#                      group's rows (ReLU applied on-the-fly) into double partials.
#                      Splitting each group over multiple blocks keeps the GPU
#                      saturated (N*G alone is only 128 blocks).
#   2) affine kernel:  folds partials + gamma/beta into per-(n,c) scale/shift a,b.
#   3) apply kernel:   y = relu(x)*a + b; each thread owns one channel-quad (float4)
#                      across many spatial rows, so a/b are loaded once per thread.
fused_relu_groupnorm_nhwc_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <climits>

#define STATS_BLOCK 256
#define APPLY_BLOCK 256
#define APPLY_ROWS 8

template <int BLOCK>
__device__ __forceinline__ void block_reduce2(float fs, float fq,
                                              double* out_s, double* out_q) {
    double ds = (double)fs;
    double dq = (double)fq;
    #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 shs[BLOCK / 32];
    __shared__ double shq[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) { shs[wid] = ds; shq[wid] = dq; }
    __syncthreads();
    if (wid == 0) {
        constexpr int NW = BLOCK / 32;
        ds = (lane < NW) ? shs[lane] : 0.0;
        dq = (lane < NW) ? shq[lane] : 0.0;
        #pragma unroll
        for (int off = NW / 2; off > 0; off >>= 1) {
            ds += __shfl_down_sync(0xffffffffu, ds, off);
            dq += __shfl_down_sync(0xffffffffu, dq, off);
        }
        if (lane == 0) { *out_s = ds; *out_q = dq; }
    }
}

__device__ __forceinline__ void relu4_accum(const float4 v, float& fs, float& fq) {
    const float a0 = fmaxf(v.x, 0.0f);
    const float a1 = fmaxf(v.y, 0.0f);
    const float a2 = fmaxf(v.z, 0.0f);
    const float a3 = fmaxf(v.w, 0.0f);
    fs += (a0 + a1) + (a2 + a3);
    fq += (a0 * a0 + a1 * a1) + (a2 * a2 + a3 * a3);
}

// ------------------------- stats kernels -------------------------
// grid = (splits, G, N). Each block reduces rows [s0, s1) of one group.
// V = Cg/4 float4s per spatial row (Cg channels per group).

template <int VSHIFT>
__global__ void stats_pow2_kernel(const float* __restrict__ x,
                                  double* __restrict__ partials,
                                  int S, int C, int Cg, int G,
                                  int rows_per_split, int splits) {
    constexpr int V = 1 << VSHIFT;
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    const int c4 = threadIdx.x & (V - 1);
    const int r0 = threadIdx.x >> VSHIFT;
    const int rstride = STATS_BLOCK >> VSHIFT;
    float fs = 0.0f, fq = 0.0f;
    for (int r = r0; r < s1 - s0; r += rstride) {
        const float4 v = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + r) * C + (c4 << 2));
        relu4_accum(v, fs, fq);
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void stats_vec4_kernel(const float* __restrict__ x,
                                  double* __restrict__ partials,
                                  int S, int C, int Cg, int G, int V,
                                  int rows_per_split, int splits) {
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    float fs = 0.0f, fq = 0.0f;
    const int count = (s1 - s0) * V;
    for (int idx = threadIdx.x; idx < count; idx += STATS_BLOCK) {
        const int row = idx / V;
        const int c4  = idx - row * V;
        const float4 v = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + row) * C + (c4 << 2));
        relu4_accum(v, fs, fq);
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void stats_scalar_kernel(const float* __restrict__ x,
                                    double* __restrict__ partials,
                                    int S, int C, int Cg, int G,
                                    int rows_per_split, int splits) {
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    float fs = 0.0f, fq = 0.0f;
    for (int r = s0; r < s1; ++r) {
        const float* rowp = base + (size_t)r * C;
        for (int c = threadIdx.x; c < Cg; c += STATS_BLOCK) {
            const float v = fmaxf(rowp[c], 0.0f);
            fs += v;
            fq += v * v;
        }
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

// ---------------------- finalize + affine kernel ----------------------
// Folds the per-split partials into mean/rstd and then into per-(n,c)
// scale a and shift b:  y = relu(x) * a[n,c] + b[n,c].
__global__ void finalize_affine_kernel(const double* __restrict__ partials,
                                       const float* __restrict__ gamma,
                                       const float* __restrict__ beta,
                                       float* __restrict__ a,
                                       float* __restrict__ b,
                                       int NC, int C, int Cg, int G,
                                       int splits, double inv_gs, double eps) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= NC) return;
    const int n = idx / C;
    const int c = idx - n * C;
    const int g = c / Cg;
    const double* p = partials + (size_t)(n * G + g) * splits * 2;
    double s = 0.0, q = 0.0;
    for (int j = 0; j < splits; ++j) {
        s += p[2 * j];
        q += p[2 * j + 1];
    }
    const double m = s * inv_gs;
    double var = q * inv_gs - m * m;
    if (var < 0.0) var = 0.0;
    const double rstd = 1.0 / sqrt(var + eps);
    const float av = (float)(rstd * (double)gamma[c]);
    a[idx] = av;
    b[idx] = (float)((double)beta[c] - m * (double)av);
}

// ------------------------- apply kernels -------------------------
// grid = (blocks_x, N). Each thread owns one channel-quad (c4) of image n
// and walks spatial rows with a fixed stride, so a/b are loaded once.

template <int VSHIFT>
__global__ void apply_pow2_kernel(const float* __restrict__ x,
                                  const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int S, int C) {
    constexpr int V = 1 << VSHIFT;  // = C/4
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c4  = tid & (V - 1);
    const int s0  = tid >> VSHIFT;
    const int sstride = (gridDim.x * APPLY_BLOCK) >> VSHIFT;
    if (s0 >= S) return;
    const size_t img = (size_t)n * S * V;
    const float4 av = *reinterpret_cast<const float4*>(a + (size_t)n * C + (c4 << 2));
    const float4 bv = *reinterpret_cast<const float4*>(b + (size_t)n * C + (c4 << 2));
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * V + c4;
        const float4 v = *reinterpret_cast<const float4*>(x + (off << 2));
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), av.x, bv.x);
        o.y = fmaf(fmaxf(v.y, 0.0f), av.y, bv.y);
        o.z = fmaf(fmaxf(v.z, 0.0f), av.z, bv.z);
        o.w = fmaf(fmaxf(v.w, 0.0f), av.w, bv.w);
        *reinterpret_cast<float4*>(y + (off << 2)) = o;
    }
}

__global__ void apply_vec4_kernel(const float* __restrict__ x,
                                  const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int S, int C, int V) {
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c4  = tid % V;
    const int s0  = tid / V;
    const int sstride = (gridDim.x * APPLY_BLOCK) / V;
    if (s0 >= S || sstride <= 0) return;
    const size_t img = (size_t)n * S * V;
    const float4 av = *reinterpret_cast<const float4*>(a + (size_t)n * C + (c4 << 2));
    const float4 bv = *reinterpret_cast<const float4*>(b + (size_t)n * C + (c4 << 2));
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * V + c4;
        const float4 v = *reinterpret_cast<const float4*>(x + (off << 2));
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), av.x, bv.x);
        o.y = fmaf(fmaxf(v.y, 0.0f), av.y, bv.y);
        o.z = fmaf(fmaxf(v.z, 0.0f), av.z, bv.z);
        o.w = fmaf(fmaxf(v.w, 0.0f), av.w, bv.w);
        *reinterpret_cast<float4*>(y + (off << 2)) = o;
    }
}

__global__ void apply_scalar_kernel(const float* __restrict__ x,
                                    const float* __restrict__ a,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    size_t total, int C) {
    const int n = blockIdx.y;
    const size_t base = (size_t)n * total;
    const size_t ab   = (size_t)n * C;
    const size_t stride = (size_t)gridDim.x * APPLY_BLOCK;
    for (size_t p = (size_t)blockIdx.x * APPLY_BLOCK + threadIdx.x;
         p < total; p += stride) {
        const int c = (int)(p % (size_t)C);
        const float v = fmaxf(x[base + p], 0.0f);
        y[base + p] = fmaf(v, a[ab + c], b[ab + c]);
    }
}

// ------------------------------ host ------------------------------
static inline int log2_int(int v) {
    int s = 0;
    while (v > 1) { v >>= 1; ++s; }
    return s;
}

torch::Tensor fused_relu_groupnorm_nhwc(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(at::MemoryFormat::ChannelsLast3d);
    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;
    TORCH_CHECK(S > 0 && S <= INT_MAX && C <= INT_MAX && Cg <= INT_MAX,
                "bad tensor dims");
    TORCH_CHECK(N <= 65535 && G <= 65535, "N/G too large for grid");
    TORCH_CHECK(gc.numel() == C && bc.numel() == C, "gamma/beta size mismatch");

    auto y = at::empty_like(xc);  // preserves channels_last_3d layout
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // ---- stats ----
    const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t NG = N * G;
    int splits = (int)((8LL * sm_count + NG - 1) / NG);
    if (splits < 1) splits = 1;
    if (splits > 32) splits = 32;
    if ((int64_t)splits > S) splits = (int)S;
    int rows_per_split = (int)((S + splits - 1) / splits);
    splits = (int)((S + rows_per_split - 1) / rows_per_split);

    auto partials = at::empty({NG * splits * 2},
                              xc.options().dtype(at::kDouble));
    const dim3 sgrid(splits, (unsigned)G, (unsigned)N);
    const float* xp = xc.data_ptr<float>();
    double* pp = partials.data_ptr<double>();

    const bool vec_ok = (C % 4 == 0) && (Cg % 4 == 0);
    const int Vs = (int)(Cg / 4);
    const bool vpow2 = vec_ok && Vs > 0 && (Vs & (Vs - 1)) == 0;
    if (vec_ok && vpow2) {
        const int sh = log2_int(Vs);
        switch (sh) {
            case 0: stats_pow2_kernel<0><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 1: stats_pow2_kernel<1><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 2: stats_pow2_kernel<2><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 3: stats_pow2_kernel<3><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 4: stats_pow2_kernel<4><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 5: stats_pow2_kernel<5><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 6: stats_pow2_kernel<6><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            default: stats_vec4_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, Vs, rows_per_split, splits); break;
        }
    } else if (vec_ok) {
        stats_vec4_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, Vs, rows_per_split, splits);
    } else {
        stats_scalar_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits);
    }

    // ---- finalize + affine ----
    const int NC = (int)(N * C);
    auto a = at::empty({NC}, xc.options());
    auto b = at::empty({NC}, xc.options());
    finalize_affine_kernel<<<(NC + 255) / 256, 256, 0, stream>>>(
        pp, gc.data_ptr<float>(), bc.data_ptr<float>(),
        a.data_ptr<float>(), b.data_ptr<float>(),
        NC, (int)C, (int)Cg, (int)G, splits,
        1.0 / (double)(Cg * S), eps);

    // ---- apply ----
    float* yp = y.data_ptr<float>();
    const float* ap = a.data_ptr<float>();
    const float* bp = b.data_ptr<float>();
    if (C % 4 == 0) {
        const int V = (int)(C / 4);
        int64_t blocks_x = (S * V + (int64_t)APPLY_BLOCK * APPLY_ROWS - 1) /
                           ((int64_t)APPLY_BLOCK * APPLY_ROWS);
        const int64_t min_blocks = (V + APPLY_BLOCK - 1) / APPLY_BLOCK;
        if (blocks_x < min_blocks) blocks_x = min_blocks;
        if (blocks_x < 1) blocks_x = 1;
        const dim3 agrid((unsigned)blocks_x, (unsigned)N, 1);
        const bool vp2 = (V & (V - 1)) == 0;
        if (vp2) {
            const int sh = log2_int(V);
            switch (sh) {
                case 0: apply_pow2_kernel<0><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 1: apply_pow2_kernel<1><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 2: apply_pow2_kernel<2><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 3: apply_pow2_kernel<3><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 4: apply_pow2_kernel<4><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 5: apply_pow2_kernel<5><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 6: apply_pow2_kernel<6><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 7: apply_pow2_kernel<7><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                default: apply_vec4_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C, V); break;
            }
        } else {
            apply_vec4_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C, V);
        }
    } else {
        const size_t total = (size_t)S * (size_t)C;
        int64_t blocks_x = ((int64_t)total + (int64_t)APPLY_BLOCK * APPLY_ROWS - 1) /
                           ((int64_t)APPLY_BLOCK * APPLY_ROWS);
        if (blocks_x < 1) blocks_x = 1;
        const dim3 agrid((unsigned)blocks_x, (unsigned)N, 1);
        apply_scalar_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, total, (int)C);
    }

    return y;
}
"""

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

fused_relu_groupnorm_nhwc = load_inline(
    name="fused_relu_groupnorm_nhwc_ext",
    cpp_sources=fused_relu_groupnorm_nhwc_cpp_source,
    cuda_sources=fused_relu_groupnorm_nhwc_cuda_source,
    functions=["fused_relu_groupnorm_nhwc"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: runs ConvTranspose3d fully in channels_last_3d so cuDNN
    uses its NHWC tensor-core kernels without any layout transposes, and the
    ReLU + GroupNorm chain is a single fused custom CUDA operator that works
    directly on the NHWC activations (ReLU is computed on-the-fly and 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)
        with torch.no_grad():
            # Keep weights in channels_last_3d so cuDNN does not re-transpose
            # them on every call.
            self.conv_transpose.weight.data = self.conv_transpose.weight.data.to(
                memory_format=torch.channels_last_3d
            )
        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_nhwc

    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last_3d)
        x = self.conv_transpose(x)
        x = self.fused.fused_relu_groupnorm_nhwc(
            x, self.group_norm.weight, self.group_norm.bias, self.groups, self.eps
        )
        return x
```