```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

torch.backends.cudnn.benchmark = True

# Custom fused CUDA kernel: Mish -> Add -> Hardtanh -> Scale in a single pass.
# Mish(x) = x * tanh(softplus(x)). Using t = exp(x):
#   tanh(softplus(x)) = (t*(t+2)) / (t*(t+2) + 2)
# which needs only one expf and a few FLOPs (no log/tanh), keeping the
# elementwise epilogue purely memory-bound.
fused_epilogue_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float epilogue(float v, float add_value, float scale) {
    // mish(x) = x * tanh(softplus(x)); clamp exp arg to avoid overflow
    float t = expf(fminf(v, 20.0f));
    float u = t * (t + 2.0f);
    float mish = v * (u / (u + 2.0f));
    float y = mish + add_value;
    y = fminf(fmaxf(y, -1.0f), 1.0f);   // hardtanh in [-1, 1]
    return y * scale;
}

__global__ void fused_epilogue_vec_kernel(const float4* __restrict__ in,
                                          float4* __restrict__ out,
                                          int n4,
                                          float add_value,
                                          float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < n4; i += stride) {
        float4 v = in[i];
        float4 r;
        r.x = epilogue(v.x, add_value, scale);
        r.y = epilogue(v.y, add_value, scale);
        r.z = epilogue(v.z, add_value, scale);
        r.w = epilogue(v.w, add_value, scale);
        out[i] = r;
    }
}

__global__ void fused_epilogue_kernel(const float* __restrict__ in,
                                      float* __restrict__ out,
                                      int n,
                                      float add_value,
                                      float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < n; i += stride) {
        out[i] = epilogue(in[i], add_value, scale);
    }
}

torch::Tensor fused_epilogue_cuda(torch::Tensor x, double add_value, double scale) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);
    const long n = xc.numel();
    if (n == 0) return out;

    auto stream = at::cuda::getCurrentCUDAStream();
    const int block_size = 256;

    if ((n & 3) == 0) {
        const int n4 = (int)(n >> 2);
        const int num_blocks = (n4 + block_size - 1) / block_size;
        fused_epilogue_vec_kernel<<<num_blocks, block_size, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            n4, (float)add_value, (float)scale);
    } else {
        const int num_blocks = ((int)n + block_size - 1) / block_size;
        fused_epilogue_kernel<<<num_blocks, block_size, 0, stream>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(),
            (int)n, (float)add_value, (float)scale);
    }
    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_epilogue_cuda(torch::Tensor x, double add_value, double scale);"
)

fused_epilogue = load_inline(
    name="fused_mish_add_hardtanh_scale",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_source,
    functions=["fused_epilogue_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN-accelerated ConvTranspose2d and fuses
    Mish + add + Hardtanh + scale into a single custom CUDA kernel, removing
    three full memory round-trips over the (large) activation tensor.
    """
    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
        self.fused_epilogue = fused_epilogue

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.fused_epilogue.fused_epilogue_cuda(x, self.add_value, self.scale)
        return x
```