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

# CUDA kernels:
#  1) nchw_to_nhwc_pad_kernel: fused NCHW -> NHWC transpose + zero pad of the
#     bottom/right edges (feeds the 2x2 stride-1 reformulated convolution).
#  2) fused_subpixel_epilogue_kernel: reads the NHWC conv output where each
#     group of 4 consecutive channels encodes a 2x2 sub-pixel block, applies
#     bias + Mish + add_value + Hardtanh + scale, and writes the final NHWC
#     tensor with the sub-pixels interleaved (pixel-shuffle fused for free).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.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)); with t = exp(x): tanh = u/(u+2), u = t*(t+2)
    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;
}

// Fused NCHW -> NHWC transpose with bottom/right zero padding by 1.
// x:   (N, C, H, W) contiguous fp32
// out: (N, H+1, W+1, C) contiguous fp32 (i.e. channels_last (N, C, H+1, W+1))
__global__ void nchw_to_nhwc_pad_kernel(const float* __restrict__ x,
                                        float* __restrict__ out,
                                        int N, int C, int H, int W) {
    extern __shared__ float tile[];  // (W) * (C + 1) floats
    const int n = blockIdx.x;
    const int h = blockIdx.y;  // in [0, H]
    const bool hv = (h < H);
    const int tid = threadIdx.x;

    for (int idx = tid; idx < C * W; idx += blockDim.x) {
        int w = idx % W;
        int c = idx / W;
        tile[w * (C + 1) + c] = hv ? x[((long)(n * C + c) * H + h) * W + w] : 0.0f;
    }
    __syncthreads();

    long base = ((long)(n * (H + 1) + h) * (W + 1)) * C;
    for (int idx = tid; idx < (W + 1) * C; idx += blockDim.x) {
        int c = idx % C;
        int w = idx / C;
        out[base + (long)w * C + c] = (w < W) ? tile[w * (C + 1) + c] : 0.0f;
    }
}

// Fused: pixel-shuffle(2) + bias + Mish + add + Hardtanh + scale.
// in:  (N, Hi, Wi, 4*Co) fp32, channels ordered as co*4 + (dy*2+dx)
// out: (N, 2*Hi, 2*Wi, Co) fp32 (channels_last view of (N, Co, 2Hi, 2Wi))
// Fast path: Co divides blockDim.x (e.g. Co=64 with 256 threads).
__global__ void fused_subpixel_epilogue_kernel(const float4* __restrict__ in,
                                               float* __restrict__ out,
                                               const float* __restrict__ bias,
                                               int N, int Co, int Hi, int Wi,
                                               float add_value, float scale) {
    const int jpb = blockDim.x / Co;             // js per block along Wi
    const int co = threadIdx.x % Co;
    const int js = threadIdx.x / Co;
    const int j = blockIdx.x * jpb + js;
    const int i = blockIdx.y;
    const int n = blockIdx.z;
    if (j >= Wi) return;

    long e = ((long)(n * Hi + i) * Wi + j) * Co + co;   // float4 index
    float4 v = in[e];
    float b = bias[co];

    float r0 = epilogue(v.x, b, add_value, scale);  // (dy=0, dx=0) -> (2i,   2j)
    float r1 = epilogue(v.y, b, add_value, scale);  // (dy=0, dx=1) -> (2i,   2j+1)
    float r2 = epilogue(v.z, b, add_value, scale);  // (dy=1, dx=0) -> (2i+1, 2j)
    float r3 = epilogue(v.w, b, add_value, scale);  // (dy=1, dx=1) -> (2i+1, 2j+1)

    const int Wo = 2 * Wi;
    long p = (((long)n * 2 * Hi + 2 * i) * Wo + 2 * j) * Co + co;
    long rstride = (long)Wo * Co;
    out[p] = r0;
    out[p + Co] = r1;
    out[p + rstride] = r2;
    out[p + rstride + Co] = r3;
}

