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

# ---------------------------------------------------------------------------
# Key idea:
#   ConvTranspose2d(k=4, s=2, p=1, op=1) is mathematically equivalent to:
#     1) A regular 2x2 Conv2d producing 4*out_channels feature maps
#        (the "sub-pixel" decomposition of a transposed convolution), which
#        hits cuDNN's highly-optimized forward-conv path (tensor cores).
#     2) A phase-interleave (pixel-shuffle-like) remap of those feature maps.
#   Step (2) is fused together with the bias subtraction and the tanh
#   activation into a single custom CUDA kernel, so the epilogue costs exactly
#   one read + one write of the activations instead of multiple full passes.
#
#   Weight transform (exact):
#     W_eq[(oc*2 + ry)*2 + rx, ic, dy, dx] = W_t[ic, oc, 3 - 2*dy - ry, 3 - 2*dx - rx]
#   and out[b, oc, 2Y+ry, 2X+rx] = conv[b, (oc,ry,rx), Y+ry, X+rx].
# ---------------------------------------------------------------------------

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

// Fused: phase-interleave (pixel-shuffle) + bias subtract + tanh.
//   conv: [B, OC*4, Hc, Wc]   (result of the equivalent 2x2 conv, pad=1)
//   bias: [OC]                (effective bias to subtract per output channel)
//   out:  [B, OC, 2*Hc-1, 2*Wc-1]
__global__ void interleave_bias_tanh_kernel(
    const float* __restrict__ conv,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int OC, int Hc, int Wc, int Ho, int Wo, int xPairs)
{
    int xp = blockIdx.x * blockDim.x + threadIdx.x;  // pair index along width
    int y  = blockIdx.y * blockDim.y + threadIdx.y;  // output row
    if (xp >= xPairs || y >= Ho) return;

    int oc = blockIdx.z % OC;
    int b  = blockIdx.z / OC;

    int ry = y & 1;
    int sy = (y >> 1) + ry;

    long long base0 = (((long long)b * (OC * 4) + (oc * 4 + ry * 2 + 0)) * Hc + sy) * Wc + xp;
    float bv = __ldg(bias + oc);

    long long outOff = (((long long)b * OC + oc) * Ho + y) * Wo + 2 * xp;
    out[outOff] = tanhf(__ldg(conv + base0) - bv);

    if (2 * xp + 1 < Wo) {
        long long base1 = (((long long)b * (OC * 4) + (oc * 4 + ry * 2 + 1)) * Hc + sy) * Wc + xp + 1;
        out[outOff + 1] = tanhf(__ldg(conv + base1) - bv);
    }
}

// Fused bias-subtract + tanh for the generic fallback path.
__global__ void bias_tanh_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    float* __restrict__ out,
    long long total, int C, int HW)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (int)((idx / HW) % C);
    out[idx] = tanhf(x[idx] - __ldg(bias + c));
}

torch::Tensor interleave_bias_tanh(torch::Tensor conv, torch::Tensor bias)
{
    TORCH_CHECK(conv.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(conv.scalar_type() == at::kFloat, "conv must be float32");

    auto conv_c = conv.contiguous();
    auto bias_c = bias.contiguous();

    int64_t B  = conv_c.size(0);
    int64_t C4 = conv_c.size(1);
    int64_t Hc = conv_c.size(2);
    int64_t Wc = conv_c.size(3);
    int64_t OC = C4 / 4;
    int64_t Ho = 2 * Hc - 1;
    int64_t Wo = 2 * Wc - 1;

    auto out = at::empty({B, OC, Ho, Wo}, conv_c.options());

    int64_t xPairs = (Wo + 1) / 2;

    dim3 block(128, 8);
    dim3 grid((unsigned int)((xPairs + block.x - 1) / block.x),
              (unsigned int)((Ho + block.y - 1) / block.y),
              (unsigned int)(B * OC));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    interleave_bias_tanh_kernel<<<grid, block, 0, stream>>>(
        conv_c.data_ptr<float>(), bias_c.data_ptr<float>(), out.data_ptr<float>(),
        (int)OC, (int)Hc, (int)Wc, (int)Ho, (int)Wo, (int)xPairs);

    return out;
}

torch::Tensor bias_tanh(torch::Tensor x, torch::Tensor bias)
{
    TORCH_CHECK(x.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    auto x_c = x.contiguous();
    auto bias_c = bias.contiguous();
    auto out = at::empty_like(x_c);

    long long total = (long long)x_c.numel();
    int C  = (int)x_c.size(1);
    int HW = (int)(x_c.size(2) * x_c.size(3));

    int threads = 256;
    long long blocks = (total + threads - 1) / threads;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    bias_tanh_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        x_c.data_ptr<float>(), bias_c.data_ptr<float>(), out.data_ptr<float>(),
        total, C, HW);

    return out;
}
"""

_fused_cpp_src = (
    "torch::Tensor interleave_bias_tanh(torch::Tensor conv, torch::Tensor bias);\n"
    "torch::Tensor bias_tanh(torch::Tensor x, torch::Tensor bias);"
)

_fused_ops = load_inline(
    name="fused_subpixel_deconv_epilogue",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["interleave_bias_tanh", "bias_tanh"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)

torch.backends.cudnn.benchmark = True


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - ConvTranspose2d(k=4, s=2, p=1, op=1) is computed as an equivalent
        2x2 Conv2d with 4*out_channels (sub-pixel decomposition, exact math),
        which uses the faster forward-convolution kernels.
      - The phase interleave (pixel shuffle), bias subtraction and tanh are
        fused into a single custom CUDA kernel (one read + one write pass).
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape,
                 stride=2, padding=1, output_padding=1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, output_padding=output_padding)
        self.bias = nn.Parameter(torch.randn(bias_shape))

        self.in_channels = in_channels
        self.out_channels = out_channels

        k = kernel_size[0] if isinstance(kernel_size, (tuple, list)) else kernel_size
        s = stride[0] if isinstance(stride, (tuple, list)) else stride
        p = padding[0] if isinstance(padding, (tuple, list)) else padding
        op = output_padding[0] if isinstance(output_padding, (tuple, list)) else output_padding
        self._fast_path = (k == 4 and s == 2 and p == 1 and op == 1)

    def forward(self, x):
        if self._fast_path:
            w = self.conv_transpose.weight  # [IC, OC, 4, 4]
            IC, OC = self.in_channels, self.out_channels
            # W_eq[(oc,ry,rx), ic, dy, dx] = W_t[ic, oc, 3-2dy-ry, 3-2dx-rx]
            w6 = w.flip([2, 3]).reshape(IC, OC, 2, 2, 2, 2)  # [ic, oc, dy, ry, dx, rx]
            w_eq = w6.permute(1, 3, 5, 0, 2, 4).reshape(OC * 4, IC, 2, 2).contiguous()

            y = F.conv2d(x, w_eq, padding=1)  # [B, 4*OC, Hin+1, Win+1]

            # Fold the ConvTranspose2d bias into the subtraction bias:
            # out = tanh(conv + conv_bias - bias) = tanh(conv - eff_bias)
            if self.conv_transpose.bias is not None:
                eff_bias = self.bias.view(-1) - self.conv_transpose.bias
            else:
                eff_bias = self.bias.view(-1)
            return _fused_ops.interleave_bias_tanh(y, eff_bias)

        # Generic fallback for other configurations.
        y = self.conv_transpose(x)
        return _fused_ops.bias_tanh(y, self.bias.view(-1))
```