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

# Let cuDNN pick the fastest convolution algorithm for the fixed input shape.
torch.backends.cudnn.benchmark = True

# Custom CUDA kernels:
#   1) group_stats_kernel: computes per-(batch, group) mean / rstd for GroupNorm.
#   2) fused_act_lse_*_kernel: fuses GroupNorm application + Tanh + HardSwish +
#      residual add + online (numerically stable) LogSumExp over the channel dim
#      into a single pass, never materializing the intermediate [B, C, H, W]
#      tensors (x_norm, x_tanh, x_hard_swish, x_res).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr int kBlockSize = 256;

__device__ __forceinline__ float gn_tanh_hardswish_res(float v, float mn, float rs,
                                                       float w, float b) {
    // GroupNorm affine
    float gn = (v - mn) * rs * w + b;
    // Tanh
    float t = tanhf(gn);
    // HardSwish: x * relu6(x + 3) / 6
    float hs = t * (fminf(fmaxf(t + 3.0f, 0.0f), 6.0f) * 0.1666666666666667f);
    // Residual add with the pre-norm conv output
    return v + hs;
}

__device__ __forceinline__ void lse_update(float v, float& m, float& s) {
    // Online log-sum-exp update (numerically stable).
    if (v > m) {
        s = s * expf(m - v) + 1.0f;
        m = v;
    } else {
        s += expf(v - m);
    }
}

__global__ void group_stats_kernel(const float* __restrict__ x,
                                   float* __restrict__ mean,
                                   float* __restrict__ rstd,
                                   long long group_size,
                                   float eps) {
    const int g = blockIdx.x;
    const int b = blockIdx.y;
    const int G = gridDim.x;
    const long long base = ((long long)b * G + g) * group_size;
    const float* p = x + base;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;

    if ((group_size & 3LL) == 0LL) {
        const long long n4 = group_size >> 2;
        const float4* p4 = reinterpret_cast<const float4*>(p);
        for (long long i = threadIdx.x; i < n4; i += blockDim.x) {
            const float4 v = p4[i];
            s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
            q0 = fmaf(v.x, v.x, q0);
            q1 = fmaf(v.y, v.y, q1);
            q2 = fmaf(v.z, v.z, q2);
            q3 = fmaf(v.w, v.w, q3);
        }
    } else {
        for (long long i = threadIdx.x; i < group_size; i += blockDim.x) {
            const float v = p[i];
            s0 += v;
            q0 = fmaf(v, v, q0);
        }
    }

    // Promote partial sums to double for an accurate block reduction.
    double s = (double)((s0 + s1) + (s2 + s3));
    double sq = (double)((q0 + q1) + (q2 + q3));

    __shared__ double sh_s[32];
    __shared__ double sh_q[32];
    const unsigned int mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        s += __shfl_down_sync(mask, s, off);
        sq += __shfl_down_sync(mask, sq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh_s[wid] = s;
        sh_q[wid] = sq;
    }
    __syncthreads();
    if (wid == 0) {
        const int nw = (blockDim.x + 31) >> 5;
        s = (lane < nw) ? sh_s[lane] : 0.0;
        sq = (lane < nw) ? sh_q[lane] : 0.0;
        for (int off = 16; off > 0; off >>= 1) {
            s += __shfl_down_sync(mask, s, off);
            sq += __shfl_down_sync(mask, sq, off);
        }
        if (lane == 0) {
            const double n = (double)group_size;
            const double m = s / n;
            double var = sq / n - m * m;
            if (var < 0.0) var = 0.0;
            const int idx = b * G + g;
            mean[idx] = (float)m;
            rstd[idx] = (float)(1.0 / sqrt(var + (double)eps));
        }
    }
}

