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

# Optimized pipeline (all activations kept in channels_last_3d end-to-end):
#   0) Fast custom NCHW->NHWC transpose for the input (the harness feeds NCHW).
#   1) ConvTranspose3d rewritten as an equivalent forward Conv3d (stride-1,
#      pad-0 transposed conv == conv3d with flipped/swapped kernel, pad=k-1),
#      so cuDNN uses the fast NHWC implicit-GEMM forward path. Verified once at
#      runtime against the real conv_transpose with a permanent fallback.
#   2) Optional fp16 tensor-core convolution path: input is transposed+cast to
#      fp16 NHWC in one kernel, the conv runs in fp16 (2x tensor-core
#      throughput over TF32), and the fused ReLU+GroupNorm kernels read the
#      fp16 activations directly while accumulating statistics in fp32/fp64
#      and producing an fp32 output. Enabled only after a runtime accuracy
#      check (vs the fp32 pipeline) AND a runtime speed check both pass;
#      otherwise the proven fp32 path is used permanently.
#   3) ReLU+GroupNorm fused: stats kernel (ReLU applied on the fly), a tiny
#      finalize kernel producing per-(n,c) scale/shift, and a vectorized
#      apply kernel. The ReLU output is never materialized.
fastpath_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <climits>

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

struct alignas(16) H8 { __half2 h[4]; };

// ------------------------- NCHW -> NHWC transposes -------------------------
__global__ void nchw2nhwc_kernel(const float* __restrict__ src,
                                 float* __restrict__ dst,
                                 int C, int S) {
    __shared__ float tile[32][33];
    const int n  = blockIdx.z;
    const int c0 = blockIdx.x * 32;
    const int s0 = blockIdx.y * 32;
    const size_t base = (size_t)n * C * S;
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int c = c0 + ty + i;
        const int s = s0 + tx;
        if (c < C && s < S) tile[ty + i][tx] = src[base + (size_t)c * S + s];
    }
    __syncthreads();
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int s = s0 + ty + i;
        const int c = c0 + tx;
        if (s < S && c < C) dst[base + (size_t)s * C + c] = tile[tx][ty + i];
    }
}

__global__ void nchw2nhwc_half_kernel(const float* __restrict__ src,
                                      __half* __restrict__ dst,
                                      int C, int S) {
    __shared__ float tile[32][33];
    const int n  = blockIdx.z;
    const int c0 = blockIdx.x * 32;
    const int s0 = blockIdx.y * 32;
    const size_t base = (size_t)n * C * S;
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int c = c0 + ty + i;
        const int s = s0 + tx;
        if (c < C && s < S) tile[ty + i][tx] = src[base + (size_t)c * S + s];
    }
    __syncthreads();
    const size_t dbase = (size_t)n * C * S;
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int s = s0 + ty + i;
        const int c = c0 + tx;
        if (s < S && c < C) {
            dst[dbase + (size_t)s * C + c] = __float2half(tile[tx][ty + i]);
        }
    }
}

torch::Tensor nchw_to_nhwc(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat && x.dim() == 5,
                "bad input");
    auto xc = x.contiguous();
    const int64_t N = xc.size(0), C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    TORCH_CHECK(S <= INT_MAX && C <= INT_MAX && N <= 65535, "too big");
    auto y = at::empty_like(xc, at::MemoryFormat::ChannelsLast3d);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const dim3 block(32, 8, 1);
    const dim3 grid((unsigned)((C + 31) / 32), (unsigned)((S + 31) / 32),
                    (unsigned)N);
    nchw2nhwc_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), y.data_ptr<float>(), (int)C, (int)S);
    return y;
}

