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

# Custom fused CUDA kernels:
#   1) group_stats_*_kernel: computes per-(batch, group) mean/rstd for GroupNorm
#      in a single pass over the conv output (vectorized float4 loads when possible).
#   2) fused_norm_act_res_lse_kernel: fuses GroupNorm affine + Tanh + HardSwish +
#      residual add + online LogSumExp over the channel dimension into one kernel.
# This eliminates all intermediate global-memory traffic of the original op chain.
_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math.h>

#define THREADS 256

__device__ __forceinline__ void block_reduce_sum2(float &sum, float &sumsq) {
    __shared__ float s_sum[THREADS];
    __shared__ float s_sq[THREADS];
    s_sum[threadIdx.x] = sum;
    s_sq[threadIdx.x] = sumsq;
    __syncthreads();
    for (int s = THREADS >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + s];
            s_sq[threadIdx.x]  += s_sq[threadIdx.x + s];
        }
        __syncthreads();
    }
    sum = s_sum[0];
    sumsq = s_sq[0];
}

__global__ void group_stats_scalar_kernel(const float* __restrict__ x,
                                          float* __restrict__ mean,
                                          float* __restrict__ rstd,
                                          int group_size, float eps) {
    const long gid = blockIdx.x;
    const float* ptr = x + gid * (long)group_size;
    float sum = 0.f, sumsq = 0.f;
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        float v = ptr[i];
        sum += v;
        sumsq += v * v;
    }
    block_reduce_sum2(sum, sumsq);
    if (threadIdx.x == 0) {
        float m = sum / (float)group_size;
        float var = sumsq / (float)group_size - m * m;
        var = fmaxf(var, 0.f);
        mean[gid] = m;
        rstd[gid] = 1.0f / sqrtf(var + eps);
    }
}

__global__ void group_stats_vec4_kernel(const float* __restrict__ x,
                                        float* __restrict__ mean,
                                        float* __restrict__ rstd,
                                        int group_size, float eps) {
    const long gid = blockIdx.x;
    const float4* ptr = reinterpret_cast<const float4*>(x + gid * (long)group_size);
    const int n4 = group_size >> 2;
    float sum = 0.f, sumsq = 0.f;
    for (int i = threadIdx.x; i < n4; i += blockDim.x) {
        float4 v = ptr[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;
    }
    block_reduce_sum2(sum, sumsq);
    if (threadIdx.x == 0) {
        float m = sum / (float)group_size;
        float var = sumsq / (float)group_size - m * m;
        var = fmaxf(var, 0.f);
        mean[gid] = m;
        rstd[gid] = 1.0f / sqrtf(var + eps);
    }
}

__global__ void fused_norm_act_res_lse_kernel(const float* __restrict__ x_conv,
                                              const float* __restrict__ mean,
                                              const float* __restrict__ rstd,
                                              const float* __restrict__ weight,
                                              const float* __restrict__ bias,
                                              float* __restrict__ out,
                                              int C, int HW, int groups, int cpg) {
    const int hw = blockIdx.x * blockDim.x + threadIdx.x;
    if (hw >= HW) return;
    const int n = blockIdx.y;

    const float* base = x_conv + (long)n * C * HW + hw;
    const float* mptr = mean + n * groups;
    const float* rptr = rstd + n * groups;

    float m = -INFINITY;
    float s = 0.f;
    int g = 0;
    int next_boundary = cpg;
    float cur_mean = mptr[0];
    float cur_rstd = rptr[0];

    #pragma unroll 4
    for (int c = 0; c < C; ++c) {
        if (c == next_boundary) {
            ++g;
            next_boundary += cpg;
            cur_mean = mptr[g];
            cur_rstd = rptr[g];
        }
        const float xv = base[(long)c * HW];
        // GroupNorm affine
        const float norm = (xv - cur_mean) * cur_rstd * weight[c] + bias[c];
        // Tanh
        const float t = tanhf(norm);
        // HardSwish: t * relu6(t + 3) / 6
        const float hs = t * fminf(fmaxf(t + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
        // Residual add
        const float r = xv + hs;
        // Online LogSumExp over channels
        const float m_new = fmaxf(m, r);
        s = s * expf(m - m_new) + expf(r - m_new);
        m = m_new;
    }
    out[(long)n * HW + hw] = m + logf(s);
}

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

    auto xc = x_conv.contiguous();
    auto w = weight.contiguous();
    auto b = 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;
    const int cpg = C / G;
    const int group_size = cpg * HW;

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

    const float* xptr = xc.data_ptr<float>();
    float* mptr = mean.data_ptr<float>();
    float* rptr = rstd.data_ptr<float>();

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

    const int nstat_blocks = B * G;
    if ((group_size & 3) == 0) {
        group_stats_vec4_kernel<<<nstat_blocks, THREADS, 0, stream>>>(
            xptr, mptr, rptr, group_size, (float)eps);
    } else {
        group_stats_scalar_kernel<<<nstat_blocks, THREADS, 0, stream>>>(
            xptr, mptr, rptr, group_size, (float)eps);
    }

    dim3 grid((HW + THREADS - 1) / THREADS, B);
    fused_norm_act_res_lse_kernel<<<grid, THREADS, 0, stream>>>(
        xptr, mptr, rptr, w.data_ptr<float>(), b.data_ptr<float>(),
        out.data_ptr<float>(), C, HW, G, cpg);

    return out;
}
"""

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

_fused_ext = load_inline(
    name="fused_gn_act_res_lse_ext",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_gn_act_res_lse"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: Conv2d (cuDNN) + a fully fused custom CUDA kernel that performs
    GroupNorm + Tanh + HardSwish + Residual Add + LogSumExp in two kernel launches
    with a single extra pass over the conv output for the group statistics.
    """
    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)
        if x_conv.is_cuda and x_conv.dtype == torch.float32:
            return _fused_ext.fused_gn_act_res_lse(
                x_conv,
                self.group_norm.weight,
                self.group_norm.bias,
                self.groups,
                self.eps,
            )
        # Fallback reference path (non-CUDA / non-float32 inputs)
        x_norm = self.group_norm(x_conv)
        x_tanh = torch.tanh(x_norm)
        x_hard_swish = F.hardswish(x_tanh)
        x_res = x_conv + x_hard_swish
        return torch.logsumexp(x_res, dim=1, keepdim=True)
```