```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: one-pass per-(batch, group) mean / rstd for GroupNorm.
#   2) fused_gn_act_lse_kernel: fuses GroupNorm affine + Tanh + HardSwish +
#      residual add + channel-wise LogSumExp (online/single-pass logsumexp)
#      into a single kernel, reading the conv output only once.
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define BLOCK_THREADS 256
#define NUM_WARPS (BLOCK_THREADS / 32)

__global__ void group_stats_kernel(const float* __restrict__ x,
                                   float* __restrict__ mean,
                                   float* __restrict__ rstd,
                                   int C, int G, int cpg,
                                   long spatial, float eps) {
    const int g = blockIdx.x % G;
    const int n = blockIdx.x / G;
    const int c0 = g * cpg;
    const float* xb = x + (long)n * C * spatial;

    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;

    const long nvec = spatial >> 2;
    for (int c = 0; c < cpg; ++c) {
        const float* p = xb + (long)(c0 + c) * spatial;
        const float4* p4 = reinterpret_cast<const float4*>(p);
        for (long i = threadIdx.x; i < nvec; i += blockDim.x) {
            float4 v = p4[i];
            s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
            q0 += v.x * v.x; q1 += v.y * v.y; q2 += v.z * v.z; q3 += v.w * v.w;
        }
        for (long i = (nvec << 2) + threadIdx.x; i < spatial; i += blockDim.x) {
            float v = p[i];
            s0 += v; q0 += v * v;
        }
    }

    float sum = (s0 + s1) + (s2 + s3);
    float sumsq = (q0 + q1) + (q2 + q3);

    __shared__ float sh_sum[NUM_WARPS];
    __shared__ float sh_sq[NUM_WARPS];
    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) { sh_sum[wid] = sum; sh_sq[wid] = sumsq; }
    __syncthreads();

    if (wid == 0) {
        sum   = (lane < NUM_WARPS) ? sh_sum[lane] : 0.f;
        sumsq = (lane < NUM_WARPS) ? sh_sq[lane] : 0.f;
        #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 inv = 1.f / (float)((long)cpg * spatial);
            const float m = sum * inv;
            float var = sumsq * inv - m * m;
            var = fmaxf(var, 0.f);
            mean[blockIdx.x] = m;
            rstd[blockIdx.x] = rsqrtf(var + eps);
        }
    }
}

__global__ void fused_gn_act_lse_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 G, int cpg, long spatial) {
    const int n = blockIdx.y;
    const long s = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= spatial) return;

    const float* xn = x + (long)n * C * spatial;
    const float* mn = mean + n * G;
    const float* rn = rstd + n * G;

    float m = -INFINITY;
    float l = 0.f;

    int g = 0;
    int cnt = 0;
    float mg = mn[0];
    float rg = rn[0];

    for (int c = 0; c < C; ++c) {
        const float v = xn[(long)c * spatial + s];
        // GroupNorm affine
        float y = (v - mg) * rg;
        y = y * weight[c] + bias[c];
        // Tanh
        const float t = tanhf(y);
        // HardSwish: t * relu6(t + 3) / 6
        const float r6 = fminf(fmaxf(t + 3.f, 0.f), 6.f);
        // Residual add
        const float r = v + t * r6 * (1.f / 6.f);
        // Online logsumexp across channels
        if (r > m) {
            l = l * expf(m - r) + 1.f;
            m = r;
        } else {
            l += expf(r - m);
        }
        if (++cnt == cpg) {
            cnt = 0;
            ++g;
            if (g < G) { mg = mn[g]; rg = rn[g]; }
        }
    }
    out[(long)n * spatial + s] = m + logf(l);
}

torch::Tensor fused_gn_act_lse(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                               int64_t groups, double eps) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32, "x must be CUDA float32");
    auto xc = x.contiguous();
    const int N = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const long spatial = (long)xc.size(2) * (long)xc.size(3);
    const int G = (int)groups;
    const int cpg = 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, xc.size(2), xc.size(3)}, opts);

    dim3 stat_grid(N * G);
    group_stats_kernel<<<stat_grid, BLOCK_THREADS>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        C, G, cpg, spatial, (float)eps);

    dim3 grid((unsigned int)((spatial + BLOCK_THREADS - 1) / BLOCK_THREADS), N);
    fused_gn_act_lse_kernel<<<grid, BLOCK_THREADS>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),
        C, G, cpg, spatial);

    return out;
}
"""

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

_ext = load_inline(
    name="fused_gn_act_lse_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_gn_act_lse"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: convolution stays in cuDNN, while GroupNorm + Tanh +
    HardSwish + residual add + LogSumExp are fused into custom CUDA kernels
    (one statistics pass + one fully fused elementwise/reduction pass).
    """
    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()
        self.groups = groups
        self.eps = eps

    def forward(self, x):
        x_conv = self.conv(x)
        if x_conv.is_cuda and x_conv.dtype == torch.float32:
            return _ext.fused_gn_act_lse(
                x_conv.contiguous(),
                self.group_norm.weight,
                self.group_norm.bias,
                self.groups,
                float(self.eps),
            )
        # Fallback reference path (non-CUDA tensors)
        x_norm = self.group_norm(x_conv)
        x_tanh = self.tanh(x_norm)
        x_hard_swish = self.hard_swish(x_tanh)
        x_res = x_conv + x_hard_swish
        return torch.logsumexp(x_res, dim=1, keepdim=True)
```