torch::Tensor nchw_to_nhwc_half(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat && x.dim() == 5,
                "bad input");
    auto xc = x.contiguous();
    const int64_t N = xc.size(0), C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    TORCH_CHECK(S <= INT_MAX && C <= INT_MAX && N <= 65535, "too big");
    auto y = at::empty(xc.sizes(),
                       xc.options().dtype(at::kHalf).memory_format(
                           at::MemoryFormat::ChannelsLast3d));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const dim3 block(32, 8, 1);
    const dim3 grid((unsigned)((C + 31) / 32), (unsigned)((S + 31) / 32),
                    (unsigned)N);
    nchw2nhwc_half_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), reinterpret_cast<__half*>(y.data_ptr<at::Half>()),
        (int)C, (int)S);
    return y;
}

// --------------------------- shared reduce helper --------------------------
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);
}

__device__ __forceinline__ void relu8_accum_h(const H8& v, float& fs, float& fq) {
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        const float2 f = __half22float2(v.h[i]);
        const float a0 = fmaxf(f.x, 0.0f);
        const float a1 = fmaxf(f.y, 0.0f);
        fs += a0 + a1;
        fq += fmaf(a0, a0, a1 * a1);
    }
}

// --------------------------- fp32 stats kernels ----------------------------
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;
    const int rows = s1 - s0;
    int r = r0;
    for (; r + rstride < rows; r += 2 * rstride) {
        const float4 v0 = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + r) * C + (c4 << 2));
        const float4 v1 = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + r + rstride) * C + (c4 << 2));
        relu4_accum(v0, fs, fq);
        relu4_accum(v1, fs, fq);
    }
    for (; r < rows; 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;
    }
}

// --------------------------- fp16 stats kernels ----------------------------
template <int VSHIFT>
__global__ void stats_h_pow2_kernel(const __half* __restrict__ x,
                                    double* __restrict__ partials,
                                    int S, int C, int Cg, int G,
                                    int rows_per_split, int splits) {
    constexpr int V = 1 << VSHIFT;  // octets of halfs per group = Cg/8
    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 __half* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    const int c8 = threadIdx.x & (V - 1);
    const int r0 = threadIdx.x >> VSHIFT;
    const int rstride = STATS_BLOCK >> VSHIFT;
    float fs = 0.0f, fq = 0.0f;
    const int rows = s1 - s0;
    int r = r0;
    for (; r + rstride < rows; r += 2 * rstride) {
        const H8 v0 = *reinterpret_cast<const H8*>(
            base + (size_t)(s0 + r) * C + (c8 << 3));
        const H8 v1 = *reinterpret_cast<const H8*>(
            base + (size_t)(s0 + r + rstride) * C + (c8 << 3));
        relu8_accum_h(v0, fs, fq);
        relu8_accum_h(v1, fs, fq);
    }
    for (; r < rows; r += rstride) {
        const H8 v = *reinterpret_cast<const H8*>(
            base + (size_t)(s0 + r) * C + (c8 << 3));
        relu8_accum_h(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_h_vec_kernel(const __half* __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 __half* 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 c8  = idx - row * V;
        const H8 v = *reinterpret_cast<const H8*>(
            base + (size_t)(s0 + row) * C + (c8 << 3));
        relu8_accum_h(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;
    }
}

// ------------------------------ finalize ----------------------------------
__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);
}

// --------------------------- fp32 apply kernels ----------------------------
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;
    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]);
    }
}

