import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, applies Mish activation, adds a value, 
    applies Hardtanh activation, and scales the output.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
        self.add_value = add_value
        self.scale = scale

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.nn.functional.mish(x) # Mish activation
        x = x + self.add_value
        x = torch.nn.functional.hardtanh(x, min_val=-1, max_val=1) # Hardtanh activation
        x = x * self.scale # Scaling
        return x

batch_size = 128
in_channels  = 64  
out_channels = 64  
height = width = 128  
kernel_size  = 3
stride       = 2  
padding      = 1
output_padding = 1
add_value = 0.5
scale = 2

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale]

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 fused_mish_add_hardtanh_scale(torch::Tensor x, double add_value, double scale);
"""

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

__device__ __forceinline__ float softplus_stable(float v) {
    // Matches the usual stable softplus formulation used by PyTorch closely
    // enough for fp32 tolerance, while avoiding overflow for large positives.
    if (v > 20.0f) return v;
    if (v < -20.0f) return __expf(v);
    return log1pf(__expf(v));
}

__device__ __forceinline__ float op_one(float v, float addv, float scale) {
    // mish(v) = v * tanh(softplus(v))
    float sp = softplus_stable(v);
    float y = v * tanhf(sp);
    y += addv;
    // hardtanh with min=-1, max=1, then scale
    y = fminf(1.0f, fmaxf(-1.0f, y));
    return y * scale;
}

__global__ void fused_mish_add_hardtanh_scale_kernel(float* __restrict__ data,
                                                     int64_t n,
                                                     float addv,
                                                     float scale) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < n; idx += stride) {
        data[idx] = op_one(data[idx], addv, scale);
    }
}

torch::Tensor fused_mish_add_hardtanh_scale(torch::Tensor x, double add_value, double scale) {
    const int64_t n = x.numel();
    if (n == 0) return x;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    // Keep the grid reasonably sized; the grid-stride loop covers the rest.
    if (blocks > 65535) blocks = 65535;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_mish_add_hardtanh_scale_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), n, (float)add_value, (float)scale);
    return x;
}
"""

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


class ModelNew(nn.Module):
    """
    ConvTranspose2d followed by Mish, add, Hardtanh and scaling.
    The expensive convolution is delegated to the same optimized backend as
    PyTorch, while the four following pointwise operations are fused into a
    single custom CUDA kernel that operates in-place on the convolution output.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
        self.add_value = add_value
        self.scale = scale

    def forward(self, x):
        y = self.conv_transpose(x)
        return _ext.fused_mish_add_hardtanh_scale(y, float(self.add_value), float(self.scale))