// Generic fallback (any Co / shape).
__global__ void fused_subpixel_epilogue_generic_kernel(const float* __restrict__ in,
                                                       float* __restrict__ out,
                                                       const float* __restrict__ bias,
                                                       int N, int Co, int Hi, int Wi,
                                                       float add_value, float scale) {
    const int Wo = 2 * Wi;
    const long total = (long)N * Hi * Wi * Co;
    const long stride = (long)gridDim.x * blockDim.x;
    for (long t = (long)blockIdx.x * blockDim.x + threadIdx.x; t < total; t += stride) {
        int co = (int)(t % Co);
        long tmp = t / Co;
        int j = (int)(tmp % Wi); tmp /= Wi;
        int i = (int)(tmp % Hi);
        int n = (int)(tmp / Hi);
        const float* src = in + (((long)(n * Hi + i) * Wi + j) * 4 * Co + 4 * co);
        float b = bias[co];
        float r0 = epilogue(src[0], b, add_value, scale);
        float r1 = epilogue(src[1], b, add_value, scale);
        float r2 = epilogue(src[2], b, add_value, scale);
        float r3 = epilogue(src[3], b, add_value, scale);
        long p = (((long)n * 2 * Hi + 2 * i) * Wo + 2 * j) * Co + co;
        long rstride = (long)Wo * Co;
        out[p] = r0;
        out[p + Co] = r1;
        out[p + rstride] = r2;
        out[p + rstride + Co] = r3;
    }
}

torch::Tensor nchw_to_nhwc_pad_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32, "x must be CUDA fp32");
    TORCH_CHECK(x.dim() == 4 && x.is_contiguous(), "x must be contiguous NCHW");
    const int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    auto out = torch::empty({N, C, H + 1, W + 1}, x.options())
                   .contiguous(at::MemoryFormat::ChannelsLast);
    size_t smem = (size_t)W * (C + 1) * sizeof(float);
    TORCH_CHECK(smem <= 96 * 1024, "tile too large for shared memory");
    dim3 grid(N, H + 1);
    auto stream = at::cuda::getCurrentCUDAStream();
    nchw_to_nhwc_pad_kernel<<<grid, 256, smem, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), N, C, H, W);
    return out;
}

