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

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

// ---------------------------------------------------------------------------
// Fused ConvTranspose2d(k=4,s=2,p=1,op=1) + bias-subtract + tanh.
//
// The transposed convolution is decomposed into 4 parity "phases"
// (po = oh & 1, pq = ow & 1).  For each phase the operation is an exact
// GEMM over output pixels M = (n, oh2, ow2) with K = (ci, a, b):
//
//   out = sum_{ci,a,b} x[n, ci, oh2+po-a, ow2+pq-b] * w[ci, co, 2a+1-po, 2b+1-pq]
//
// computed with a tiled fp32 implicit-GEMM kernel (shared-memory staged,
// register 4x4 micro-tiles), with (+conv_bias, -bias, tanh) fused epilogue.
// ---------------------------------------------------------------------------

#define TILE_M 64
#define TILE_N 64
#define TILE_K 32
#define NTHREADS 256

// wp[phase][k=(ci,a,b)][co]  with  k = ci*4 + a*2 + b
__global__ void rearrange_weight_kernel(const float* __restrict__ w,
                                        float* __restrict__ wp) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;   // 4*256*64 elements
    if (idx >= 4 * 256 * 64) return;
    int co = idx & 63;
    int k  = (idx >> 6) & 255;
    int p  = idx >> 14;
    int po = p >> 1, pq = p & 1;
    int ci = k >> 2;
    int a  = (k >> 1) & 1;
    int b  = k & 1;
    int kh = 2 * a + 1 - po;
    int kw = 2 * b + 1 - pq;
    wp[idx] = w[(long)ci * 1024 + co * 16 + kh * 4 + kw];
}

__global__ void deconv_fused_kernel(const float* __restrict__ x,
                                    const float* __restrict__ wp,
                                    const float* __restrict__ conv_bias,
                                    const float* __restrict__ sub_bias,
                                    float* __restrict__ y,
                                    int N, int H, int W, int Ho, int Wo,
                                    int has_cb) {
    __shared__ float As[TILE_K][TILE_M];
    __shared__ float Bs[TILE_K][TILE_N];
    __shared__ int soff[TILE_M][4];
    __shared__ long long sbase[TILE_M];
    __shared__ long long sorow[TILE_M];

    const int phase = blockIdx.z;
    const int po = phase >> 1, pq = phase & 1;
    const int Hp = (Ho - po + 1) >> 1;
    const int Wp2 = (Wo - pq + 1) >> 1;
    const long M = (long)N * Hp * Wp2;
    const long row0 = (long)blockIdx.x * TILE_M;
    if (row0 >= M) return;

    const int tid = threadIdx.x;
    const long xsc = (long)H * W;
    const long xsn = 64 * xsc;
    const long ytc = (long)Ho * Wo;
    const long ytn = 64 * ytc;

    if (tid < TILE_M) {
        long row = row0 + tid;
        if (row < M) {
            long hw = (long)Hp * Wp2;
            long n = row / hw;
            long rem = row - n * hw;
            int oh2 = (int)(rem / Wp2);
            int ow2 = (int)(rem - (long)oh2 * Wp2);
            sbase[tid] = (long long)n * xsn;
            #pragma unroll
            for (int ab = 0; ab < 4; ++ab) {
                int a = ab >> 1, b = ab & 1;
                int ih = oh2 + po - a;
                int iw = ow2 + pq - b;
                bool v = (ih >= 0 && ih < H && iw >= 0 && iw < W);
                soff[tid][ab] = v ? (ih * W + iw) : -1;
            }
            int oh = 2 * oh2 + po;
            int ow = 2 * ow2 + pq;
            sorow[tid] = (long long)n * ytn + (long)oh * Wo + ow;
        } else {
            sbase[tid] = 0;
            #pragma unroll
            for (int ab = 0; ab < 4; ++ab) soff[tid][ab] = -1;
            sorow[tid] = -1;
        }
    }
    __syncthreads();

    const int tx = tid & 15;   // M dimension (4 rows each)
    const int ty = tid >> 4;   // N dimension (4 cols each)
    const int m0 = tx * 4;
    const int c0 = ty * 4;

    float acc[4][4];
    #pragma unroll
    for (int r = 0; r < 4; ++r)
        #pragma unroll
        for (int c = 0; c < 4; ++c)
            acc[r][c] = 0.0f;

    const float* wphase = wp + (long)phase * (256 * 64);

    for (int kc = 0; kc < 256 / TILE_K; ++kc) {
        // stage A tile (gather from x)
        #pragma unroll
        for (int it = 0; it < (TILE_K * TILE_M) / NTHREADS; ++it) {
            int idx = tid + it * NTHREADS;
            int m  = idx & (TILE_M - 1);
            int kk = idx >> 6;                 // 0..31
            int kg = kc * TILE_K + kk;
            int ci = kg >> 2;
            int ab = kg & 3;
            int off = soff[m][ab];
            float v = 0.0f;
            if (off >= 0) v = __ldg(x + sbase[m] + (long)ci * xsc + off);
            As[kk][m] = v;
        }
        // stage B tile (rearranged weights, coalesced)
        #pragma unroll
        for (int it = 0; it < (TILE_K * TILE_N) / NTHREADS; ++it) {
            int idx = tid + it * NTHREADS;
            int co = idx & 63;
            int kk = idx >> 6;
            Bs[kk][co] = __ldg(wphase + (long)(kc * TILE_K + kk) * 64 + co);
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < TILE_K; ++kk) {
            float4 av = *reinterpret_cast<const float4*>(&As[kk][m0]);
            float4 bv = *reinterpret_cast<const float4*>(&Bs[kk][c0]);
            acc[0][0] += av.x * bv.x; acc[0][1] += av.x * bv.y;
            acc[0][2] += av.x * bv.z; acc[0][3] += av.x * bv.w;
            acc[1][0] += av.y * bv.x; acc[1][1] += av.y * bv.y;
            acc[1][2] += av.y * bv.z; acc[1][3] += av.y * bv.w;
            acc[2][0] += av.z * bv.x; acc[2][1] += av.z * bv.y;
            acc[2][2] += av.z * bv.z; acc[2][3] += av.z * bv.w;
            acc[3][0] += av.w * bv.x; acc[3][1] += av.w * bv.y;
            acc[3][2] += av.w * bv.z; acc[3][3] += av.w * bv.w;
        }
        __syncthreads();
    }

    // fused epilogue: + conv bias, - bias, tanh
    #pragma unroll
    for (int r = 0; r < 4; ++r) {
        int m = m0 + r;
        long long orow = sorow[m];
        if (orow < 0) continue;
        #pragma unroll
        for (int c = 0; c < 4; ++c) {
            int co = c0 + c;
            float v = acc[r][c];
            if (has_cb) v += __ldg(conv_bias + co);
            v -= __ldg(sub_bias + co);
            y[orow + (long long)co * ytc] = tanhf(v);
        }
    }
}

