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

torch.backends.cudnn.benchmark = True

# Custom CUDA kernels:
#   Pass 1: per-(batch, group) GroupNorm statistics (mean / rstd), vectorized.
#   Pass 2: fully fused GroupNorm-apply + Tanh + HardSwish + residual add +
#           online LogSumExp reduction over the channel dimension.
# This keeps only the cuDNN convolution from PyTorch and replaces all the
# remaining pointwise / reduction ops with two fused kernels, so the large
# intermediate activation tensor is only written once (by the conv) and read
# twice, instead of being read/written by six separate ops.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

// ------------------------- device helpers -------------------------

__device__ __forceinline__ void gn_act_lse_update(float v, float mu, float rs,
                                                  float w, float b,
                                                  float& m, float& s) {
    // GroupNorm affine
    float y = (v - mu) * rs;
    y = y * w + b;
    // tanh(y) = 1 - 2 / (exp(2y) + 1)  (absolute error < ~1e-6 vs tanhf)
    float e2 = __expf(2.f * y);
    float t = 1.f - __fdividef(2.f, e2 + 1.f);
    // HardSwish: t * relu6(t + 3) / 6, then residual add with conv output v
    float r = v + t * fminf(fmaxf(t + 3.f, 0.f), 6.f) * (1.f / 6.f);
    // online log-sum-exp update
    if (r > m) {
        s = s * __expf(m - r) + 1.f;
        m = r;
    } else {
        s += __expf(r - m);
    }
}

// ------------------------- pass 1: stats -------------------------

__global__ void gn_stats_kernel(const float* __restrict__ x,
                                float* __restrict__ mean,
                                float* __restrict__ rstd,
                                int C, int HW, float eps) {
    const int g = blockIdx.x;
    const int n = blockIdx.y;
    const int G = gridDim.x;
    const int Cg = C / G;
    const long count = (long)Cg * HW;
    const float* p = x + ((long)n * C + (long)g * Cg) * HW;

    float sum = 0.f;
    float sumsq = 0.f;

    if (((reinterpret_cast<uintptr_t>(p) & 15) == 0) && ((count & 3) == 0)) {
        const float4* p4 = reinterpret_cast<const float4*>(p);
        const long count4 = count >> 2;
        for (long i = threadIdx.x; i < count4; i += blockDim.x) {
            float4 v = p4[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 (long i = threadIdx.x; i < count; i += blockDim.x) {
            float v = p[i];
            sum += v;
            sumsq += v * v;
        }
    }

    __shared__ float sh_sum[32];
    __shared__ float sh_sq[32];
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, off);
        sumsq += __shfl_down_sync(0xffffffffu, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) { sh_sum[wid] = sum; sh_sq[wid] = sumsq; }
    __syncthreads();
    if (wid == 0) {
        const int nw = (blockDim.x + 31) >> 5;
        sum = (lane < nw) ? sh_sum[lane] : 0.f;
        sumsq = (lane < nw) ? sh_sq[lane] : 0.f;
        #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) {
            const float m = sum / (float)count;
            const float var = sumsq / (float)count - m * m;
            mean[n * G + g] = m;
            rstd[n * G + g] = rsqrtf(var + eps);
        }
    }
}

// ------------------------- pass 2: fused apply + reduce -------------------------

