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 mish_bn_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                              torch::Tensor running_mean, torch::Tensor running_var,
                              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>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")

static inline int div_up_int64(int64_t a, int b) {
    return (int)((a + b - 1) / b);
}

__device__ __forceinline__ float mish_func(float x) {
    if (x > 20.0f) return x;
    float e = expf(x);
    return x * (e * (e + 2.0f)) / (e * e + 2.0f * e + 2.0f);
}

template<int BLOCK>
__global__ void act_partial_kernel(float* __restrict__ x,
                                   float* __restrict__ psum,
                                   float* __restrict__ psq,
                                   int C, int HW, int M, int parts) {
    int c = blockIdx.x;
    int p = blockIdx.y;
    int tid = threadIdx.x;
    int64_t start = ((int64_t)p * M) / parts;
    int64_t end = ((int64_t)(p + 1) * M) / parts;
    float sum = 0.0f, sq = 0.0f;
    for (int64_t m = start + tid; m < end; m += BLOCK) {
        int n = (int)(m / HW);
        int hw = (int)(m - (int64_t)n * HW);
        int64_t idx = ((int64_t)n * C + c) * HW + hw;
        float a = mish_func(x[idx]);
        x[idx] = a;
        sum += a;
        sq += a * a;
    }
    __shared__ float ssum[BLOCK];
    __shared__ float ssq[BLOCK];
    ssum[tid] = sum;
    ssq[tid] = sq;
    __syncthreads();
    for (int stride = BLOCK / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ssum[tid] += ssum[tid + stride];
            ssq[tid] += ssq[tid + stride];
        }
        __syncthreads();
    }
    if (tid == 0) {
        int off = c * parts + p;
        psum[off] = ssum[0];
        psq[off] = ssq[0];
    }
}

template<int BLOCK>
__global__ void finalize_kernel(const float* __restrict__ psum,
                                const float* __restrict__ psq,
                                const float* __restrict__ weight,
                                const float* __restrict__ bias,
                                float* __restrict__ running_mean,
                                float* __restrict__ running_var,
                                float* __restrict__ scale_out,
                                float* __restrict__ shift_out,
                                int C, int M, int parts, float momentum, float eps) {
    int c = blockIdx.x;
    int tid = threadIdx.x;
    float sum = 0.0f, sq = 0.0f;
    for (int p = tid; p < parts; p += BLOCK) {
        int off = c * parts + p;
        sum += psum[off];
        sq += psq[off];
    }
    __shared__ float ssum[BLOCK];
    __shared__ float ssq[BLOCK];
    ssum[tid] = sum;
    ssq[tid] = sq;
    __syncthreads();
    for (int stride = BLOCK / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ssum[tid] += ssum[tid + stride];
            ssq[tid] += ssq[tid + stride];
        }
        __syncthreads();
    }
    if (tid == 0) {
        float mean = ssum[0] / (float)M;
        float var = ssq[0] / (float)M - mean * mean;
        if (var < 0.0f) var = 0.0f;
        float inv = rsqrtf(var + eps);
        float w = weight[c];
        scale_out[c] = w * inv;
        shift_out[c] = bias[c] - mean * w * inv;
        running_mean[c] = running_mean[c] * (1.0f - momentum) + momentum * mean;
        running_var[c] = running_var[c] * (1.0f - momentum) + momentum * (var * ((float)M / (float)(M - 1)));
    }
}

__global__ void norm_kernel(float* __restrict__ x,
                            const float* __restrict__ scale,
                            const float* __restrict__ shift,
                            int64_t total, int C, int HW) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = (int)((idx / HW) % C);
        x[idx] = x[idx] * scale[c] + shift[c];
    }
}

__global__ void eval_kernel(float* __restrict__ x,
                            const float* __restrict__ weight,
                            const float* __restrict__ bias,
                            const float* __restrict__ running_mean,
                            const float* __restrict__ running_var,
                            int64_t total, int C, int HW, float eps) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = (int)((idx / HW) % C);
        float a = mish_func(x[idx]);
        x[idx] = (a - running_mean[c]) * rsqrtf(running_var[c] + eps) * weight[c] + bias[c];
    }
}

torch::Tensor mish_bn_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                              torch::Tensor running_mean, torch::Tensor running_var,
                              bool training, double momentum, double eps) {
    CHECK_CUDA(x);
    CHECK_CUDA(weight);
    CHECK_CUDA(bias);
    CHECK_CUDA(running_mean);
    CHECK_CUDA(running_var);
    TORCH_CHECK(x.dtype() == torch::kFloat32, "float32 only");
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int HW = (int)(x.size(2) * x.size(3));
    int M = N * HW;
    int64_t total = x.numel();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    const int parts = 16;
    if (training) {
        auto opts = x.options();
        auto partial_sum = torch::empty({C, parts}, opts);
        auto partial_sq = torch::empty({C, parts}, opts);
        auto scale = torch::empty({C}, opts);
        auto shift = torch::empty({C}, opts);
        act_partial_kernel<256><<<dim3(C, parts), threads, 0, stream>>>(
            x.data_ptr<float>(), partial_sum.data_ptr<float>(), partial_sq.data_ptr<float>(), C, HW, M, parts);
        finalize_kernel<256><<<C, threads, 0, stream>>>(
            partial_sum.data_ptr<float>(), partial_sq.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(),
            C, M, parts, (float)momentum, (float)eps);
        int blocks = (int)std::min<int64_t>(4096, div_up_int64(total, threads));
        norm_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(), total, C, HW);
    } else {
        int blocks = (int)std::min<int64_t>(4096, div_up_int64(total, threads));
        eval_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(), total, C, HW, (float)eps);
    }
    return x;
}
'''

_ext = load_inline(
    name='mish_bn_final_ext',
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=['mish_bn_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, 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):
        y = self.conv(x)
        return _ext.mish_bn_forward(
            y, self.bn.weight, self.bn.bias, self.bn.running_mean, self.bn.running_var,
            self.training, float(self.bn.momentum), float(self.bn.eps)
        )