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

# Fused epilogue kernel: out = tanh(conv_transpose_out + conv_bias - bias)
# This replaces three separate GPU passes (cuDNN bias add, the bias subtraction,
# and the tanh) with a single memory-bound pass over the activation tensor.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

__global__ void fused_bias_sub_tanh_kernel(const float* __restrict__ x,
                                           float* __restrict__ out,
                                           const float* __restrict__ conv_bias,
                                           const float* __restrict__ sub_bias,
                                           const int plane_size,
                                           const int channels) {
    // blockIdx.y indexes the (batch * channel) plane
    const int bc = blockIdx.y;
    const int c = bc % channels;
    const float b = __ldg(conv_bias + c) - __ldg(sub_bias + c);

    const long base = (long)bc * (long)plane_size;
    const float* xp = x + base;
    float* op = out + base;

    // xp and op have identical alignment; compute leading scalars until 16B
    const int mis = (int)((size_t)xp & 15);
    const int head = ((16 - mis) & 15) >> 2;  // 0..3

    const int tid = threadIdx.x;
    const int gtid = blockIdx.x * blockDim.x + tid;

    // Handle unaligned head and the tail with scalar accesses (block 0 only)
    if (blockIdx.x == 0) {
        for (int i = tid; i < head; i += blockDim.x) {
            op[i] = tanhf(xp[i] + b);
        }
        const int nvec0 = (plane_size - head) >> 2;
        const int tail_start = head + (nvec0 << 2);
        for (int i = tail_start + tid; i < plane_size; i += blockDim.x) {
            op[i] = tanhf(xp[i] + b);
        }
    }

    // Vectorized body with float4
    const int nvec = (plane_size - head) >> 2;
    const float4* __restrict__ x4 = reinterpret_cast<const float4*>(xp + head);
    float4* __restrict__ o4 = reinterpret_cast<float4*>(op + head);
    const int stride = gridDim.x * blockDim.x;
    for (int i = gtid; i < nvec; i += stride) {
        float4 v = x4[i];
        v.x = tanhf(v.x + b);
        v.y = tanhf(v.y + b);
        v.z = tanhf(v.z + b);
        v.w = tanhf(v.w + b);
        o4[i] = v;
    }
}

torch::Tensor fused_bias_sub_tanh_cuda(torch::Tensor x,
                                       torch::Tensor conv_bias,
                                       torch::Tensor sub_bias) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(conv_bias.numel() == x.size(1), "conv_bias size mismatch");
    TORCH_CHECK(sub_bias.numel() == x.size(1), "sub_bias size mismatch");

    const c10::cuda::OptionalCUDAGuard guard(x.device());

    auto out = at::empty_like(x);
    const long HW = (long)x.size(2) * x.size(3);
    const long BC = (long)x.size(0) * x.size(1);
    const int plane_size = (int)HW;
    const int channels = (int)x.size(1);

    const int threads = 256;
    const long nvec_max = (HW + 3) >> 2;
    long blocks_x = (nvec_max + threads - 1) / threads;
    if (blocks_x < 1) blocks_x = 1;
    dim3 grid((unsigned)blocks_x, (unsigned)BC);

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_bias_sub_tanh_kernel<<<grid, threads, 0, stream>>>(
        x.data_ptr<float>(),
        out.data_ptr<float>(),
        conv_bias.data_ptr<float>(),
        sub_bias.data_ptr<float>(),
        plane_size,
        channels);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_bias_sub_tanh_cuda(torch::Tensor x, torch::Tensor conv_bias, torch::Tensor sub_bias);"
)

# Compile the inline CUDA epilogue
fused_epilogue = load_inline(
    name="fused_bias_sub_tanh",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_bias_sub_tanh_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Strategy:
      - Keep the transposed convolution on the highly-optimized cuDNN path
        (with cudnn autotuning enabled), but launch it WITHOUT its bias so no
        extra bias-add pass is performed.
      - Fuse the remaining operators (conv-bias add, parameter subtraction,
        and tanh activation) into a single custom CUDA epilogue kernel that
        makes exactly one read + one write pass over the activation tensor
        using vectorized float4 accesses.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape,
                 stride=2, padding=1, output_padding=1):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        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.stride = stride
        self.padding = padding
        self.output_padding = output_padding
        self.epilogue = fused_epilogue

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            # cuDNN transposed conv without bias (bias handled in the epilogue)
            out = F.conv_transpose2d(
                x, self.conv_transpose.weight, None,
                self.stride, self.padding, self.output_padding)
            # Single fused pass: tanh(out + conv_bias - bias)
            return self.epilogue.fused_bias_sub_tanh_cuda(
                out.contiguous(),
                self.conv_transpose.bias,
                self.bias.view(-1))
        # Fallback (reference path)
        x = self.conv_transpose(x)
        x = x - self.bias
        return torch.tanh(x)
```