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

torch.backends.cudnn.benchmark = True
torch.backends.cudnn.benchmark_limit = 0  # exhaustive cudnn algo search
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True

# Fused epilogue operating directly on the channels_last (NHWC) output of the
# transposed convolution. In NHWC a flat index maps to channel via (idx % C),
# so one float4 spans 4 consecutive channels and the bias is an aligned float4
# load. Fuses: bias add -> Mish -> add_value -> Hardtanh -> scale into a single
# streaming pass with cache-streaming loads/stores (data is used exactly once).
fused_epilogue_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

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

__device__ __forceinline__ float4 epilogue4(float4 val, float4 b, float add_value, float scale) {
    float4 r;
    r.x = epilogue(val.x, b.x, add_value, scale);
    r.y = epilogue(val.y, b.y, add_value, scale);
    r.z = epilogue(val.z, b.z, add_value, scale);
    r.w = epilogue(val.w, b.w, add_value, scale);
    return r;
}

// Vectorized path (numel % 4 == 0 and C % 4 == 0). Each thread handles two
// float4s in an interleaved pattern to keep full coalescing while increasing
// memory-level parallelism. Streaming cache hints avoid L2 thrash.
__global__ __launch_bounds__(256) void fused_epilogue_nhwc_vec4_kernel(
        const float4* __restrict__ in,
        float4* __restrict__ out,
        const float4* __restrict__ bias4,
        long n4, long c4,
        float add_value, float scale) {
    long stride = (long)gridDim.x * blockDim.x * 2;
    for (long base = (long)blockIdx.x * blockDim.x * 2 + threadIdx.x; base < n4; base += stride) {
        long v0 = base;
        long v1 = base + blockDim.x;
        {
            float4 val = __ldcs(in + v0);
            float4 b = __ldg(bias4 + (v0 % c4));
            __stcs(out + v0, epilogue4(val, b, add_value, scale));
        }
        if (v1 < n4) {
            float4 val = __ldcs(in + v1);
            float4 b = __ldg(bias4 + (v1 % c4));
            __stcs(out + v1, epilogue4(val, b, add_value, scale));
        }
    }
}

// Generic fallback path.
__global__ void fused_epilogue_nhwc_kernel(const float* __restrict__ in,
                                           float* __restrict__ out,
                                           const float* __restrict__ bias,
                                           long n, long C,
                                           float add_value, float scale) {
    long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < n; i += stride) {
        long c = i % C;
        out[i] = epilogue(in[i], bias[c], add_value, scale);
    }
}

torch::Tensor fused_epilogue_nhwc_cuda(torch::Tensor x, torch::Tensor bias, 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");
    TORCH_CHECK(x.is_contiguous(at::MemoryFormat::ChannelsLast), "x must be channels_last contiguous");
    TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kFloat32, "bias must be CUDA float32");

    const c10::cuda::OptionalCUDAGuard guard(x.device());
    auto b = bias.contiguous();
    const long C = x.size(1);
    TORCH_CHECK(b.numel() == C, "bias size mismatch");

    auto out = torch::empty_like(x);  // preserves channels_last layout
    const long n = x.numel();
    if (n == 0) return out;

    auto stream = at::cuda::getCurrentCUDAStream();

    if ((n % 4 == 0) && (C % 4 == 0) &&
        ((reinterpret_cast<uintptr_t>(x.data_ptr<float>()) & 0xF) == 0) &&
        ((reinterpret_cast<uintptr_t>(out.data_ptr<float>()) & 0xF) == 0) &&
        ((reinterpret_cast<uintptr_t>(b.data_ptr<float>()) & 0xF) == 0)) {
        const int block = 256;
        const long n4 = n >> 2;
        long blocks = (n4 + (long)block * 2 - 1) / ((long)block * 2);
        if (blocks > 2147483647L) blocks = 2147483647L;
        fused_epilogue_nhwc_vec4_kernel<<<(unsigned int)blocks, block, 0, stream>>>(
            reinterpret_cast<const float4*>(x.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            reinterpret_cast<const float4*>(b.data_ptr<float>()),
            n4, C >> 2,
            (float)add_value, (float)scale);
    } else {
        const int block = 256;
        long blocks = (n + block - 1) / block;
        if (blocks > 2147483647L) blocks = 2147483647L;
        fused_epilogue_nhwc_kernel<<<(unsigned int)blocks, block, 0, stream>>>(
            x.data_ptr<float>(), out.data_ptr<float>(), b.data_ptr<float>(),
            n, C, (float)add_value, (float)scale);
    }
    return out;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model:
      * Runs the transposed convolution in channels_last so cudnn's fast NHWC
        dgrad GEMM is used without the expensive NHWC->NCHW output transpose.
      * Caches a channels_last copy of the conv weight so neither ATen nor
        cudnn has to re-transpose the filter on every forward pass.
      * Calls F.conv_transpose2d with bias=None and folds the bias into a
        single fused CUDA epilogue (bias + Mish + add + Hardtanh + scale),
        removing ATen's separate bias-add pass and three memory round-trips.
      * The conv_transpose module (incl. its bias parameter) is kept intact so
        the state_dict remains identical to the reference model.
    """
    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
        self._w_cl = None
        self._w_key = None
        self._zero_bias = None

    def _weight_channels_last(self, w):
        key = (w.data_ptr(), w._version, w.device, w.dtype)
        if self._w_cl is None or self._w_key != key:
            self._w_cl = w.detach().contiguous(memory_format=torch.channels_last)
            self._w_key = key
        return self._w_cl

    def forward(self, x):
        if (not x.is_cuda) or x.dtype != torch.float32:
            y = self.conv_transpose(x)
            y = F.mish(y)
            y = y + self.add_value
            y = F.hardtanh(y, min_val=-1, max_val=1)
            return y * self.scale

        ct = self.conv_transpose
        x = x.contiguous(memory_format=torch.channels_last)
        w = self._weight_channels_last(ct.weight)
        y = F.conv_transpose2d(
            x,
            w,
            None,  # bias folded into the fused epilogue
            ct.stride,
            ct.padding,
            ct.output_padding,
            ct.groups,
            ct.dilation,
        )
        if not y.is_contiguous(memory_format=torch.channels_last):
            y = y.contiguous(memory_format=torch.channels_last)
        bias = ct.bias
        if bias is None:
            if self._zero_bias is None or self._zero_bias.device != y.device or self._zero_bias.numel() != y.size(1):
                self._zero_bias = torch.zeros(y.size(1), device=y.device, dtype=torch.float32)
            bias = self._zero_bias
        return self.fused_epilogue.fused_epilogue_nhwc_cuda(y, bias, float(self.add_value), float(self.scale))
```