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

# Fused CUDA kernel: out = tanh(x - bias[c]) for an NCHW tensor with a
# per-channel bias of shape (C, 1, 1). This replaces two full memory passes
# (bias subtraction + tanh) with a single pass over the conv output.
fused_sub_bias_tanh_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

// Vectorized kernel: grid.y indexes the (N*C) planes so the channel id is
// obtained without per-element integer division. The bulk of each plane is
// processed with 16-byte aligned float4 loads/stores.
__global__ void sub_bias_tanh_vec_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int C,
    const int HW)
{
    const int plane = blockIdx.y;
    const float b = __ldg(bias + (plane % C));
    const size_t base = (size_t)plane * (size_t)HW;
    const float* xp = x + base;
    float* op = out + base;

    // Number of leading scalar elements until xp is 16B aligned.
    const uintptr_t addr = reinterpret_cast<uintptr_t>(xp);
    int head = (int)(((16 - (addr & 15)) & 15) >> 2);
    if (head > HW) head = HW;
    const int nvec = (HW - head) >> 2;

    const float4* __restrict__ xv = reinterpret_cast<const float4*>(xp + head);
    float4* __restrict__ ov = reinterpret_cast<float4*>(op + head);

    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = gridDim.x * blockDim.x;
    for (int i = tid; i < nvec; i += stride) {
        float4 v = xv[i];
        float4 r;
        r.x = tanhf(v.x - b);
        r.y = tanhf(v.y - b);
        r.z = tanhf(v.z - b);
        r.w = tanhf(v.w - b);
        ov[i] = r;
    }

    // Scalar peel for the unaligned head and the tail (handled by block 0).
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const int i = threadIdx.x;
            op[i] = tanhf(xp[i] - b);
        }
        const int tailStart = head + (nvec << 2);
        const int tailCount = HW - tailStart;
        if (threadIdx.x >= 128 && threadIdx.x < 128 + tailCount) {
            const int i = tailStart + (threadIdx.x - 128);
            op[i] = tanhf(xp[i] - b);
        }
    }
}

// Scalar fallback for tiny planes.
__global__ void sub_bias_tanh_scalar_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int C,
    const int HW)
{
    const int plane = blockIdx.y;
    const float b = __ldg(bias + (plane % C));
    const size_t base = (size_t)plane * (size_t)HW;
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < HW) {
        out[base + i] = tanhf(x[base + i] - b);
    }
}

torch::Tensor sub_bias_tanh_cuda(torch::Tensor x, torch::Tensor bias)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() >= 2, "x must have at least 2 dims");

    auto xc = x.contiguous();
    auto bc = bias.contiguous();
    auto out = torch::empty_like(xc);

    const int C = (int)xc.size(1);
    int64_t hw = 1;
    for (int i = 2; i < xc.dim(); ++i) hw *= xc.size(i);
    const int HW = (int)hw;
    const int64_t planes = xc.numel() / hw;
    TORCH_CHECK(planes <= 65535, "too many planes");

    const int block = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (HW >= 8) {
        const int nvecMax = (HW + 3) / 4;
        int gx = (nvecMax + block - 1) / block;
        if (gx > 65535) gx = 65535;  // grid-stride loop covers the rest
        dim3 grid(gx, (unsigned)planes);
        sub_bias_tanh_vec_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(),
            out.data_ptr<float>(), C, HW);
    } else {
        dim3 grid((HW + block - 1) / block, (unsigned)planes);
        sub_bias_tanh_scalar_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(),
            out.data_ptr<float>(), C, HW);
    }
    return out;
}
"""

fused_sub_bias_tanh_cpp_source = (
    "torch::Tensor sub_bias_tanh_cuda(torch::Tensor x, torch::Tensor bias);"
)

fused_sub_bias_tanh = load_inline(
    name="fused_sub_bias_tanh",
    cpp_sources=fused_sub_bias_tanh_cpp_source,
    cuda_sources=fused_sub_bias_tanh_cuda_source,
    functions=["sub_bias_tanh_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN transposed convolution and replaces the
    separate bias-subtraction and tanh passes with a single fused custom CUDA
    kernel (one memory pass instead of two).
    """
    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.fused_op = fused_sub_bias_tanh

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.fused_op.sub_bias_tanh_cuda(x, self.bias)
        return x
```