```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 kernels:
#  1) group_stats_kernel: computes per-(batch, group) mean / rstd of the conv output.
#  2) fused_act_lse_*_kernel: fuses GroupNorm affine + Tanh + HardSwish + residual add
#     + LogSumExp over channels into a single pass over the data (vectorized float4 path
#     plus a scalar fallback for odd spatial sizes).
_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)

__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) {
    // online logsumexp update
    if (r > m) {
        s = s * __expf(m - r) + 1.0f;
        m = r;
    } else {
        s += __expf(r - m);
    }
}

// Kernel 1: per (n, group) mean and rstd of the conv output.
__global__ void group_stats_kernel(const float *__restrict__ x,
                                   float *__restrict__ mean,
                                   float *__restrict__ rstd,
                                   int C, int HW, int G, int CperG, float eps) {
    const int ng = blockIdx.x;
    const int n = ng / G;
    const int g = ng - n * G;
    const int c0 = g * CperG;
    const float *base = x + (size_t)n * (size_t)C * (size_t)HW;
    float sum = 0.0f, sumsq = 0.0f;
    if ((HW & 3) == 0) {
        const int HW4 = HW >> 2;
        for (int c = c0; c < c0 + CperG; ++c) {
            const float4 *p = reinterpret_cast<const float4 *>(base + (size_t)c * (size_t)HW);
            for (int i = threadIdx.x; i < HW4; i += blockDim.x) {
                float4 v = p[i];
                sum += (v.x + v.y) + (v.z + v.w);
                sumsq += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
            }
        }
    } else {
        for (int c = c0; c < c0 + CperG; ++c) {
            const float *p = base + (size_t)c * (size_t)HW;
            for (int i = threadIdx.x; i < HW; i += blockDim.x) {
                float v = p[i];
                sum += v;
                sumsq += 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);
        sumsq += __shfl_down_sync(0xffffffffu, sumsq, o);
    }
    if (lane == 0) {
        ssum[wid] = sum;
        sssq[wid] = sumsq;
    }
    __syncthreads();
    if (wid == 0) {
        const int nw = THREADS / 32;
        sum = (lane < nw) ? ssum[lane] : 0.0f;
        sumsq = (lane < nw) ? sssq[lane] : 0.0f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, o);
            sumsq += __shfl_down_sync(0xffffffffu, sumsq, o);
        }
        if (lane == 0) {
            const float cnt = (float)((long long)CperG * (long long)HW);
            const float m = sum / cnt;
            float var = sumsq / cnt - m * m;
            var = fmaxf(var, 0.0f);
            mean[ng] = m;
            rstd[ng] = rsqrtf(var + eps);
        }
    }
}

// Kernel 2 (float4 vectorized): fused GroupNorm + Tanh + HardSwish + residual + LogSumExp.
__global__ void fused_act_lse_vec4_kernel(const float *__restrict__ x,
                                          const float *__restrict__ w,
                                          const float *__restrict__ b,
                                          const float *__restrict__ mean,
                                          const float *__restrict__ rstd,
                                          float *__restrict__ out,
                                          int C, int HW, int G, int CperG) {
    const int n = blockIdx.y;
    const int HW4 = HW >> 2;
    extern __shared__ float sm[];
    float *sa = sm;        // per-channel scale  = gamma * rstd(group)
    float *sb = sm + C;    // per-channel offset = beta - mean(group) * scale
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        const int ng = n * G + c / CperG;
        const float a = w[c] * rstd[ng];
        sa[c] = a;
        sb[c] = b[c] - mean[ng] * a;
    }
    __syncthreads();
    const int hw4 = blockIdx.x * blockDim.x + threadIdx.x;
    if (hw4 >= HW4) return;
    const float4 *xp = reinterpret_cast<const float4 *>(x + (size_t)n * (size_t)C * (size_t)HW) + hw4;
    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;
    const size_t stride = (size_t)HW4;
    for (int c = 0; c < C; ++c) {
        const float4 v = xp[c * stride];
        const float a = sa[c], bb = sb[c];
        lse_update(act_res(v.x, a, bb), m0, s0);
        lse_update(act_res(v.y, a, bb), m1, s1);
        lse_update(act_res(v.z, a, bb), m2, s2);
        lse_update(act_res(v.w, a, bb), m3, s3);
    }
    float4 o;
    o.x = m0 + logf(s0);
    o.y = m1 + logf(s1);
    o.z = m2 + logf(s2);
    o.w = m3 + logf(s3);
    reinterpret_cast<float4 *>(out + (size_t)n * (size_t)HW)[hw4] = o;
}

// Kernel 2 (scalar fallback for spatial size not divisible by 4).
__global__ void fused_act_lse_kernel(const float *__restrict__ x,
                                     const float *__restrict__ w,
                                     const float *__restrict__ b,
                                     const float *__restrict__ mean,
                                     const float *__restrict__ rstd,
                                     float *__restrict__ out,
                                     int C, int HW, int G, int CperG) {
    const int n = blockIdx.y;
    extern __shared__ float sm[];
    float *sa = sm;
    float *sb = sm + C;
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        const int ng = n * G + c / CperG;
        const float a = w[c] * rstd[ng];
        sa[c] = a;
        sb[c] = b[c] - mean[ng] * a;
    }
    __syncthreads();
    const int hw = blockIdx.x * blockDim.x + threadIdx.x;
    if (hw >= HW) return;
    const float *xp = x + (size_t)n * (size_t)C * (size_t)HW + hw;
    float m = NEG_INF, s = 0.0f;
    for (int c = 0; c < C; ++c) {
        const float v = xp[(size_t)c * (size_t)HW];
        lse_update(act_res(v, sa[c], sb[c]), m, s);
    }
    out[(size_t)n * (size_t)HW + hw] = m + logf(s);
}

torch::Tensor fused_gn_act_lse(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");
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();
    const int N = xc.size(0), C = xc.size(1), H = xc.size(2), W = xc.size(3);
    const int HW = H * W;
    const int G = (int)groups;
    TORCH_CHECK(C % G == 0, "channels must be divisible by groups");
    const int CperG = C / G;
    auto opts = xc.options();
    auto mean = at::empty({N * G}, opts);
    auto rstd = at::empty({N * G}, opts);
    auto out = at::empty({N, 1, H, W}, opts);
    auto stream = at::cuda::getCurrentCUDAStream();
    group_stats_kernel<<<N * G, THREADS, 0, stream>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        C, HW, G, CperG, (float)eps);
    const int smem = 2 * C * (int)sizeof(float);
    if ((HW & 3) == 0) {
        dim3 grid((HW / 4 + THREADS - 1) / THREADS, N);
        fused_act_lse_vec4_kernel<<<grid, THREADS, smem, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            out.data_ptr<float>(), C, HW, G, CperG);
    } else {
        dim3 grid((HW + THREADS - 1) / THREADS, N);
        fused_act_lse_kernel<<<grid, THREADS, smem, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            out.data_ptr<float>(), C, HW, G, CperG);
    }
    return out;
}
"""

_cpp_src = "torch::Tensor fused_gn_act_lse(torch::Tensor x, torch::Tensor w, torch::Tensor b, int64_t groups, double eps);"

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


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN convolution, and replaces
    GroupNorm + Tanh + HardSwish + Residual Add + LogSumExp with a pair of
    fused custom CUDA kernels (stats kernel + fused elementwise/reduction kernel).
    """
    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):
        # Convolution (cuDNN)
        x_conv = self.conv(x)
        if x_conv.is_cuda and x_conv.dtype == torch.float32:
            # Fused GroupNorm + Tanh + HardSwish + Residual + LogSumExp
            return _fused.fused_gn_act_lse(
                x_conv,
                self.group_norm.weight,
                self.group_norm.bias,
                self.groups,
                self.eps,
            )
        # Fallback (e.g. CPU tensors): original PyTorch operator sequence
        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)
```