__global__ void fused_act_lse_vec4_kernel(const float* __restrict__ x,
                                          const float* __restrict__ mean,
                                          const float* __restrict__ rstd,
                                          const float* __restrict__ weight,
                                          const float* __restrict__ bias,
                                          float* __restrict__ out,
                                          int C, int HW, int G, int cpg) {
    const int HW4 = HW >> 2;
    const int b = blockIdx.y;
    const int r = blockIdx.x * blockDim.x + threadIdx.x;
    if (r >= HW4) return;

    const float4* xp = reinterpret_cast<const float4*>(x) + ((long long)b * C) * HW4 + r;
    const float* mp = mean + b * G;
    const float* rp = rstd + b * G;

    float m0 = -3.402823466e38f, m1 = m0, m2 = m0, m3 = m0;
    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;

    for (int gg = 0; gg < G; ++gg) {
        const float mn = mp[gg];
        const float rs = rp[gg];
        const int c0 = gg * cpg;
        for (int cc = 0; cc < cpg; ++cc) {
            const int c = c0 + cc;
            const float w = weight[c];
            const float bb = bias[c];
            const float4 v = xp[(long long)c * HW4];
            lse_update(gn_tanh_hardswish_res(v.x, mn, rs, w, bb), m0, s0);
            lse_update(gn_tanh_hardswish_res(v.y, mn, rs, w, bb), m1, s1);
            lse_update(gn_tanh_hardswish_res(v.z, mn, rs, w, bb), m2, s2);
            lse_update(gn_tanh_hardswish_res(v.w, mn, rs, w, bb), m3, s3);
        }
    }

    float4 o;
    o.x = logf(s0) + m0;
    o.y = logf(s1) + m1;
    o.z = logf(s2) + m2;
    o.w = logf(s3) + m3;
    reinterpret_cast<float4*>(out)[(long long)b * HW4 + r] = o;
}

__global__ void fused_act_lse_scalar_kernel(const float* __restrict__ x,
                                            const float* __restrict__ mean,
                                            const float* __restrict__ rstd,
                                            const float* __restrict__ weight,
                                            const float* __restrict__ bias,
                                            float* __restrict__ out,
                                            long long P, int C, int HW, int G, int cpg) {
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= P) return;
    const int b = (int)(idx / HW);
    const int hw = (int)(idx - (long long)b * HW);
    const float* xp = x + ((long long)b * C) * HW + hw;
    const float* mp = mean + b * G;
    const float* rp = rstd + b * G;

    float m = -3.402823466e38f;
    float s = 0.f;
    for (int gg = 0; gg < G; ++gg) {
        const float mn = mp[gg];
        const float rs = rp[gg];
        const int c0 = gg * cpg;
        for (int cc = 0; cc < cpg; ++cc) {
            const int c = c0 + cc;
            const float v = xp[(long long)c * HW];
            lse_update(gn_tanh_hardswish_res(v, mn, rs, weight[c], bias[c]), m, s);
        }
    }
    out[idx] = logf(s) + m;
}

}  // namespace

torch::Tensor fused_gn_tanh_hardswish_res_lse(torch::Tensor x,
                                              torch::Tensor weight,
                                              torch::Tensor bias,
                                              int64_t groups,
                                              double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto w = weight.contiguous();
    auto bb = bias.contiguous();

    const int B = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int H = (int)xc.size(2);
    const int W = (int)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 cpg = C / G;

    auto mean = at::empty({B, G}, xc.options());
    auto rstd = at::empty({B, G}, xc.options());
    auto out = at::empty({B, 1, H, W}, xc.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const long long group_size = (long long)cpg * (long long)HW;
    dim3 grid_stats(G, B);
    group_stats_kernel<<<grid_stats, kBlockSize, 0, stream>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        group_size, (float)eps);

    if ((HW & 3) == 0) {
        const int HW4 = HW >> 2;
        dim3 grid_apply((HW4 + kBlockSize - 1) / kBlockSize, B);
        fused_act_lse_vec4_kernel<<<grid_apply, kBlockSize, 0, stream>>>(
            xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
            w.data_ptr<float>(), bb.data_ptr<float>(), out.data_ptr<float>(),
            C, HW, G, cpg);
    } else {
        const long long P = (long long)B * HW;
        const int blocks = (int)((P + kBlockSize - 1) / kBlockSize);
        fused_act_lse_scalar_kernel<<<blocks, kBlockSize, 0, stream>>>(
            xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
            w.data_ptr<float>(), bb.data_ptr<float>(), out.data_ptr<float>(),
            P, C, HW, G, cpg);
    }
    return out;
}
"""

cpp_source = "torch::Tensor fused_gn_tanh_hardswish_res_lse(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t groups, double eps);"

# Compile the inline CUDA extension.
fused_ext = load_inline(
    name="fused_gn_act_lse_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_gn_tanh_hardswish_res_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 two fused
    custom CUDA kernels (a group-statistics kernel and a single fused
    normalize/activation/residual/online-logsumexp 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.tanh = nn.Tanh()
        self.hard_swish = nn.Hardswish()

    def forward(self, x):
        # Convolution (cuDNN)
        x_conv = self.conv(x)
        # Fused GroupNorm + Tanh + HardSwish + Residual Add + LogSumExp
        return fused_ext.fused_gn_tanh_hardswish_res_lse(
            x_conv,
            self.group_norm.weight,
            self.group_norm.bias,
            self.group_norm.num_groups,
            self.group_norm.eps,
        )
```