// --------------------------- fp16 apply kernels ----------------------------
template <int VSHIFT>
__global__ void apply_h_pow2_kernel(const __half* __restrict__ x,
                                    const float* __restrict__ a,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    int S, int C) {
    constexpr int V = 1 << VSHIFT;  // octets per row = C/8
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c8  = tid & (V - 1);
    const int s0  = tid >> VSHIFT;
    const int sstride = (gridDim.x * APPLY_BLOCK) >> VSHIFT;
    if (s0 >= S) return;
    const float* ap = a + (size_t)n * C + (c8 << 3);
    const float* bp = b + (size_t)n * C + (c8 << 3);
    const float4 a0 = *reinterpret_cast<const float4*>(ap);
    const float4 a1 = *reinterpret_cast<const float4*>(ap + 4);
    const float4 b0 = *reinterpret_cast<const float4*>(bp);
    const float4 b1 = *reinterpret_cast<const float4*>(bp + 4);
    const size_t img = (size_t)n * S * C;
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * C + (c8 << 3);
        const H8 v = *reinterpret_cast<const H8*>(x + off);
        const float2 f0 = __half22float2(v.h[0]);
        const float2 f1 = __half22float2(v.h[1]);
        const float2 f2 = __half22float2(v.h[2]);
        const float2 f3 = __half22float2(v.h[3]);
        float4 o0, o1;
        o0.x = fmaf(fmaxf(f0.x, 0.0f), a0.x, b0.x);
        o0.y = fmaf(fmaxf(f0.y, 0.0f), a0.y, b0.y);
        o0.z = fmaf(fmaxf(f1.x, 0.0f), a0.z, b0.z);
        o0.w = fmaf(fmaxf(f1.y, 0.0f), a0.w, b0.w);
        o1.x = fmaf(fmaxf(f2.x, 0.0f), a1.x, b1.x);
        o1.y = fmaf(fmaxf(f2.y, 0.0f), a1.y, b1.y);
        o1.z = fmaf(fmaxf(f3.x, 0.0f), a1.z, b1.z);
        o1.w = fmaf(fmaxf(f3.y, 0.0f), a1.w, b1.w);
        *reinterpret_cast<float4*>(y + off) = o0;
        *reinterpret_cast<float4*>(y + off + 4) = o1;
    }
}

__global__ void apply_h_vec_kernel(const __half* __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 c8  = tid % V;
    const int s0  = tid / V;
    const int sstride = (gridDim.x * APPLY_BLOCK) / V;
    if (s0 >= S || sstride <= 0) return;
    const float* ap = a + (size_t)n * C + (c8 << 3);
    const float* bp = b + (size_t)n * C + (c8 << 3);
    const float4 a0 = *reinterpret_cast<const float4*>(ap);
    const float4 a1 = *reinterpret_cast<const float4*>(ap + 4);
    const float4 b0 = *reinterpret_cast<const float4*>(bp);
    const float4 b1 = *reinterpret_cast<const float4*>(bp + 4);
    const size_t img = (size_t)n * S * C;
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * C + (c8 << 3);
        const H8 v = *reinterpret_cast<const H8*>(x + off);
        const float2 f0 = __half22float2(v.h[0]);
        const float2 f1 = __half22float2(v.h[1]);
        const float2 f2 = __half22float2(v.h[2]);
        const float2 f3 = __half22float2(v.h[3]);
        float4 o0, o1;
        o0.x = fmaf(fmaxf(f0.x, 0.0f), a0.x, b0.x);
        o0.y = fmaf(fmaxf(f0.y, 0.0f), a0.y, b0.y);
        o0.z = fmaf(fmaxf(f1.x, 0.0f), a0.z, b0.z);
        o0.w = fmaf(fmaxf(f1.y, 0.0f), a0.w, b0.w);
        o1.x = fmaf(fmaxf(f2.x, 0.0f), a1.x, b1.x);
        o1.y = fmaf(fmaxf(f2.y, 0.0f), a1.y, b1.y);
        o1.z = fmaf(fmaxf(f3.x, 0.0f), a1.z, b1.z);
        o1.w = fmaf(fmaxf(f3.y, 0.0f), a1.w, b1.w);
        *reinterpret_cast<float4*>(y + off) = o0;
        *reinterpret_cast<float4*>(y + off + 4) = o1;
    }
}

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

static inline void pick_splits(int64_t N, int64_t G, int64_t S,
                               int& splits, int& rows_per_split) {
    const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t NG = N * G;
    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;
    rows_per_split = (int)((S + splits - 1) / splits);
    splits = (int)((S + rows_per_split - 1) / rows_per_split);
}

// ------------------------- fused op: fp32 activations ----------------------
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);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    int splits, rows_per_split;
    pick_splits(N, G, S, splits, rows_per_split);
    const int64_t NG = N * G;

    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