torch::Tensor fused_subpixel_epilogue_cuda(torch::Tensor y, torch::Tensor bias,
                                           double add_value, double scale) {
    TORCH_CHECK(y.is_cuda() && y.scalar_type() == torch::kFloat32, "y must be CUDA fp32");
    TORCH_CHECK(y.is_contiguous(at::MemoryFormat::ChannelsLast), "y must be channels_last");
    const int N = y.size(0), C4 = y.size(1), Hi = y.size(2), Wi = y.size(3);
    TORCH_CHECK(C4 % 4 == 0, "channels must be divisible by 4");
    const int Co = C4 / 4;
    TORCH_CHECK(bias.numel() == Co, "bias size mismatch");
    auto out = torch::empty({N, Co, 2 * Hi, 2 * Wi}, y.options())
                   .contiguous(at::MemoryFormat::ChannelsLast);
    if (out.numel() == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    const int block = 256;
    if (block % Co == 0) {
        dim3 grid((Wi + (block / Co) - 1) / (block / Co), Hi, N);
        fused_subpixel_epilogue_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(y.data_ptr<float>()),
            out.data_ptr<float>(), bias.data_ptr<float>(),
            N, Co, Hi, Wi, (float)add_value, (float)scale);
    } else {
        long total = (long)N * Hi * Wi * Co;
        long blocks = (total + block - 1) / block;
        if (blocks > 2147483647L) blocks = 2147483647L;
        fused_subpixel_epilogue_generic_kernel<<<(unsigned int)blocks, block, 0, stream>>>(
            y.data_ptr<float>(), out.data_ptr<float>(), bias.data_ptr<float>(),
            N, Co, Hi, Wi, (float)add_value, (float)scale);
    }
    return out;
}
"""

cpp_source = (
    "torch::Tensor nchw_to_nhwc_pad_cuda(torch::Tensor x);\n"
    "torch::Tensor fused_subpixel_epilogue_cuda(torch::Tensor y, torch::Tensor bias, "
    "double add_value, double scale);"
)

_fused = load_inline(
    name="fused_deconv_subpixel_epilogue_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["nchw_to_nhwc_pad_cuda", "fused_subpixel_epilogue_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model.

    Key idea: for stride=2, kernel=3, padding=1, output_padding=1 the
    transposed convolution decomposes exactly into 4 small stride-1
    sub-convolutions (one per 2x2 sub-pixel phase) applied to the input,
    followed by a 2x2 interleave (pixel shuffle). Each sub-convolution uses
    only taps from a 2x2 window of the input padded by one row/column on the
    bottom/right, so the whole deconv equals:

        conv2d(pad(x, (0,1,0,1)), W2, stride=1)  with W2: (4*Cout, Cin, 2, 2)

    where output channel co*4+(dy*2+dx) is sub-pixel (dy,dx). This replaces
    cudnn's strided dgrad with a much faster forward implicit-GEMM conv.

    Pipeline (all NHWC / channels_last, TF32 tensor cores):
      1) custom kernel: NCHW input -> NHWC, zero-padded by 1 on bottom/right
      2) cudnn conv2d 2x2 stride 1 (bias folded into epilogue)
      3) custom kernel: fused pixel-shuffle + bias + Mish + add + Hardtanh
         + scale, writing the final NHWC tensor with float4 reads and fully
         coalesced writes.

    The conv_transpose module (incl. parameters) is kept intact so the
    state_dict matches the reference model; W2 is rebuilt lazily whenever the
    weight version changes. For parameter combinations outside the supported
    pattern, falls back to the original ops.
    """
    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._w2 = None
        self._w2_version = None

    def _supported(self):
        ct = self.conv_transpose
        return (
            ct.in_channels == ct.out_channels  # not required, but keep general check below
            and ct.kernel_size == (3, 3)
            and ct.stride == (2, 2)
            and ct.padding == (1, 1)
            and ct.output_padding == (1, 1)
            and ct.dilation == (1, 1)
            and ct.groups == 1
        )

    def _build_w2(self):
        w = self.conv_transpose.weight.detach()  # (Cin, Cout, 3, 3)
        wt = w.permute(1, 0, 2, 3).contiguous()  # (Cout, Cin, 3, 3)
        Cout, Cin = wt.shape[0], wt.shape[1]
        w2 = torch.empty(Cout * 4, Cin, 2, 2, device=wt.device, dtype=wt.dtype)
        # sub-pixel (dy, dx) -> channel co*4 + dy*2 + dx
        w2[0::4, :, 0, 0] = wt[:, :, 1, 1]  # (0,0)
        w2[1::4, :, 0, 0] = wt[:, :, 1, 2]  # (0,1)
        w2[1::4, :, 0, 1] = wt[:, :, 1, 0]
        w2[2::4, :, 0, 0] = wt[:, :, 2, 1]  # (1,0)
        w2[2::4, :, 1, 0] = wt[:, :, 0, 1]
        w2[3::4, :, 0, 0] = wt[:, :, 2, 2]  # (1,1)
        w2[3::4, :, 0, 1] = wt[:, :, 2, 0]
        w2[3::4, :, 1, 0] = wt[:, :, 0, 2]
        w2[3::4, :, 1, 1] = wt[:, :, 0, 0]
        self._w2 = w2.contiguous(memory_format=torch.channels_last)
        self._w2_version = self.conv_transpose.weight._version

    def forward(self, x):
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and self._supported()
            and x.size(2) >= 1
            and x.size(3) >= 1
        ):
            if self._w2 is None or self._w2_version != self.conv_transpose.weight._version:
                self._build_w2()
            x = x.contiguous()
            xp = _fused.nchw_to_nhwc_pad_cuda(x)  # (N, C, H+1, W+1) channels_last
            y = F.conv2d(xp, self._w2, None, stride=1, padding=0)
            if not y.is_contiguous(memory_format=torch.channels_last):
                y = y.contiguous(memory_format=torch.channels_last)
            bias = self.conv_transpose.bias
            if bias is None:
                bias = torch.zeros(y.size(1) // 4, device=y.device, dtype=y.dtype)
            bias = bias.contiguous()
            return _fused.fused_subpixel_epilogue_cuda(y, bias, float(self.add_value), float(self.scale))

        # Reference fallback path.
        x = self.conv_transpose(x)
        x = F.mish(x)
        x = x + self.add_value
        x = F.hardtanh(x, min_val=-1, max_val=1)
        x = x * self.scale
        return x
```