import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, subtracts a value, applies HardSwish, MaxPool, and Mish activation functions.
    """
    def __init__(self, in_channels, out_channels, kernel_size, subtract_value, pool_kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.subtract_value = subtract_value
        self.pool = nn.MaxPool2d(pool_kernel_size)

    def forward(self, x):
        x = self.conv(x)
        x = x - self.subtract_value
        x = torch.nn.functional.hardswish(x)
        x = self.pool(x)
        x = torch.nn.functional.mish(x)
        return x

batch_size = 128
in_channels = 64
out_channels = 128
height = width = 128
kernel_size = 3
subtract_value = 0.5
pool_kernel_size = 2

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

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

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor postprocess_forward(torch::Tensor x, double subtract_value, int pool_k);
"""

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

__device__ __forceinline__ float hswish_sub(float v, float sub) {
    float x = v - sub;
    if (x <= -3.0f) return 0.0f;
    if (x >= 3.0f) return x;
    return x * (x + 3.0f) * 0.16666666666666667f;
}

__device__ __forceinline__ float mishf_oneexp(float x) {
    // mish(x)=x*tanh(log(1+exp(x))).  tanh(log(1+e)) simplifies to e*(e+2)/(e*(e+2)+2),
    // saving a log and tanh while matching fp32 PyTorch well within tolerance.
    if (x > 20.0f) return x;
    float e = expf(x);
    return x * (e * (2.0f + e)) / (2.0f + e * (2.0f + e));
}

__global__ void post_pool2_3d_kernel(const float* __restrict__ in, float* __restrict__ out,
                                     int H, int W, int OH, int OW, float sub) {
    int ow = blockIdx.x * blockDim.x + threadIdx.x;
    int oh = blockIdx.y * blockDim.y + threadIdx.y;
    int nc = blockIdx.z;
    if (ow >= OW || oh >= OH) return;

    long base = ((long)nc * H + (oh << 1)) * (long)W + (ow << 1);
    float a = hswish_sub(in[base], sub);
    float b = hswish_sub(in[base + 1], sub);
    float c = hswish_sub(in[base + W], sub);
    float d = hswish_sub(in[base + W + 1], sub);
    float m = fmaxf(fmaxf(a, b), fmaxf(c, d));

    long oidx = ((long)nc * OH + oh) * (long)OW + ow;
    out[oidx] = mishf_oneexp(m);
}

__global__ void post_pool2_1d_kernel(const float* __restrict__ in, float* __restrict__ out,
                                     long total, int H, int W, int OH, int OW, float sub) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    long t = idx / OW;
    int oh = t % OH;
    long nc = t / OH;

    long base = (nc * (long)H + (oh << 1)) * (long)W + (ow << 1);
    float a = hswish_sub(in[base], sub);
    float b = hswish_sub(in[base + 1], sub);
    float c = hswish_sub(in[base + W], sub);
    float d = hswish_sub(in[base + W + 1], sub);
    float m = fmaxf(fmaxf(a, b), fmaxf(c, d));
    out[idx] = mishf_oneexp(m);
}

__global__ void post_poolk_kernel(const float* __restrict__ in, float* __restrict__ out,
                                  long total, int H, int W, int OH, int OW, int K, float sub) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    long t = idx / OW;
    int oh = t % OH;
    long nc = t / OH;

    long base0 = (nc * (long)H + oh * K) * (long)W + ow * K;
    float m = -INFINITY;
    for (int r = 0; r < K; ++r) {
        long base = base0 + r * (long)W;
        for (int c = 0; c < K; ++c) {
            m = fmaxf(m, hswish_sub(in[base + c], sub));
        }
    }
    out[idx] = mishf_oneexp(m);
}

torch::Tensor postprocess_forward(torch::Tensor x, double subtract_value, int pool_k) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    int64_t N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    int64_t OH = H / pool_k, OW = W / pool_k;
    auto out = torch::empty({N, C, OH, OW}, x.options());
    long total = N * C * OH * OW;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (pool_k == 2 && N * C <= 65535) {
        dim3 threads(32, 8, 1);
        dim3 blocks((OW + 31) / 32, (OH + 7) / 8, N * C);
        post_pool2_3d_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), out.data_ptr<float>(), (int)H, (int)W, (int)OH, (int)OW,
            (float)subtract_value);
    } else if (pool_k == 2) {
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        post_pool2_1d_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), out.data_ptr<float>(), total, (int)H, (int)W, (int)OH, (int)OW,
            (float)subtract_value);
    } else {
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        post_poolk_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), out.data_ptr<float>(), total, (int)H, (int)W, (int)OH, (int)OW,
            pool_k, (float)subtract_value);
    }
    return out;
}
"""

_ext = load_inline(
    name="fused_post_hswish_pool_mish_ext_final",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["postprocess_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, subtract_value, pool_kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.subtract_value = subtract_value
        self.pool_kernel_size = pool_kernel_size if isinstance(pool_kernel_size, int) else pool_kernel_size[0]

    def forward(self, x):
        y = F.conv2d(x, self.conv.weight, self.conv.bias)
        return _ext.postprocess_forward(y, float(self.subtract_value), int(self.pool_kernel_size))