// Fallback fused elementwise: y = tanh(x - sub_bias) per channel
__global__ void fused_bias_tanh_kernel(const float* __restrict__ x,
                                       const float* __restrict__ sub_bias,
                                       float* __restrict__ y,
                                       long HW, long total, int CO) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int co = (int)((idx / HW) % CO);
    y[idx] = tanhf(x[idx] - __ldg(sub_bias + co));
}

torch::Tensor deconv_forward(torch::Tensor x, torch::Tensor weight,
                             c10::optional<torch::Tensor> conv_bias,
                             torch::Tensor sub_bias,
                             int64_t stride, int64_t padding,
                             int64_t output_padding) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32, "x must be CUDA fp32");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = sub_bias.contiguous();

    const int N = xc.size(0);
    const int CI = xc.size(1);
    const int H = xc.size(2);
    const int W = xc.size(3);
    const int CO = wc.size(1);
    const int KH = wc.size(2);
    const int KW = wc.size(3);

    const long Ho = (long)(H - 1) * stride - 2 * padding + KH + output_padding;
    const long Wo = (long)(W - 1) * stride - 2 * padding + KW + output_padding;

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

    torch::Tensor cb;
    const float* cb_ptr = nullptr;
    int has_cb = 0;
    if (conv_bias.has_value() && conv_bias->defined()) {
        cb = conv_bias->contiguous();
        cb_ptr = cb.data_ptr<float>();
        has_cb = 1;
    }

    auto y = torch::empty({N, CO, Ho, Wo}, xc.options());

    bool fast = (CI == 64) && (CO == 64) && (KH == 4) && (KW == 4) &&
                (stride == 2) && (padding == 1) && (output_padding == 1) &&
                (Ho <= INT_MAX) && (Wo <= INT_MAX);

    if (!fast) {
        // safe fallback: cuDNN/ATen conv (with its bias) + one fused kernel
        auto tmp = at::conv_transpose2d(xc, wc, conv_bias.has_value() ? cb : torch::Tensor(),
                                        {stride, stride}, {padding, padding},
                                        {output_padding, output_padding}, 1);
        long total = tmp.numel();
        long HW = Ho * Wo;
        int threads = 256;
        long blocks = (total + threads - 1) / threads;
        fused_bias_tanh_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            tmp.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(),
            HW, total, CO);
        return y;
    }

    // rearranged weights: 4 phases x 256 x 64
    auto wp = torch::empty({4 * 256 * 64}, xc.options());
    rearrange_weight_kernel<<<256, 256, 0, stream>>>(wc.data_ptr<float>(), wp.data_ptr<float>());

    long maxM = 0;
    for (int p = 0; p < 4; ++p) {
        int po = p >> 1, pq = p & 1;
        long Hp = (Ho - po + 1) >> 1;
        long Wp2 = (Wo - pq + 1) >> 1;
        long Mp = (long)N * Hp * Wp2;
        if (Mp > maxM) maxM = Mp;
    }
    long tiles = (maxM + TILE_M - 1) / TILE_M;

    dim3 grid((unsigned)tiles, 1, 4);
    deconv_fused_kernel<<<grid, NTHREADS, 0, stream>>>(
        xc.data_ptr<float>(), wp.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
        y.data_ptr<float>(), N, H, W, (int)Ho, (int)Wo, has_cb);
    return y;
}
"""

cpp_source = """
torch::Tensor deconv_forward(torch::Tensor x, torch::Tensor weight,
                             c10::optional<torch::Tensor> conv_bias,
                             torch::Tensor sub_bias,
                             int64_t stride, int64_t padding,
                             int64_t output_padding);
"""

deconv_mod = load_inline(
    name="deconv_fused_fp32",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["deconv_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: transposed convolution via a fused fp32 implicit-GEMM
    CUDA kernel with (+conv_bias, -bias, tanh) fused in the epilogue.
    Falls back to ATen conv_transpose2d + a fused elementwise kernel for
    non-standard configurations.
    """
    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.stride = stride
        self.padding = padding
        self.output_padding = output_padding

    def forward(self, x):
        return deconv_mod.deconv_forward(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            self.bias.view(-1),
            self.stride,
            self.padding,
            self.output_padding,
        )
```