```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 tail for the NHWC (channels_last) convolution output:
#   Kernel 1 (stats_nhwc_kernel): fully-coalesced per-(batch, group) sum / sumsq
#             accumulation using float4 loads and atomics.
#   Kernel 2 (lse_nhwc_kernel):   fused GroupNorm affine + Tanh + HardSwish +
#             residual add + LogSumExp over channels. One warp per spatial
#             location; channel weights are cached in registers.
# Running the whole tail in channels_last avoids the expensive NHWC<->NCHW
# transposes that cuDNN inserts around its NHWC tensor-core convolution.
_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 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.
__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: fused GroupNorm + Tanh + HardSwish + residual + LogSumExp (NHWC).
// One warp per spatial location; generic over C (C <= 32 * MAX_CPL).
__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, 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 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;
    TORCH_CHECK(C % 4 == 0 && CperG % 4 == 0 && (C / 4) <= THREADS && (THREADS % (C / 4)) == 0,
                "unsupported channel configuration");
    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>();

    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 int lchunk = 512;
    const float inv_cnt = 1.0f / (float)((long long)CperG * (long long)HW);
    dim3 lgrid((HW + lchunk - 1) / lchunk, N);
    lse_nhwc_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 fused_tail_nhwc(torch::Tensor x, torch::Tensor w, torch::Tensor b, int64_t groups, double eps);"

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


class ModelNew(nn.Module):
    """
    Optimized model: channels_last cuDNN convolution (native NHWC tensor-core
    path with no layout transposes), followed by 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
            fast_ok = (
                C % g == 0
                and C % 4 == 0
                and cperg % 4 == 0
                and (C // 4) <= 256
                and C <= 128
                and g <= 64
            )
            if fast_ok:
                # channels_last input -> cuDNN runs native NHWC kernels and the
                # output stays channels_last (no 512MB transposes).
                x_cl = x.contiguous(memory_format=torch.channels_last)
                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 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)
```