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

torch.backends.cudnn.benchmark = True

# Custom CUDA pipeline (channels_last):
#   1) nchw_to_nhwc_kernel       : smem-tiled transpose of the (small) input so cuDNN
#                                  runs its native NHWC tensor-core conv (~600us vs
#                                  ~1.7ms for the NCHW path).
#   2) stats_nhwc_kernel         : fully-coalesced per-(batch, group) sum / sumsq
#                                  accumulation (float4 loads + smem + atomics).
#   3) lse_nhwc_kernel           : fused GroupNorm affine + Tanh + HardSwish +
#                                  residual add + LogSumExp over channels.
#                                  REWRITE vs previous version: two-phase reduction
#                                  (warp max-reduce with cheap fmaxf shuffles, then
#                                  exp/sum) instead of online logsumexp merging at
#                                  every butterfly level. Cuts SFU __expf work ~3.5x
#                                  (56 -> 16 expf per lane-iter) and breaks the serial
#                                  dependency chains; also emits one float4 store.
#   (+ generic fallback kernels for uncommon channel configs)
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256
#define NEG_INF __int_as_float(0xff800000)
#define MAX_CPL 4

__device__ __forceinline__ float fast_tanh(float y) {
    // tanh(y) = sign(y) * (1 - e^{-2|y|}) / (1 + e^{-2|y|}); no overflow possible.
    float e = __expf(-2.0f * fabsf(y));
    float t = __fdividef(1.0f - e, 1.0f + e);
    return copysignf(t, y);
}