template <int CG>
__global__ void fused_vec4_kernel(const float* __restrict__ x,
                                  const float* __restrict__ mean,
                                  const float* __restrict__ rstd,
                                  const float* __restrict__ w,
                                  const float* __restrict__ b,
                                  float* __restrict__ out,
                                  int C, int G, int HW, int Cg) {
    extern __shared__ float sh[];
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        sh[c] = w[c];
        sh[C + c] = b[c];
    }
    __syncthreads();

    const int n = blockIdx.y;
    const int hw = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if (hw >= HW) return;

    const float* xn = x + (size_t)n * C * HW + hw;
    const float* mnp = mean + (size_t)n * G;
    const float* rnp = rstd + (size_t)n * G;
    const int cg = (CG > 0) ? CG : Cg;

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

    for (int g = 0; g < G; ++g) {
        const float mu = mnp[g];
        const float rs = rnp[g];
        #pragma unroll
        for (int ci = 0; ci < cg; ++ci) {
            const int c = g * cg + ci;
            const float4 v = *reinterpret_cast<const float4*>(xn + (size_t)c * HW);
            const float wc = sh[c];
            const float bc = sh[C + c];
            gn_act_lse_update(v.x, mu, rs, wc, bc, m0, s0);
            gn_act_lse_update(v.y, mu, rs, wc, bc, m1, s1);
            gn_act_lse_update(v.z, mu, rs, wc, bc, m2, s2);
            gn_act_lse_update(v.w, mu, rs, wc, bc, 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 * HW + hw) = o;
}

template <int CG>
__global__ void fused_scalar_kernel(const float* __restrict__ x,
                                    const float* __restrict__ mean,
                                    const float* __restrict__ rstd,
                                    const float* __restrict__ w,
                                    const float* __restrict__ b,
                                    float* __restrict__ out,
                                    int C, int G, int HW, int Cg) {
    extern __shared__ float sh[];
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        sh[c] = w[c];
        sh[C + c] = b[c];
    }
    __syncthreads();

    const int n = blockIdx.y;
    const int hw = blockIdx.x * blockDim.x + threadIdx.x;
    if (hw >= HW) return;

    const float* xn = x + (size_t)n * C * HW + hw;
    const float* mnp = mean + (size_t)n * G;
    const float* rnp = rstd + (size_t)n * G;
    const int cg = (CG > 0) ? CG : Cg;

    float m = -INFINITY, s = 0.f;
    for (int g = 0; g < G; ++g) {
        const float mu = mnp[g];
        const float rs = rnp[g];
        #pragma unroll
        for (int ci = 0; ci < cg; ++ci) {
            const int c = g * cg + ci;
            const float v = xn[(size_t)c * HW];
            gn_act_lse_update(v, mu, rs, sh[c], sh[C + c], m, s);
        }
    }
    out[(size_t)n * HW + hw] = m + __logf(s);
}

// ------------------------- host wrapper -------------------------

torch::Tensor fused_forward(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.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = 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;
    TORCH_CHECK(G > 0 && C % G == 0, "channels must be divisible by groups");
    const int Cg = 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();

    // pass 1: GroupNorm statistics, one block per (batch, group)
    {
        dim3 grid(G, N);
        gn_stats_kernel<<<grid, 256, 0, stream>>>(
            xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
            C, HW, (float)eps);
    }

    // pass 2: fused normalize + tanh + hardswish + residual + logsumexp
    const int block = 256;
    const size_t shmem = (size_t)2 * C * sizeof(float);
    if ((HW & 3) == 0) {
        const int hw4 = HW >> 2;
        dim3 grid((hw4 + block - 1) / block, N);
        if (Cg == 4) {
            fused_vec4_kernel<4><<<grid, block, shmem, stream>>>(
                xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
                wc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(),
                C, G, HW, Cg);
        } else {
            fused_vec4_kernel<0><<<grid, block, shmem, stream>>>(
                xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
                wc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(),
                C, G, HW, Cg);
        }
    } else {
        dim3 grid((HW + block - 1) / block, N);
        if (Cg == 4) {
            fused_scalar_kernel<4><<<grid, block, shmem, stream>>>(
                xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
                wc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(),
                C, G, HW, Cg);
        } else {
            fused_scalar_kernel<0><<<grid, block, shmem, stream>>>(
                xc.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
                wc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(),
                C, G, HW, Cg);
        }
    }

    return out;
}
"""

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

fused_ops = load_inline(
    name="fused_gn_tanh_hswish_lse",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN convolution followed by two fused custom CUDA
    kernels that replace GroupNorm, Tanh, HardSwish, the residual addition,
    and the LogSumExp reduction.
    """
    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 = eps

    def forward(self, x):
        x_conv = self.conv(x)
        return fused_ops.fused_forward(
            x_conv,
            self.group_norm.weight,
            self.group_norm.bias,
            self.groups,
            self.eps,
        )
```