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

# Fused CUDA implementation of:
#   GroupNorm -> Tanh -> HardSwish -> Residual Add -> LogSumExp(dim=1)
# Strategy:
#   Pass 1 (groupnorm_stats_kernel): one block per (batch, group) computes
#          mean / rstd with a vectorized (float4) block reduction.
#   Pass 2 (fused_act_lse_*_kernel): each thread owns one (or four) spatial
#          location(s), loops over all channels applying the fully fused
#          normalize + tanh + hardswish + residual chain, and accumulates an
#          online (single-pass, numerically stable) log-sum-exp. The whole
#          activation chain + reduction touches the conv output exactly once
#          more and writes only the tiny [N,1,H,W] result.

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

#define FULL_MASK 0xffffffffu

__device__ __forceinline__ float chain_op(float v, float mu, float rs, float w, float b) {
    // GroupNorm affine
    float normed = (v - mu) * rs * w + b;
    // Tanh
    float t = tanhf(normed);
    // HardSwish: t * relu6(t + 3) / 6
    float hs = t * fminf(fmaxf(t + 3.0f, 0.0f), 6.0f) * 0.16666666666666666f;
    // Residual add
    return v + hs;
}

__device__ __forceinline__ void lse_update(float r, float &m, float &s) {
    // Online log-sum-exp update: log(s * e^m + e^r)
    if (r > m) {
        s = s * expf(m - r) + 1.0f;
        m = r;
    } else {
        s += expf(r - m);
    }
}

__global__ void groupnorm_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ mean,
    float* __restrict__ rstd,
    int C, int HW, int G, float eps)
{
    const int n = blockIdx.x / G;
    const int g = blockIdx.x % G;
    const int cg = C / G;
    const int count = cg * HW;
    // Group channels are contiguous in NCHW layout
    const float* base = x + ((long)n * C + (long)g * cg) * (long)HW;

    float sum = 0.0f;
    float sumsq = 0.0f;

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

    __shared__ float s_sum[32];
    __shared__ float s_sq[32];

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum += __shfl_down_sync(FULL_MASK, sum, off);
        sumsq += __shfl_down_sync(FULL_MASK, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[warp] = sum;
        s_sq[warp] = sumsq;
    }
    __syncthreads();
    const int nwarp = (blockDim.x + 31) >> 5;
    if (warp == 0) {
        float tsum = (lane < nwarp) ? s_sum[lane] : 0.0f;
        float tsq = (lane < nwarp) ? s_sq[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            tsum += __shfl_down_sync(FULL_MASK, tsum, off);
            tsq += __shfl_down_sync(FULL_MASK, tsq, off);
        }
        if (lane == 0) {
            float mu = tsum / (float)count;
            float var = tsq / (float)count - mu * mu;
            var = fmaxf(var, 0.0f);
            mean[blockIdx.x] = mu;
            rstd[blockIdx.x] = rsqrtf(var + eps);
        }
    }
}

__global__ void fused_act_lse_vec4_kernel(
    const float4* __restrict__ x4,
    const float* __restrict__ mean,
    const float* __restrict__ rstd,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float4* __restrict__ out4,
    int C, int HW4, int G)
{
    const int n = blockIdx.y;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= HW4) return;

    const float4* xn = x4 + (long)n * C * HW4;
    const float* mn = mean + n * G;
    const float* rn = rstd + n * G;
    const int cg = C / G;

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

    for (int c = 0; c < C; ++c) {
        const int g = c / cg;
        const float mu = mn[g];
        const float rs = rn[g];
        const float w = weight[c];
        const float b = bias[c];
        const float4 v = xn[(long)c * HW4 + idx];
        lse_update(chain_op(v.x, mu, rs, w, b), m0, s0);
        lse_update(chain_op(v.y, mu, rs, w, b), m1, s1);
        lse_update(chain_op(v.z, mu, rs, w, b), m2, s2);
        lse_update(chain_op(v.w, mu, rs, w, b), m3, s3);
    }

    float4 o;
    o.x = m0 + logf(s0);
    o.y = m1 + logf(s1);
    o.z = m2 + logf(s2);
    o.w = m3 + logf(s3);
    out4[(long)n * HW4 + idx] = 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,
    int C, int HW, int G)
{
    const int n = blockIdx.y;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= HW) return;

    const float* xn = x + (long)n * C * HW;
    const float* mn = mean + n * G;
    const float* rn = rstd + n * G;
    const int cg = C / G;

    float m = -INFINITY;
    float s = 0.0f;
    for (int c = 0; c < C; ++c) {
        const int g = c / cg;
        const float v = xn[(long)c * HW + idx];
        lse_update(chain_op(v, mn[g], rn[g], weight[c], bias[c]), m, s);
    }
    out[(long)n * HW + idx] = m + logf(s);
}

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 must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(weight.is_contiguous() && bias.is_contiguous(),
                "affine params must be contiguous");

    const int64_t N = x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int HW = H * W;
    const int G = (int)groups;

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

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

    // Pass 1: per-(batch, group) mean / rstd
    groupnorm_stats_kernel<<<(int)(N * G), 256, 0, stream>>>(
        x.data_ptr<float>(), mean.data_ptr<float>(), rstd.data_ptr<float>(),
        C, HW, G, (float)eps);

    // Pass 2: fused normalize + tanh + hardswish + residual + logsumexp
    const int block = 256;
    if ((HW & 3) == 0) {
        const int HW4 = HW >> 2;
        dim3 grid((HW4 + block - 1) / block, (unsigned)N);
        fused_act_lse_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(x.data_ptr<float>()),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            weight.data_ptr<float>(), bias.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            C, HW4, G);
    } else {
        dim3 grid((HW + block - 1) / block, (unsigned)N);
        fused_act_lse_scalar_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(),
            mean.data_ptr<float>(), rstd.data_ptr<float>(),
            weight.data_ptr<float>(), bias.data_ptr<float>(),
            out.data_ptr<float>(),
            C, HW, G);
    }

    return out;
}
"""

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

_fused = load_inline(
    name="fused_gn_tanh_hswish_res_lse_ext",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_gn_act_lse"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - Convolution kept as-is (cuDNN).
      - GroupNorm + Tanh + HardSwish + Residual Add + LogSumExp fused into
        two custom CUDA kernels (stats pass + single-pass fused elementwise /
        online logsumexp 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.groups = groups
        self.eps = eps

    def forward(self, x):
        x_conv = self.conv(x).contiguous()
        return _fused.fused_gn_act_lse(
            x_conv,
            self.group_norm.weight,
            self.group_norm.bias,
            self.groups,
            self.eps,
        )
```