```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
from torch.utils.cpp_extension import load_inline

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

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <cstdint>

#define BLOCK_SIZE 256

// ---------------------------------------------------------------------------
// Online log-sum-exp update
// ---------------------------------------------------------------------------
__device__ __forceinline__ void lse_update(float r, float& m, float& s) {
    float nm = fmaxf(m, r);
    s = s * expf(m - nm) + expf(r - nm);
    m = nm;
}

// ---------------------------------------------------------------------------
// Kernel 1: per-(sample, group) mean / rstd for GroupNorm.
// One thread block per (n, g); double accumulators for accuracy.
// ---------------------------------------------------------------------------
__global__ void group_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ mean,
    float* __restrict__ rstd,
    const int C, const int HW, const int G, const int CperG, const float eps)
{
    const int gid = blockIdx.x;
    const int n = gid / G;
    const int g = gid % G;
    const float* base = x + ((long)n * C + (long)g * CperG) * HW;
    const int numel = CperG * HW;

    double sum = 0.0;
    double sumsq = 0.0;

    int start_scalar = 0;
    if ((reinterpret_cast<uintptr_t>(base) & 15) == 0) {
        const int numel4 = numel >> 2;
        const float4* base4 = reinterpret_cast<const float4*>(base);
        for (int i = threadIdx.x; i < numel4; i += BLOCK_SIZE) {
            float4 v = base4[i];
            sum   += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
            sumsq += (double)v.x * v.x + (double)v.y * v.y +
                     (double)v.z * v.z + (double)v.w * v.w;
        }
        start_scalar = numel4 << 2;
    }
    for (int j = start_scalar + threadIdx.x; j < numel; j += BLOCK_SIZE) {
        double v = (double)base[j];
        sum += v;
        sumsq += v * v;
    }

    __shared__ double sh_sum[32];
    __shared__ double sh_sumsq[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(0xffffffffu, sum, off);
        sumsq += __shfl_down_sync(0xffffffffu, sumsq, off);
    }
    if (lane == 0) {
        sh_sum[warp] = sum;
        sh_sumsq[warp] = sumsq;
    }
    __syncthreads();

    if (warp == 0) {
        const int nw = BLOCK_SIZE >> 5;
        sum   = (lane < nw) ? sh_sum[lane] : 0.0;
        sumsq = (lane < nw) ? sh_sumsq[lane] : 0.0;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            sum   += __shfl_down_sync(0xffffffffu, sum, off);
            sumsq += __shfl_down_sync(0xffffffffu, sumsq, off);
        }
        if (lane == 0) {
            double inv = 1.0 / (double)numel;
            double m = sum * inv;
            double var = sumsq * inv - m * m;
            if (var < 0.0) var = 0.0;
            mean[gid] = (float)m;
            rstd[gid] = (float)(1.0 / sqrt(var + (double)eps));
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel 2 (scalar fallback): GroupNorm apply + Tanh + HardSwish + residual
// add + per-position LogSumExp over channels. One thread per (n, h, w).
// ---------------------------------------------------------------------------
__global__ void fused_postop_scalar(
    const float* __restrict__ x,
    const float* __restrict__ mean,
    const float* __restrict__ rstd,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int N, const int C, const int HW, const int G, const int CperG)
{
    const long idx = (long)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    if (idx >= (long)N * HW) return;
    const int n = (int)(idx / HW);
    const int s = (int)(idx % HW);

    const float* base = x + (long)n * C * HW + s;
    const float* mn = mean + (long)n * G;
    const float* rs = rstd + (long)n * G;

    float m = -CUDART_INF_F;
    float acc = 0.f;

    for (int g = 0; g < G; ++g) {
        const float cm = mn[g];
        const float cr = rs[g];
        const float* gb = base + (long)g * CperG * HW;
        for (int cc = 0; cc < CperG; ++cc) {
            const int c = g * CperG + cc;
            const float v = gb[(long)cc * HW];
            const float y = (v - cm) * cr * weight[c] + bias[c];
            const float t = tanhf(y);
            const float hs = t * fminf(fmaxf(t + 3.f, 0.f), 6.f) * (1.f / 6.f);
            lse_update(v + hs, m, acc);
        }
    }
    out[idx] = m + logf(acc);
}

// ---------------------------------------------------------------------------
// Kernel 2 (vectorized): same as above, but each thread processes 4 adjacent
// spatial locations via float4. Requires HW % 4 == 0.
// ---------------------------------------------------------------------------
__global__ void fused_postop_vec4(
    const float4* __restrict__ x,
    const float* __restrict__ mean,
    const float* __restrict__ rstd,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float4* __restrict__ out,
    const int N, const int C, const int HW4, const int G, const int CperG)
{
    const long idx = (long)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    if (idx >= (long)N * HW4) return;
    const int n = (int)(idx / HW4);
    const int s = (int)(idx % HW4);

    const float4* base = x + (long)n * C * HW4 + s;
    const float* mn = mean + (long)n * G;
    const float* rs = rstd + (long)n * G;

    float m[4]   = {-CUDART_INF_F, -CUDART_INF_F, -CUDART_INF_F, -CUDART_INF_F};
    float acc[4] = {0.f, 0.f, 0.f, 0.f};

    for (int g = 0; g < G; ++g) {
        const float cm = mn[g];
        const float cr = rs[g];
        const float4* gb = base + (long)g * CperG * HW4;
        for (int cc = 0; cc < CperG; ++cc) {
            const int c = g * CperG + cc;
            const float w = weight[c];
            const float b = bias[c];
            const float4 v = gb[(long)cc * HW4];
            #pragma unroll
            for (int k = 0; k < 4; ++k) {
                const float vc = (&v.x)[k];
                const float y = (vc - cm) * cr * w + b;
                const float t = tanhf(y);
                const float hs = t * fminf(fmaxf(t + 3.f, 0.f), 6.f) * (1.f / 6.f);
                lse_update(vc + hs, m[k], acc[k]);
            }
        }
    }

    float4 o;
    o.x = m[0] + logf(acc[0]);
    o.y = m[1] + logf(acc[1]);
    o.z = m[2] + logf(acc[2]);
    o.w = m[3] + logf(acc[3]);
    out[idx] = o;
}

// ---------------------------------------------------------------------------
// Host wrapper
// ---------------------------------------------------------------------------
torch::Tensor fused_gn_tanh_hswish_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 b = bias.contiguous();

    const int N = (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;
    const int CperG = C / G;

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

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

    group_stats_kernel<<<N * G, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        C, HW, G, CperG, (float)eps);

    if (HW % 4 == 0) {
        const int HW4 = HW / 4;
        const long total = (long)N * HW4;
        const long blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
        fused_postop_vec4<<<(unsigned int)blocks, BLOCK_SIZE, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            w.data_ptr<float>(), b.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            N, C, HW4, G, CperG);
    } else {
        const long total = (long)N * HW;
        const long blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
        fused_postop_scalar<<<(unsigned int)blocks, BLOCK_SIZE, 0, stream>>>(
            xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
            w.data_ptr<float>(), b.data_ptr<float>(),
            out.data_ptr<float>(), N, C, HW, G, CperG);
    }

    return out;
}
"""

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


class ModelNew(nn.Module):
    """
    Optimized model: the convolution is left to cuDNN, while GroupNorm
    (statistics + affine), Tanh, HardSwish, the residual addition and the
    channel-wise LogSumExp are all fused into custom CUDA kernels, so the
    convolution output only needs to be read twice and no intermediate
    activation tensors are ever materialized.
    """
    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):
        x_conv = self.conv(x)
        return _fused_ops.fused_gn_tanh_hswish_res_lse(
            x_conv,
            self.group_norm.weight,
            self.group_norm.bias,
            self.groups,
            self.eps,
        )
```