__device__ __forceinline__ float act_res(float v, float a, float b) {
    // groupnorm affine: y = v * a + b, then tanh, then hardswish, then residual add.
    float y = v * a + b;
    float t = fast_tanh(y);
    float hs = t * fminf(fmaxf(t + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
    return v + hs;
}

__device__ __forceinline__ void lse_update(float r, float &m, float &s) {
    if (r > m) {
        s = s * __expf(m - r) + 1.0f;
        m = r;
    } else {
        s += __expf(r - m);
    }
}

// ---------------------------------------------------------------------------
// Kernel 0: NCHW -> NHWC transpose, smem tiled. Each block handles 256 spatial
// positions x 8 channels of one image. Both global reads and writes coalesced.
// ---------------------------------------------------------------------------
__global__ void nchw_to_nhwc_kernel(const float *__restrict__ in,
                                    float *__restrict__ out,
                                    int C, int HW) {
    __shared__ float sm[8 * 260];
    const int hw0 = blockIdx.x * 256;
    const int c0 = blockIdx.y * 8;
    const int n = blockIdx.z;
    const int nc = min(8, C - c0);
    const int t = threadIdx.x;
    const bool hv = (hw0 + t) < HW;
    const float *in_n = in + (size_t)n * (size_t)C * (size_t)HW;
    float vals[8];
    #pragma unroll
    for (int c = 0; c < 8; ++c) {
        vals[c] = (c < nc && hv) ? in_n[(size_t)(c0 + c) * (size_t)HW + hw0 + t] : 0.0f;
    }
    #pragma unroll
    for (int c = 0; c < 8; ++c) {
        if (c < nc) sm[c * 260 + t] = vals[c];
    }
    __syncthreads();
    float *out_n = out + (size_t)n * (size_t)HW * (size_t)C + (size_t)hw0 * (size_t)C + c0;
    const int total = 256 * nc;
    for (int f = t; f < total; f += blockDim.x) {
        const int i = f / nc;
        const int c = f - i * nc;
        if (hw0 + i < HW) {
            out_n[(size_t)i * (size_t)C + c] = sm[c * 260 + i];
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel 1 (fast): per-(n, group) sum / sumsq of the NHWC conv output.
// Requires: C % 4 == 0, CperG % 4 == 0, (C/4) <= THREADS, THREADS % (C/4) == 0.
// ---------------------------------------------------------------------------
__global__ void stats_nhwc_kernel(const float4 *__restrict__ x4,
                                  float *__restrict__ acc,
                                  int HW, int C, int G, int CperG, int chunk) {
    const int n = blockIdx.y;
    const int hw0 = blockIdx.x * chunk;
    const int span = min(chunk, HW - hw0);
    if (span <= 0) return;
    const int f4perHw = C >> 2;
    const int f4idx = threadIdx.x % f4perHw;
    const int hwSub = threadIdx.x / f4perHw;
    const int hwStep = blockDim.x / f4perHw;
    const int g = (f4idx << 2) / CperG;
    const float4 *base = x4 + ((size_t)n * (size_t)HW + hw0) * (size_t)f4perHw;
    float sum = 0.0f, sq = 0.0f;
    for (int hw = hwSub; hw < span; hw += hwStep) {
        float4 v = base[(size_t)hw * (size_t)f4perHw + f4idx];
        sum += (v.x + v.y) + (v.z + v.w);
        sq  += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
    }
    __shared__ float sm[64 * 2];
    for (int i = threadIdx.x; i < 2 * G; i += blockDim.x) sm[i] = 0.0f;
    __syncthreads();
    atomicAdd(&sm[g * 2 + 0], sum);
    atomicAdd(&sm[g * 2 + 1], sq);
    __syncthreads();
    float *accg = acc + (size_t)n * (size_t)G * 2;
    for (int i = threadIdx.x; i < G; i += blockDim.x) {
        atomicAdd(&accg[i * 2 + 0], sm[i * 2 + 0]);
        atomicAdd(&accg[i * 2 + 1], sm[i * 2 + 1]);
    }
}

// ---------------------------------------------------------------------------
// Kernel 1 (generic fallback): one block per (n, group), scalar strided loop.
// ---------------------------------------------------------------------------
__global__ void stats_nhwc_generic_kernel(const float *__restrict__ x,
                                          float *__restrict__ acc,
                                          int HW, int C, int G, int CperG) {
    const int ng = blockIdx.x;
    const int n = ng / G;
    const int g = ng - n * G;
    const int c0 = g * CperG;
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    const int total = CperG * HW;
    float sum = 0.0f, sq = 0.0f;
    for (int i = threadIdx.x; i < total; i += blockDim.x) {
        const int hw = i / CperG;
        const int c = c0 + (i - hw * CperG);
        const float v = xb[(size_t)hw * (size_t)C + c];
        sum += v;
        sq += v * v;
    }
    __shared__ float ssum[THREADS / 32];
    __shared__ float sssq[THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, o);
        sq += __shfl_down_sync(0xffffffffu, sq, o);
    }
    if (lane == 0) { ssum[wid] = sum; sssq[wid] = sq; }
    __syncthreads();
    if (wid == 0) {
        const int nw = THREADS / 32;
        sum = (lane < nw) ? ssum[lane] : 0.0f;
        sq = (lane < nw) ? sssq[lane] : 0.0f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, o);
            sq += __shfl_down_sync(0xffffffffu, sq, o);
        }
        if (lane == 0) {
            acc[ng * 2 + 0] = sum;
            acc[ng * 2 + 1] = sq;
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel 2 (fast): fused GroupNorm + Tanh + HardSwish + residual + LogSumExp,
// NHWC, TWO-PHASE reduction. One warp handles 4 consecutive spatial positions;
// lane l owns channels (2l, 2l+1) via float2 loads. Requires: C even, C <= 64.
// ---------------------------------------------------------------------------
__global__ void lse_nhwc_kernel(const float *__restrict__ x,
                                const float *__restrict__ w,
                                const float *__restrict__ b,
                                const float *__restrict__ acc,
                                float *__restrict__ out,
                                int HW, int C, int G, int CperG,
                                float inv_cnt, float eps) {
    const int n = blockIdx.y;
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;
    const int c0 = lane * 2;
    const int c1 = c0 + 1;
    const bool active = c0 < C;   // C even => c0 < C implies c1 < C
    float a0 = 0.0f, b0 = 0.0f, a1 = 0.0f, b1 = 0.0f;
    if (active) {
        const float *accn = acc + (size_t)n * (size_t)G * 2;
        const int g0 = c0 / CperG;
        const float sum0 = accn[g0 * 2 + 0], sq0 = accn[g0 * 2 + 1];
        const float mean0 = sum0 * inv_cnt;
        const float var0 = fmaxf(sq0 * inv_cnt - mean0 * mean0, 0.0f);
        const float rstd0 = rsqrtf(var0 + eps);
        a0 = w[c0] * rstd0;
        b0 = b[c0] - mean0 * a0;
        const int g1 = c1 / CperG;
        if (g1 == g0) {
            a1 = w[c1] * rstd0;
            b1 = b[c1] - mean0 * a1;
        } else {
            const float sum1 = accn[g1 * 2 + 0], sq1 = accn[g1 * 2 + 1];
            const float mean1 = sum1 * inv_cnt;
            const float var1 = fmaxf(sq1 * inv_cnt - mean1 * mean1, 0.0f);
            const float rstd1 = rsqrtf(var1 + eps);
            a1 = w[c1] * rstd1;
            b1 = b[c1] - mean1 * a1;
        }
    }
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    float *outn = out + (size_t)n * (size_t)HW;
    const int quad0 = (blockIdx.x * nwarps + wid) * 4;
    const int step = gridDim.x * nwarps * 4;
    for (int hw = quad0; hw < HW; hw += step) {
        const bool full = (hw + 3) < HW;
        float2 v0, v1, v2, v3;
        if (active) {
            const float *p = xb + (size_t)hw * (size_t)C + c0;
            v0 = *reinterpret_cast<const float2 *>(p);
            if (full) {
                v1 = *reinterpret_cast<const float2 *>(p + C);
                v2 = *reinterpret_cast<const float2 *>(p + 2 * C);
                v3 = *reinterpret_cast<const float2 *>(p + 3 * C);
            } else {
                v1 = (hw + 1 < HW) ? *reinterpret_cast<const float2 *>(p + C)
                                   : make_float2(0.0f, 0.0f);
                v2 = (hw + 2 < HW) ? *reinterpret_cast<const float2 *>(p + 2 * C)
                                   : make_float2(0.0f, 0.0f);
                v3 = make_float2(0.0f, 0.0f);
            }
        } else {
            v0 = v1 = v2 = v3 = make_float2(0.0f, 0.0f);
        }
        float m0, m1, m2, m3;
        float r0, r1, r2, r3, r4, r5, r6, r7;
        if (active) {
            r0 = act_res(v0.x, a0, b0); r1 = act_res(v0.y, a1, b1);
            r2 = act_res(v1.x, a0, b0); r3 = act_res(v1.y, a1, b1);
            r4 = act_res(v2.x, a0, b0); r5 = act_res(v2.y, a1, b1);
            r6 = act_res(v3.x, a0, b0); r7 = act_res(v3.y, a1, b1);
            m0 = fmaxf(r0, r1);
            m1 = fmaxf(r2, r3);
            m2 = fmaxf(r4, r5);
            m3 = fmaxf(r6, r7);
        } else {
            r0 = r1 = r2 = r3 = r4 = r5 = r6 = r7 = 0.0f;
            m0 = m1 = m2 = m3 = NEG_INF;
        }
        // Phase 1: cheap max butterflies (no exp).
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, o));
            m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, o));
            m2 = fmaxf(m2, __shfl_xor_sync(0xffffffffu, m2, o));
            m3 = fmaxf(m3, __shfl_xor_sync(0xffffffffu, m3, o));
        }
        // Phase 2: exp against the true max, then sum butterflies.
        float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
        if (active) {
            s0 = __expf(r0 - m0) + __expf(r1 - m0);
            s1 = __expf(r2 - m1) + __expf(r3 - m1);
            s2 = __expf(r4 - m2) + __expf(r5 - m2);
            s3 = __expf(r6 - m3) + __expf(r7 - m3);
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            s0 += __shfl_xor_sync(0xffffffffu, s0, o);
            s1 += __shfl_xor_sync(0xffffffffu, s1, o);
            s2 += __shfl_xor_sync(0xffffffffu, s2, o);
            s3 += __shfl_xor_sync(0xffffffffu, s3, o);
        }
        // After xor butterflies every lane holds all 4 results; lane 0 does one
        // vectorized store (hw is a multiple of 4 => 16B aligned).
        if (lane == 0) {
            if (full) {
                float4 o4;
                o4.x = m0 + logf(s0);
                o4.y = m1 + logf(s1);
                o4.z = m2 + logf(s2);
                o4.w = m3 + logf(s3);
                *reinterpret_cast<float4 *>(outn + hw) = o4;
            } else {
                outn[hw] = m0 + logf(s0);
                if (hw + 1 < HW) outn[hw + 1] = m1 + logf(s1);
                if (hw + 2 < HW) outn[hw + 2] = m2 + logf(s2);
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel 2 (generic fallback): one warp per spatial position, C <= 32*MAX_CPL.
// ---------------------------------------------------------------------------
__global__ void lse_nhwc_generic_kernel(const float *__restrict__ x,
                                        const float *__restrict__ w,
                                        const float *__restrict__ b,
                                        const float *__restrict__ acc,
                                        float *__restrict__ out,
                                        int HW, int C, int G, int CperG,
                                        float inv_cnt, float eps, int chunk) {
    const int n = blockIdx.y;
    const int hw0 = blockIdx.x * chunk;
    const int hw1 = min(hw0 + chunk, HW);
    if (hw0 >= hw1) return;
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;
    const int cpl = (C + 31) >> 5;
    float A[MAX_CPL], B[MAX_CPL];
    #pragma unroll
    for (int j = 0; j < MAX_CPL; ++j) {
        const int c = lane + (j << 5);
        if (j < cpl && c < C) {
            const int ng = n * G + c / CperG;
            const float mean = acc[ng * 2 + 0] * inv_cnt;
            const float var = fmaxf(acc[ng * 2 + 1] * inv_cnt - mean * mean, 0.0f);
            const float a = w[c] * rsqrtf(var + eps);
            A[j] = a;
            B[j] = b[c] - mean * a;
        }
    }
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    for (int hw = hw0 + wid; hw < hw1; hw += nwarps) {
        const float *p = xb + (size_t)hw * (size_t)C;
        float m = NEG_INF, s = 0.0f;
        #pragma unroll
        for (int j = 0; j < MAX_CPL; ++j) {
            if (j < cpl) {
                const int c = lane + (j << 5);
                if (c < C) {
                    lse_update(act_res(p[c], A[j], B[j]), m, s);
                }
            }
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            const float m2 = __shfl_xor_sync(0xffffffffu, m, o);
            const float s2 = __shfl_xor_sync(0xffffffffu, s, o);
            const float M = fmaxf(m, m2);
            float ns = 0.0f;
            if (M != NEG_INF) ns = s * __expf(m - M) + s2 * __expf(m2 - M);
            m = M;
            s = ns;
        }
        if (lane == 0) out[(size_t)n * (size_t)HW + hw] = m + logf(s);
    }
}

torch::Tensor nchw_to_nhwc(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4 && x.scalar_type() == at::kFloat,
                "x must be CUDA float32 4D");
    auto xc = x.contiguous();
    const int N = xc.size(0), C = xc.size(1), H = xc.size(2), W = xc.size(3);
    const int HW = H * W;
    // Allocate DIRECTLY in channels_last format so the NHWC-ordered buffer needs
    // no further reinterpretation.
    auto out = at::empty({N, C, H, W},
                         xc.options().memory_format(at::MemoryFormat::ChannelsLast));
    TORCH_CHECK(out.is_contiguous(at::MemoryFormat::ChannelsLast),
                "channels_last allocation failed");
    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid((HW + 255) / 256, (C + 7) / 8, N);
    nchw_to_nhwc_kernel<<<grid, THREADS, 0, stream>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(), C, HW);
    return out;
}

torch::Tensor fused_tail_nhwc(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                              int64_t groups, double eps) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4 && x.scalar_type() == at::kFloat,
                "x must be a CUDA float32 4D tensor");
    TORCH_CHECK(x.is_contiguous(at::MemoryFormat::ChannelsLast),
                "x must be channels_last contiguous");
    const int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    const int HW = H * W;
    const int G = (int)groups;
    TORCH_CHECK(C % G == 0, "channels must be divisible by groups");
    TORCH_CHECK(G <= 64, "too many groups");
    const int CperG = C / G;
    const bool stats_ok = (C % 4 == 0) && (CperG % 4 == 0) && (C / 4) <= THREADS &&
                          (THREADS % (C / 4)) == 0;
    const bool lse_fast_ok = (C % 2 == 0) && (C <= 64);
    TORCH_CHECK(C <= 32 * MAX_CPL, "too many channels");

    auto opts = x.options();
    auto acc = at::zeros({(long)N * G * 2}, opts);
    auto out = at::empty({N, 1, H, W}, opts);
    auto stream = at::cuda::getCurrentCUDAStream();

    const float *xp = x.data_ptr<float>();
    const float *wp = w.data_ptr<float>();
    const float *bp = b.data_ptr<float>();
    float *accp = acc.data_ptr<float>();
    float *outp = out.data_ptr<float>();

    if (stats_ok) {
        const int schunk = 512;
        dim3 sgrid((HW + schunk - 1) / schunk, N);
        stats_nhwc_kernel<<<sgrid, THREADS, 0, stream>>>(
            reinterpret_cast<const float4 *>(xp), accp, HW, C, G, CperG, schunk);
    } else {
        stats_nhwc_generic_kernel<<<N * G, THREADS, 0, stream>>>(
            xp, accp, HW, C, G, CperG);
    }

    const float inv_cnt = 1.0f / (float)((long long)CperG * (long long)HW);
    if (lse_fast_ok) {
        // fast path: warp per 4 spatial positions, two-phase reduction
        const int nwarps = THREADS / 32;
        const int iters = 8;
        dim3 lgrid((HW + 4 * nwarps * iters - 1) / (4 * nwarps * iters), N);
        lse_nhwc_kernel<<<lgrid, THREADS, 0, stream>>>(
            xp, wp, bp, accp, outp, HW, C, G, CperG, inv_cnt, (float)eps);
    } else {
        const int lchunk = 512;
        dim3 lgrid((HW + lchunk - 1) / lchunk, N);
        lse_nhwc_generic_kernel<<<lgrid, THREADS, 0, stream>>>(
            xp, wp, bp, accp, outp, HW, C, G, CperG, inv_cnt, (float)eps, lchunk);
    }
    return out;
}
"""

_cpp_src = (
    "torch::Tensor nchw_to_nhwc(torch::Tensor x);\n"
    "torch::Tensor fused_tail_nhwc(torch::Tensor x, torch::Tensor w, torch::Tensor b, int64_t groups, double eps);"
)

_fused = load_inline(
    name="fused_tail_nhwc_v5",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["nchw_to_nhwc", "fused_tail_nhwc"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: custom NCHW->NHWC input transpose, channels_last cuDNN
    convolution (native NHWC tensor-core path), then two custom CUDA kernels that
    fuse GroupNorm statistics + GroupNorm/Tanh/HardSwish/Residual/LogSumExp.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, eps=1e-5):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.group_norm = nn.GroupNorm(groups, out_channels, eps=eps)
        self.groups = groups
        self.eps = float(eps)
        C = out_channels
        g = groups
        cperg = C // g if g > 0 else 0
        # Static check that the custom NHWC tail supports this configuration.
        self._tail_ok = (
            g > 0
            and C % g == 0
            and C <= 128
            and g <= 64
        )

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and self._tail_ok:
            # Custom fast transpose to channels_last (much faster than torch's
            # .contiguous(memory_format=channels_last)).
            x_cl = _fused.nchw_to_nhwc(x)
            # NOTE: do NOT cache this - weights may be updated between calls.
            # It is only a tiny (out_ch, in_ch, k, k) copy.
            w_cl = self.conv.weight.contiguous(memory_format=torch.channels_last)
            x_conv = F.conv2d(
                x_cl,
                w_cl,
                self.conv.bias,
                self.conv.stride,
                self.conv.padding,
                self.conv.dilation,
                self.conv.groups,
            )
            if x_conv.is_contiguous(memory_format=torch.channels_last):
                return _fused.fused_tail_nhwc(
                    x_conv,
                    self.group_norm.weight,
                    self.group_norm.bias,
                    self.groups,
                    self.eps,
                )
            x_conv = x_conv.contiguous()
        else:
            x_conv = self.conv(x)
        # Reference fallback (CPU tensors or unsupported configurations).
        x_norm = self.group_norm(x_conv)
        x_tanh = torch.tanh(x_norm)
        x_hard_swish = F.hardswish(x_tanh)
        x_res = x_conv + x_hard_swish
        return torch.logsumexp(x_res, dim=1, keepdim=True)
```