```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     : fast smem-tiled transpose of the (small) input so
#                                cuDNN can run its native NHWC tensor-core conv
#                                (measured ~600us vs ~1.7ms for the NCHW path).
#                                FIX vs previous attempt: the output tensor is now
#                                allocated directly in channels_last memory format
#                                (previously the NHWC buffer was wrapped in an NCHW
#                                tensor and then permute().contiguous() silently
#                                transposed the already-transposed data a 2nd time,
#                                scrambling the conv input -> wrong results).
#   2) stats_nhwc_kernel       : fully-coalesced per-(batch, group) sum / sumsq
#                                accumulation using float4 loads + atomics.
#   3) lse_nhwc_kernel         : fused GroupNorm affine + Tanh + HardSwish + residual
#                                add + online LogSumExp over channels. One warp per
#                                4 spatial positions, float2 loads, 4 independent
#                                accumulation chains for memory-level parallelism.
#   (+ 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.
// smem row stride 260 floats makes the store-phase bank-conflict free.
// ---------------------------------------------------------------------------
__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: accumulate per-(n, group) sum and sum-of-squares of the NHWC conv
// output. Requires: C % 4 == 0, CperG % 4 == 0, (C/4) <= THREADS,
// THREADS % (C/4) == 0, G <= 64.
// ---------------------------------------------------------------------------
__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;   // constant per thread
    const int hwSub = threadIdx.x / f4perHw;
    const int hwStep = blockDim.x / f4perHw;
    const int g = (f4idx << 2) / CperG;        // group of this thread's channels
    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 2 (fast path): fused GroupNorm + Tanh + HardSwish + residual +
// LogSumExp, NHWC. 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;
    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);
        a0 = w[c0] * rsqrtf(var0 + eps);
        b0 = b[c0] - mean0 * a0;
        const int g1 = c1 / CperG;
        if (g1 == g0) {
            a1 = w[c1] * rsqrtf(var0 + eps);
            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);
            a1 = w[c1] * rsqrtf(var1 + eps);
            b1 = b[c1] - mean1 * a1;
        }
    }
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    const int quad0 = (blockIdx.x * nwarps + wid) * 4;
    const int step = gridDim.x * nwarps * 4;
    for (int hw = quad0; hw < HW; hw += step) {
        float2 v0, v1, v2, v3;
        if (active)                 v0 = *reinterpret_cast<const float2 *>(xb + (size_t)hw * C + c0);
        else                        v0 = make_float2(0.0f, 0.0f);
        if (active && hw + 1 < HW)  v1 = *reinterpret_cast<const float2 *>(xb + (size_t)(hw + 1) * C + c0);
        else                        v1 = make_float2(0.0f, 0.0f);
        if (active && hw + 2 < HW)  v2 = *reinterpret_cast<const float2 *>(xb + (size_t)(hw + 2) * C + c0);
        else                        v2 = make_float2(0.0f, 0.0f);
        if (active && hw + 3 < HW)  v3 = *reinterpret_cast<const float2 *>(xb + (size_t)(hw + 3) * C + c0);
        else                        v3 = make_float2(0.0f, 0.0f);
        float m0 = NEG_INF, m1 = NEG_INF, m2 = NEG_INF, m3 = NEG_INF;
        float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
        if (active) {
            lse_update(act_res(v0.x, a0, b0), m0, s0);
            lse_update(act_res(v0.y, a1, b1), m0, s0);
            lse_update(act_res(v1.x, a0, b0), m1, s1);
            lse_update(act_res(v1.y, a1, b1), m1, s1);
            lse_update(act_res(v2.x, a0, b0), m2, s2);
            lse_update(act_res(v2.y, a1, b1), m2, s2);
            lse_update(act_res(v3.x, a0, b0), m3, s3);
            lse_update(act_res(v3.y, a1, b1), m3, s3);
        }
        // xor-butterfly over all 32 lanes: after this, every lane holds the full
        // reduction for each of the 4 position chains; lanes 0-3 write them out.
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            const float t0 = __shfl_xor_sync(0xffffffffu, m0, o);
            const float u0 = __shfl_xor_sync(0xffffffffu, s0, o);
            const float M0 = fmaxf(m0, t0);
            s0 = (M0 != NEG_INF) ? s0 * __expf(m0 - M0) + u0 * __expf(t0 - M0) : 0.0f;
            m0 = M0;
            const float t1 = __shfl_xor_sync(0xffffffffu, m1, o);
            const float u1 = __shfl_xor_sync(0xffffffffu, s1, o);
            const float M1 = fmaxf(m1, t1);
            s1 = (M1 != NEG_INF) ? s1 * __expf(m1 - M1) + u1 * __expf(t1 - M1) : 0.0f;
            m1 = M1;
            const float t2 = __shfl_xor_sync(0xffffffffu, m2, o);
            const float u2 = __shfl_xor_sync(0xffffffffu, s2, o);
            const float M2 = fmaxf(m2, t2);
            s2 = (M2 != NEG_INF) ? s2 * __expf(m2 - M2) + u2 * __expf(t2 - M2) : 0.0f;
            m2 = M2;
            const float t3 = __shfl_xor_sync(0xffffffffu, m3, o);
            const float u3 = __shfl_xor_sync(0xffffffffu, s3, o);
            const float M3 = fmaxf(m3, t3);
            s3 = (M3 != NEG_INF) ? s3 * __expf(m3 - M3) + u3 * __expf(t3 - M3) : 0.0f;
            m3 = M3;
        }
        float *outn = out + (size_t)n * (size_t)HW;
        if (lane == 0) outn[hw] = m0 + logf(s0);
        if (lane == 1 && hw + 1 < HW) outn[hw + 1] = m1 + logf(s1);
        if (lane == 2 && hw + 2 < HW) outn[hw + 2] = m2 + logf(s2);
        if (lane == 3 && hw + 3 < HW) outn[hw + 3] = m3 + logf(s3);
    }
}

// ---------------------------------------------------------------------------
// 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: the buffer is NHWC-ordered and
    // the tensor already carries channels_last strides, so the kernel's NHWC
    // output needs no further reinterpretation (this was the previous bug).
    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;
    TORCH_CHECK(stats_ok, "unsupported channel configuration for stats kernel");

    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>();

    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);

    const float inv_cnt = 1.0f / (float)((long long)CperG * (long long)HW);
    if (C % 2 == 0 && C <= 64) {
        // fast path: warp per 4 spatial positions
        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 {
        TORCH_CHECK(C <= 32 * MAX_CPL, "too many channels");
        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_v4",
    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, no layout conversion kernels),
    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)

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            C = self.conv.out_channels
            g = self.groups
            cperg = C // g if g > 0 else 0
            fast_ok = (
                g > 0
                and C % g == 0
                and C % 4 == 0
                and cperg % 4 == 0
                and (C // 4) <= 256
                and (256 % (C // 4)) == 0
                and C <= 128
                and g <= 64
            )
            if fast_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)
        else:
            x_conv = self.conv(x)
        # Reference fallback (CPU tensors or unsupported channel configs).
        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)
```