import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, applies Group Normalization, Tanh, HardSwish, 
    Residual Addition, and LogSumExp.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, eps=1e-5):
        super(Model, 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()

    def forward(self, x):
        # Convolution
        x_conv = self.conv(x)
        # Group Normalization
        x_norm = self.group_norm(x_conv)
        # Tanh
        x_tanh = self.tanh(x_norm)
        # HardSwish
        x_hard_swish = self.hard_swish(x_tanh)
        # Residual Addition
        x_res = x_conv + x_hard_swish
        # LogSumExp
        x_logsumexp = torch.logsumexp(x_res, dim=1, keepdim=True)
        return x_logsumexp

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128
kernel_size = 3
groups = 16

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, groups]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor fused_gn_lse_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int64_t groups, double eps);
"""

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

__global__ void gn_stats_kernel(const float* __restrict__ x,
                                float* __restrict__ mean,
                                float* __restrict__ invstd,
                                int N, int C, int H, int W, int G, float eps) {
    int bg = blockIdx.x;
    int tid = threadIdx.x;
    int cpg = C / G;
    int HW = H * W;
    int elems = cpg * HW;
    int n = bg / G;
    int g = bg - n * G;
    int c0 = g * cpg;

    float sum = 0.0f;
    float ssq = 0.0f;
    int base = (n * C + c0) * HW;
    for (int i = tid; i < elems; i += blockDim.x) {
        int cr = i / HW;
        int s = i - cr * HW;
        float v = x[base + cr * HW + s];
        sum += v;
        ssq = fmaf(v, v, ssq);
    }

    extern __shared__ float smem[];
    float* ssum = smem;
    float* sssq = smem + blockDim.x;
    ssum[tid] = sum;
    sssq[tid] = ssq;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ssum[tid] += ssum[tid + stride];
            sssq[tid] += sssq[tid + stride];
        }
        __syncthreads();
    }
    if (tid == 0) {
        float m = ssum[0] / (float)elems;
        float v = sssq[0] / (float)elems - m * m;
        if (v < 0.0f) v = 0.0f;
        mean[bg] = m;
        invstd[bg] = rsqrtf(v + eps);
    }
}

__device__ __forceinline__ float act_res(float x, float m, float inv, float gamma, float beta) {
    float z = (x - m) * inv;
    z = fmaf(z, gamma, beta);
    float t = tanhf(z);
    // Since tanh output is in [-1,1], HardSwish(t) = t * (t + 3) / 6 exactly (no clamp branch needed).
    float hs = t * (t + 3.0f) * 0.16666666666666666f;
    return x + hs;
}

__global__ void fused_lse_kernel(const float* __restrict__ x,
                                 const float* __restrict__ gamma,
                                 const float* __restrict__ beta,
                                 const float* __restrict__ mean,
                                 const float* __restrict__ invstd,
                                 float* __restrict__ y,
                                 int N, int C, int H, int W, int G) {
    int HW = H * W;
    int total = N * HW;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int n = idx / HW;
    int s = idx - n * HW;
    int cpg = C / G;
    int nC = n * C;

    float maxv = -CUDART_INF_F;
#pragma unroll
    for (int c = 0; c < 64; ++c) {
        if (c >= C) break;
        int g = c / cpg;
        float xv = x[(nC + c) * HW + s];
        float v = act_res(xv, mean[n * G + g], invstd[n * G + g], gamma[c], beta[c]);
        maxv = fmaxf(maxv, v);
    }

    float sum = 0.0f;
#pragma unroll
    for (int c = 0; c < 64; ++c) {
        if (c >= C) break;
        int g = c / cpg;
        float xv = x[(nC + c) * HW + s];
        float v = act_res(xv, mean[n * G + g], invstd[n * G + g], gamma[c], beta[c]);
        sum += expf(v - maxv);
    }
    y[idx] = maxv + logf(sum);
}

torch::Tensor fused_gn_lse_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int64_t groups, double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(gamma.is_cuda() && beta.is_cuda(), "gamma/beta must be CUDA");
    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();

    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int H = (int)xc.size(2);
    int W = (int)xc.size(3);
    int G = (int)groups;

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int stat_threads = 256;
    size_t shmem = 2 * stat_threads * sizeof(float);
    gn_stats_kernel<<<N * G, stat_threads, shmem, stream>>>(
        xc.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        N, C, H, W, G, (float)eps);

    const int threads = 128;
    int total = N * H * W;
    int blocks = (total + threads - 1) / threads;
    fused_lse_kernel<<<blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(), gc.data_ptr<float>(), bc.data_ptr<float>(),
        mean.data_ptr<float>(), invstd.data_ptr<float>(), y.data_ptr<float>(),
        N, C, H, W, G);
    return y;
}
"""

_ext = load_inline(
    name="fused_gn_lse_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_gn_lse_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    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()

    def forward(self, x):
        x_conv = self.conv(x)
        return _ext.fused_gn_lse_forward(x_conv, self.group_norm.weight, self.group_norm.bias,
                                         self.group_norm.num_groups, self.group_norm.eps)