import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a convolution, applies activation, and then applies Batch Normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels, eps=eps, momentum=momentum)

    def forward(self, x):
        x = self.conv(x)
        x = torch.multiply(torch.tanh(torch.nn.functional.softplus(x)), x)
        x = self.bn(x)
        return x

batch_size = 64
in_channels = 64
out_channels = 128
height, width = 128, 128
kernel_size = 3

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

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

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_mish_bn(torch::Tensor x, torch::Tensor running_mean, torch::Tensor running_var,
                            torch::Tensor gamma, torch::Tensor beta,
                            bool training, double momentum, double eps);
'''

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

static __device__ __forceinline__ float mish_act(float x) {
    float sp = (x > 20.0f) ? x : log1pf(expf(x));
    return x * tanhf(sp);
}

__global__ void mish_partial_kernel(const float* __restrict__ x, float* __restrict__ y,
                                    double* __restrict__ psum, double* __restrict__ psq,
                                    int N, int C, int H, int W, int parts) {
    int c = blockIdx.x, p = blockIdx.y, tid = threadIdx.x;
    int HW = H * W, M = N * HW;
    double sum = 0.0, ssq = 0.0;
    for (int m = p * blockDim.x + tid; m < M; m += parts * blockDim.x) {
        int n = m / HW;
        int hw = m - n * HW;
        int idx = (n * C + c) * HW + hw;
        float v = mish_act(x[idx]);
        y[idx] = v;
        double dv = (double)v;
        sum += dv;
        ssq += dv * dv;
    }
    extern __shared__ double smem[];
    double* s1 = smem; double* s2 = smem + blockDim.x;
    s1[tid] = sum; s2[tid] = ssq;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) { s1[tid] += s1[tid+s]; s2[tid] += s2[tid+s]; }
        __syncthreads();
    }
    if (tid == 0) { psum[c * parts + p] = s1[0]; psq[c * parts + p] = s2[0]; }
}

__global__ void finalize_kernel(const double* __restrict__ psum, const double* __restrict__ psq,
                                float* __restrict__ running_mean, float* __restrict__ running_var,
                                float* __restrict__ mean, float* __restrict__ invstd,
                                int C, int M, int parts, double momentum, double eps) {
    int c = blockIdx.x, tid = threadIdx.x;
    double s = 0.0, q = 0.0;
    for (int i = tid; i < parts; i += blockDim.x) { s += psum[c*parts+i]; q += psq[c*parts+i]; }
    __shared__ double sm[32], sv[32];
    sm[tid] = s; sv[tid] = q;
    __syncthreads();
    for (int off = 16; off > 0; off >>= 1) {
        if (tid < off) { sm[tid] += sm[tid+off]; sv[tid] += sv[tid+off]; }
        __syncthreads();
    }
    if (tid == 0) {
        double mu = sm[0] / (double)M;
        double var = sv[0] / (double)M - mu * mu;
        if (var < 0.0) var = 0.0;
        mean[c] = (float)mu;
        invstd[c] = (float)rsqrt(var + eps);
        double unb = var * ((double)M / (double)(M - 1));
        running_mean[c] = (float)((1.0 - momentum) * (double)running_mean[c] + momentum * mu);
        running_var[c] = (float)((1.0 - momentum) * (double)running_var[c] + momentum * unb);
    }
}

template<int ITEMS>
__global__ void norm_plane_kernel(float* __restrict__ y, const float* __restrict__ mean,
                                  const float* __restrict__ invstd, const float* __restrict__ gamma,
                                  const float* __restrict__ beta, int C, int HW, int tiles) {
    int c = blockIdx.x;
    int n = blockIdx.y;
    int tile = blockIdx.z;
    int base_hw = (tile * blockDim.x + threadIdx.x) * ITEMS;
    int base = (n * C + c) * HW;
    float m = mean[c], is = invstd[c], g = gamma[c], b = beta[c];
    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int hw = base_hw + i;
        if (hw < HW) {
            int idx = base + hw;
            y[idx] = (y[idx] - m) * is * g + b;
        }
    }
}

template<int ITEMS>
__global__ void eval_plane_kernel(const float* __restrict__ x, float* __restrict__ y,
                                  const float* __restrict__ running_mean, const float* __restrict__ running_var,
                                  const float* __restrict__ gamma, const float* __restrict__ beta,
                                  int C, int HW, int tiles, double eps) {
    int c = blockIdx.x;
    int n = blockIdx.y;
    int tile = blockIdx.z;
    int base_hw = (tile * blockDim.x + threadIdx.x) * ITEMS;
    int base = (n * C + c) * HW;
    float m = running_mean[c];
    float is = rsqrtf(running_var[c] + (float)eps);
    float g = gamma[c], b = beta[c];
    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int hw = base_hw + i;
        if (hw < HW) {
            int idx = base + hw;
            float v = mish_act(x[idx]);
            y[idx] = (v - m) * is * g + b;
        }
    }
}

torch::Tensor fused_mish_bn(torch::Tensor x, torch::Tensor running_mean, torch::Tensor running_var,
                            torch::Tensor gamma, torch::Tensor beta,
                            bool training, double momentum, double eps) {
    x = x.contiguous();
    auto y = torch::empty_like(x);
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int HW = H * W;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    constexpr int ITEMS = 8;
    int threads = 256;
    int tiles = (HW + threads * ITEMS - 1) / (threads * ITEMS);
    if (training) {
        int parts = 16;
        auto opts_d = x.options().dtype(torch::kFloat64);
        auto opts_f = x.options().dtype(torch::kFloat32);
        auto psum = torch::empty({C, parts}, opts_d);
        auto psq = torch::empty({C, parts}, opts_d);
        auto mean = torch::empty({C}, opts_f);
        auto invstd = torch::empty({C}, opts_f);
        mish_partial_kernel<<<dim3(C, parts), threads, 2 * threads * sizeof(double), stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), psum.data_ptr<double>(), psq.data_ptr<double>(), N, C, H, W, parts);
        finalize_kernel<<<C, 32, 0, stream>>>(psum.data_ptr<double>(), psq.data_ptr<double>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
            C, N * HW, parts, momentum, eps);
        norm_plane_kernel<ITEMS><<<dim3(C, N, tiles), threads, 0, stream>>>(
            y.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), C, HW, tiles);
    } else {
        eval_plane_kernel<ITEMS><<<dim3(C, N, tiles), threads, 0, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            gamma.data_ptr<float>(), beta.data_ptr<float>(), C, HW, tiles, eps);
    }
    return y;
}
'''

_ext = load_inline(name='fused_mish_bn_ext_v3', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['fused_mish_bn'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels, eps=eps, momentum=momentum)
    def forward(self, x):
        x = self.conv(x)
        if self.training and self.bn.track_running_stats:
            self.bn.num_batches_tracked.add_(1)
            mom = self.bn.momentum
            if mom is None:
                mom = 1.0 / float(self.bn.num_batches_tracked.item())
        else:
            mom = 0.0 if self.bn.momentum is None else float(self.bn.momentum)
        return _ext.fused_mish_bn(x, self.bn.running_mean, self.bn.running_var, self.bn.weight, self.bn.bias, bool(self.training), float(mom), float(self.bn.eps))