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_source = r'''
#include <torch/extension.h>
at::Tensor fused_forward(at::Tensor x, at::Tensor weight, at::Tensor bias, at::Tensor gamma, at::Tensor beta, int64_t groups, double eps);
'''

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

#define CI 8
#define CO 64
#define HH 128
#define WW 128
#define KK 3
#define OH_ 126
#define OW_ 126
#define SP_ 15876
#define GRP 16
#define CPG 4

__global__ void conv3_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b,
                             float* __restrict__ conv, int N) {
    long long total = (long long)N * CO * SP_;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int sp = idx % SP_;
    int ox = sp % OW_;
    int oy = sp / OW_;
    int co = (idx / SP_) & 63;
    int n = idx / ((long long)CO * SP_);
    float acc = b[co];
    const float* wp = w + co * (CI*KK*KK);
    const float* xp0 = x + ((long long)n * CI) * HH * WW + oy * WW + ox;
    #pragma unroll
    for (int ci = 0; ci < CI; ++ci) {
        const float* xp = xp0 + ci * HH * WW;
        const float* wc = wp + ci * 9;
        acc = fmaf(xp[0], wc[0], acc);
        acc = fmaf(xp[1], wc[1], acc);
        acc = fmaf(xp[2], wc[2], acc);
        acc = fmaf(xp[WW], wc[3], acc);
        acc = fmaf(xp[WW+1], wc[4], acc);
        acc = fmaf(xp[WW+2], wc[5], acc);
        acc = fmaf(xp[2*WW], wc[6], acc);
        acc = fmaf(xp[2*WW+1], wc[7], acc);
        acc = fmaf(xp[2*WW+2], wc[8], acc);
    }
    conv[idx] = acc;
}

__global__ void reduce_kernel(const float* __restrict__ conv, float* __restrict__ mean, float* __restrict__ invstd,
                              int N, float eps) {
    int bg = blockIdx.x;
    int n = bg >> 4;
    int g = bg & 15;
    constexpr int count = CPG * SP_;
    float s = 0.0f, ss = 0.0f;
    int c0 = g * CPG;
    for (int i = threadIdx.x; i < count; i += blockDim.x) {
        int lc = i / SP_;
        int sp = i - lc * SP_;
        float v = conv[((long long)n * CO + (c0 + lc)) * SP_ + sp];
        s += v;
        ss = fmaf(v, v, ss);
    }
    __shared__ float shs[256], shss[256];
    shs[threadIdx.x] = s; shss[threadIdx.x] = ss;
    __syncthreads();
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) { shs[threadIdx.x] += shs[threadIdx.x+stride]; shss[threadIdx.x] += shss[threadIdx.x+stride]; }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        float m = shs[0] * (1.0f / (float)count);
        float var = shss[0] * (1.0f / (float)count) - m*m;
        mean[bg] = m;
        invstd[bg] = rsqrtf(var + eps);
    }
}

__global__ void final_kernel(const float* __restrict__ conv, const float* __restrict__ mean, const float* __restrict__ invstd,
                             const float* __restrict__ gamma, const float* __restrict__ beta, float* __restrict__ out,
                             int N) {
    long long total = (long long)N * SP_;
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int sp = idx % SP_;
    int n = idx / SP_;
    float mval = -INFINITY;
    float vals[64];
    #pragma unroll
    for (int c = 0; c < CO; ++c) {
        float v = conv[((long long)n * CO + c) * SP_ + sp];
        int g = c >> 2;
        float z = (v - mean[(n << 4) + g]) * invstd[(n << 4) + g];
        z = z * gamma[c] + beta[c];
        float t = tanhf(z);
        float r = v + t * (t + 3.0f) * 0.16666666666666666f;
        vals[c] = r;
        mval = fmaxf(mval, r);
    }
    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < CO; ++c) sum += expf(vals[c] - mval);
    out[idx] = mval + logf(sum);
}

at::Tensor fused_forward(at::Tensor x, at::Tensor weight, at::Tensor bias, at::Tensor gamma, at::Tensor beta, int64_t groups, double eps) {
    x = x.contiguous(); weight = weight.contiguous(); bias = bias.contiguous(); gamma = gamma.contiguous(); beta = beta.contiguous();
    int N = x.size(0);
    auto opts = x.options();
    auto conv = at::empty({N, CO, OH_, OW_}, opts);
    auto mean = at::empty({N, GRP}, opts);
    auto invstd = at::empty({N, GRP}, opts);
    auto out = at::empty({N, 1, OH_, OW_}, opts);
    int threads = 256;
    long long conv_total = (long long)N * CO * SP_;
    conv3_kernel<<<(conv_total + threads - 1)/threads, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), conv.data_ptr<float>(), N);
    reduce_kernel<<<N * GRP, 256>>>(conv.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), N, (float)eps);
    long long out_total = (long long)N * SP_;
    final_kernel<<<(out_total + threads - 1)/threads, threads>>>(conv.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), N);
    return out;
}
'''

_ext = load_inline(name='fused_conv_gn_lse_ext_final', cpp_sources=[cpp_source], cuda_sources=[cuda_source], functions=['fused_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], 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.groups = groups
        self.eps = eps
    def forward(self, x):
        return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.group_norm.weight, self.group_norm.bias, self.groups, self.eps)