```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA kernels
#
# The module computes:  y = scale * hardtanh(mish(ConvTranspose2d(x)) + add)
#
# Optimizations:
#  1. The transposed convolution (k=3, s=2, p=1, op=1) is mathematically
#     reformulated as a *regular* 2x2 convolution producing 4 phase channels
#     per output channel (sub-pixel / pixel-shuffle decomposition).  This maps
#     directly onto cuDNN's tensor-core friendly forward-conv kernels.
#  2. pixel-shuffle + Mish + add + Hardtanh + scale are fused into a single
#     memory-bound CUDA kernel, eliminating 4 full passes over the 2GB+ feature
#     map.  Mish is evaluated with the algebraically equivalent, well
#     conditioned form  x * (u^2 + 2u) / (u^2 + 2u + 2),  u = exp(x).
# ---------------------------------------------------------------------------

_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float epilogue_op(float v, float add_value, float scale) {
    // mish(x) = x * tanh(softplus(x)) = x * (u^2 + 2u) / (u^2 + 2u + 2), u = e^x
    float u = __expf(v);
    u = fminf(u, 1.0e18f);                 // avoid inf/inf for very large v
    float p = fmaf(u, u, 2.0f * u);        // u^2 + 2u  (well conditioned)
    float m = v * __fdividef(p, p + 2.0f); // mish(v)
    float y = m + add_value;
    y = fminf(fmaxf(y, -1.0f), 1.0f);      // hardtanh(-1, 1)
    return y * scale;
}

// Fused pixel-shuffle (factor 2) + epilogue.
// in  : (N, C*4, Hp, Wp) view described by strides (may be a cropped view)
// out : (N, C, 2*Hp, 2*Wp) contiguous
__global__ void shuffle_epilogue_kernel(
        const float* __restrict__ in, float* __restrict__ out,
        int N, int C, int Hp, int Wp,
        long sN, long sC, long sH, long sW,
        float add_value, float scale) {
    int ow = blockIdx.x * blockDim.x + threadIdx.x;
    int oh = blockIdx.y * blockDim.y + threadIdx.y;
    int z  = blockIdx.z;
    if (ow >= Wp || oh >= Hp) return;
    int c = z % C;
    int n = z / C;

    const float* base = in + (long)n * sN + (long)(c * 4) * sC
                        + (long)oh * sH + (long)ow * sW;
    float v00 = __ldg(base);
    float v01 = __ldg(base + sC);
    float v10 = __ldg(base + 2 * sC);
    float v11 = __ldg(base + 3 * sC);

    int H = 2 * Hp, W = 2 * Wp;
    float* ob = out + (((long)(n * C + c) * H + 2 * oh) * W) + 2 * ow;
    float2 r0, r1;
    r0.x = epilogue_op(v00, add_value, scale);
    r0.y = epilogue_op(v01, add_value, scale);
    r1.x = epilogue_op(v10, add_value, scale);
    r1.y = epilogue_op(v11, add_value, scale);
    *reinterpret_cast<float2*>(ob)     = r0;
    *reinterpret_cast<float2*>(ob + W) = r1;
}

// Plain fused epilogue (mish + add + hardtanh + scale), vectorized.
__global__ void epilogue_vec4_kernel(const float4* __restrict__ in,
                                     float4* __restrict__ out, long n4,
                                     float add_value, float scale) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 v = in[i];
        v.x = epilogue_op(v.x, add_value, scale);
        v.y = epilogue_op(v.y, add_value, scale);
        v.z = epilogue_op(v.z, add_value, scale);
        v.w = epilogue_op(v.w, add_value, scale);
        out[i] = v;
    }
}

__global__ void epilogue_scalar_kernel(const float* __restrict__ in,
                                       float* __restrict__ out, long n,
                                       float add_value, float scale) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) out[i] = epilogue_op(in[i], add_value, scale);
}

torch::Tensor shuffle_epilogue(torch::Tensor x, int64_t C,
                               double add_value, double scale) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32 && x.dim() == 4,
                "shuffle_epilogue expects a CUDA float32 4D tensor");
    int64_t N  = x.size(0);
    int64_t Hp = x.size(2);
    int64_t Wp = x.size(3);
    auto out = torch::empty({N, C, 2 * Hp, 2 * Wp}, x.options());

    dim3 block(32, 8);
    dim3 grid((unsigned)((Wp + block.x - 1) / block.x),
              (unsigned)((Hp + block.y - 1) / block.y),
              (unsigned)(N * C));
    shuffle_epilogue_kernel<<<grid, block>>>(
        x.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)C, (int)Hp, (int)Wp,
        (long)x.stride(0), (long)x.stride(1), (long)x.stride(2), (long)x.stride(3),
        (float)add_value, (float)scale);
    return out;
}

torch::Tensor epilogue(torch::Tensor x, double add_value, double scale) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "epilogue expects a CUDA float32 tensor");
    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);
    long n = xc.numel();
    if ((n & 3) == 0) {
        long n4 = n >> 2;
        long blocks = (n4 + 255) / 256;
        epilogue_vec4_kernel<<<(unsigned)blocks, 256>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            n4, (float)add_value, (float)scale);
    } else {
        long blocks = (n + 255) / 256;
        epilogue_scalar_kernel<<<(unsigned)blocks, 256>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(), n,
            (float)add_value, (float)scale);
    }
    return out;
}
"""

_cpp_source = (
    "torch::Tensor shuffle_epilogue(torch::Tensor x, int64_t C, double add_value, double scale);\n"
    "torch::Tensor epilogue(torch::Tensor x, double add_value, double scale);"
)

_ext = load_inline(
    name="fused_conv_transpose_epilogue",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["shuffle_epilogue", "epilogue"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - ConvTranspose2d(k=3, s=2, p=1, op=1) is computed as a regular 2x2
        Conv2d with rearranged weights (exact sub-pixel decomposition), which
        is much friendlier to cuDNN/tensor cores.
      - pixel-shuffle + Mish + add + Hardtanh + scaling are fused into one
        custom CUDA kernel.
    """

    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 = float(add_value)
        self.scale = float(scale)
        self._out_channels = out_channels
        self._fast_path = (int(kernel_size) == 3 and int(stride) == 2
                           and int(padding) == 1 and int(output_padding) == 1)
        self._w2 = None
        self._b2 = None
        self._cache_key = None

    def _build_fast_weights(self, device, dtype):
        w = self.conv_transpose.weight.detach()          # (Cin, Cout, 3, 3)
        Cin, Cout = w.shape[0], w.shape[1]
        wp = w.permute(1, 0, 2, 3).contiguous()          # (Cout, Cin, 3, 3)
        w2 = torch.zeros(Cout * 4, Cin, 2, 2, device=device, dtype=dtype)
        # K2[oc*4 + r*2 + c, ic, dy, dx] = W[ic, oc, r+1-2dy, c+1-2dx] (0 if OOB)
        w2[0::4, :, 0, 0] = wp[:, :, 1, 1]   # phase (0,0)
        w2[1::4, :, 0, 0] = wp[:, :, 1, 2]   # phase (0,1)
        w2[1::4, :, 0, 1] = wp[:, :, 1, 0]
        w2[2::4, :, 0, 0] = wp[:, :, 2, 1]   # phase (1,0)
        w2[2::4, :, 1, 0] = wp[:, :, 0, 1]
        w2[3::4, :, 0, 0] = wp[:, :, 2, 2]   # phase (1,1)
        w2[3::4, :, 0, 1] = wp[:, :, 2, 0]
        w2[3::4, :, 1, 0] = wp[:, :, 0, 2]
        w2[3::4, :, 1, 1] = wp[:, :, 0, 0]
        b = self.conv_transpose.bias
        b2 = b.detach().repeat_interleave(4) if b is not None else None
        self._w2 = w2
        self._b2 = b2
        wsrc = self.conv_transpose.weight
        bsrc = self.conv_transpose.bias
        self._cache_key = (wsrc.data_ptr(), wsrc._version,
                           bsrc.data_ptr() if bsrc is not None else None,
                           bsrc._version if bsrc is not None else None,
                           device, dtype)

    def _reference_path(self, x):
        x = self.conv_transpose(x)
        x = F.mish(x)
        x = x + self.add_value
        x = F.hardtanh(x, min_val=-1, max_val=1)
        return x * self.scale

    def forward(self, x):
        if (not x.is_cuda) or x.dtype != torch.float32 or not self._fast_path:
            if x.is_cuda and x.dtype == torch.float32:
                # Generic CUDA fallback: cuDNN transposed conv + fused epilogue
                y = self.conv_transpose(x.contiguous())
                return _ext.epilogue(y, self.add_value, self.scale)
            return self._reference_path(x)

        w = self.conv_transpose.weight
        b = self.conv_transpose.bias
        key = (w.data_ptr(), w._version,
               b.data_ptr() if b is not None else None,
               b._version if b is not None else None,
               x.device, x.dtype)
        if self._cache_key != key:
            self._build_fast_weights(x.device, x.dtype)

        xc = x if x.is_contiguous() else x.contiguous()
        # Regular 2x2 conv with padding 1; the [1:, 1:] crop realigns phases.
        y = F.conv2d(xc, self._w2, self._b2, stride=1, padding=1)
        yv = y[:, :, 1:, 1:]
        return _ext.shuffle_epilogue(yv, self._out_channels,
                                     self.add_value, self.scale)


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]
```