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

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

using namespace nvcuda;

#define TILE_M 128
#define TILE_N 64
#define TILE_K 32
#define NTHREADS 256
#define K_CHUNKS 8

// Rearrange ConvTranspose2d weight (CI,CO,4,4) into per-phase GEMM-friendly
// layout Wp[phase][k=(ci,a,b)][co] with kh=2a+1-po, kw=2b+1-pq.
__global__ void rearrange_weight_kernel(const float* __restrict__ w,
                                        float* __restrict__ wp) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    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[ci * (64 * 16) + co * 16 + kh * 4 + kw];
}

// Implicit-GEMM transposed convolution (k=4, s=2, p=1, op=1) using TF32
// tensor cores, with fused (-bias, +conv_bias, tanh) epilogue.
__global__ void deconv_gemm_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) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    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 int M = N * Hp * Wp2;
    const int row0 = blockIdx.x * TILE_M;
    if (row0 >= M) return;

    extern __shared__ float smem[];
    float* As = smem;                                   // [2][TILE_K][TILE_M]
    float* Bs = As + 2 * TILE_K * TILE_M;               // [2][TILE_K][TILE_N]
    float* wbuf = As;                                   // reused in epilogue (32KB)
    int* soff  = (int*)(Bs + 2 * TILE_K * TILE_N);      // [TILE_M*4]
    int* sbase = soff + TILE_M * 4;                     // [TILE_M]
    int* sorow = sbase + TILE_M;                        // [TILE_M]

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

    // Per-row geometry precompute
    if (tid < TILE_M) {
        int row = row0 + tid;
        if (row < M) {
            int hw = Hp * Wp2;
            int n = row / hw;
            int rem = row - n * hw;
            int oh2 = rem / Wp2;
            int ow2 = rem - oh2 * Wp2;
            sbase[tid] = 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 * 4 + ab] = v ? (ih * W + iw) : -1;
            }
            int oh = 2 * oh2 + po;
            int ow = 2 * ow2 + pq;
            sorow[tid] = n * ytn + oh * Wo + ow;
        } else {
            sbase[tid] = 0;
            #pragma unroll
            for (int ab = 0; ab < 4; ++ab) soff[tid * 4 + ab] = -1;
            sorow[tid] = -1;
        }
    }
    __syncthreads();

    const int wid  = tid >> 5;
    const int lane = tid & 31;
    const int wm = (wid & 3) * 32;
    const int wn = (wid >> 2) * 32;

    wmma::fragment<wmma::accumulator, 16, 8, 8, float> acc[2][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

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

    // ---- prologue: load chunk 0 into buffer 0 ----
    {
        float* Asb = As;
        #pragma unroll
        for (int it = 0; it < (TILE_M * TILE_K) / NTHREADS; ++it) {
            int idx = tid + it * NTHREADS;
            int m = idx & (TILE_M - 1);
            int kk = idx >> 7;
            int ci = kk >> 2;
            int ab = kk & 3;
            int off = soff[m * 4 + ab];
            float v = 0.f;
            if (off >= 0) v = __ldg(x + sbase[m] + ci * xsc + off);
            Asb[kk * TILE_M + m] = wmma::__float_to_tf32(v);
        }
        float* Bsb = Bs;
        #pragma unroll
        for (int it = 0; it < (TILE_K * TILE_N) / NTHREADS; ++it) {
            int idx = tid + it * NTHREADS;
            Bsb[idx] = wmma::__float_to_tf32(__ldg(wphase + idx));
        }
    }
    __syncthreads();

    // ---- main K loop with double buffering ----
    #pragma unroll
    for (int kc = 0; kc < K_CHUNKS; ++kc) {
        int buf = kc & 1;
        if (kc + 1 < K_CHUNKS) {
            int nb = buf ^ 1;
            float* Asb = As + nb * (TILE_K * TILE_M);
            #pragma unroll
            for (int it = 0; it < (TILE_M * TILE_K) / NTHREADS; ++it) {
                int idx = tid + it * NTHREADS;
                int m = idx & (TILE_M - 1);
                int kk = idx >> 7;
                int ci = kk >> 2;
                int ab = kk & 3;
                int off = soff[m * 4 + ab];
                float v = 0.f;
                if (off >= 0) v = __ldg(x + sbase[m] + ci * xsc + off);
                Asb[kk * TILE_M + m] = wmma::__float_to_tf32(v);
            }
            float* Bsb = Bs + nb * (TILE_K * TILE_N);
            const float* wpp = wphase + (kc + 1) * (TILE_K * TILE_N);
            #pragma unroll
            for (int it = 0; it < (TILE_K * TILE_N) / NTHREADS; ++it) {
                int idx = tid + it * NTHREADS;
                Bsb[idx] = wmma::__float_to_tf32(__ldg(wpp + idx));
            }
        }
        {
            const float* Asb = As + buf * (TILE_K * TILE_M);
            const float* Bsb = Bs + buf * (TILE_K * TILE_N);
            #pragma unroll
            for (int ks = 0; ks < TILE_K / 8; ++ks) {
                wmma::fragment<wmma::matrix_a, 16, 8, 8, wmma::precision::tf32, wmma::col_major> af[2];
                wmma::fragment<wmma::matrix_b, 16, 8, 8, wmma::precision::tf32, wmma::row_major> bf[4];
                #pragma unroll
                for (int i = 0; i < 2; ++i)
                    wmma::load_matrix_sync(af[i], Asb + (ks * 8) * TILE_M + wm + i * 16, TILE_M);
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    wmma::load_matrix_sync(bf[j], Bsb + (ks * 8) * TILE_N + wn + j * 8, TILE_N);
                #pragma unroll
                for (int i = 0; i < 2; ++i)
                    #pragma unroll
                    for (int j = 0; j < 4; ++j)
                        wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
            }
        }
        __syncthreads();
    }

    // ---- epilogue: stage through smem, fused bias/sub/tanh ----
    float* my = wbuf + wid * 1024;   // [32(co)][32(m)] col-major staging
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::store_matrix_sync(my + (j * 8) * 32 + i * 16, acc[i][j], 32, wmma::mem_col_major);
    __syncwarp();

    int m = wm + lane;
    int orow = sorow[m];
    if (orow >= 0) {
        #pragma unroll 4
        for (int c = 0; c < 32; ++c) {
            int co = wn + c;
            float v = my[c * 32 + lane];
            if (has_cb) v += __ldg(conv_bias + co);
            v -= __ldg(sub_bias + co);
            y[orow + co * ytc] = tanhf(v);
        }
    }
#endif
}

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

static bool g_smem_configured = false;

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 int Ho = (H - 1) * (int)stride - 2 * (int)padding + KH + (int)output_padding;
    const int Wo = (W - 1) * (int)stride - 2 * (int)padding + KW + (int)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;
    }

    // check fast-path feasibility
    int dev = 0; cudaGetDevice(&dev);
    static int arch_major = -1;
    if (arch_major < 0) {
        cudaDeviceGetAttribute(&arch_major, cudaDevAttrComputeCapabilityMajor, dev);
    }
    bool fast = (arch_major >= 8) && CI == 64 && CO == 64 && KH == 4 && KW == 4 &&
                stride == 2 && padding == 1 && output_padding == 1;

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

    if (!fast) {
        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 = (long)Ho * Wo;
        int threads = 256;
        long blocks = (total + threads - 1) / threads;
        fused_bias_tanh_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            tmp.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
            y.data_ptr<float>(), HW, total, has_cb);
        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>());

    // tiles per phase (max over phases)
    int maxM = 0;
    for (int p = 0; p < 4; ++p) {
        int po = p >> 1, pq = p & 1;
        int Hp = (Ho - po + 1) >> 1;
        int Wp2 = (Wo - pq + 1) >> 1;
        int M = N * Hp * Wp2;
        if (M > maxM) maxM = M;
    }
    int tiles = (maxM + TILE_M - 1) / TILE_M;

    size_t smem = (size_t)(2 * TILE_K * TILE_M + 2 * TILE_K * TILE_N) * sizeof(float)
                + (size_t)(TILE_M * 6) * sizeof(int);
    if (!g_smem_configured) {
        cudaFuncSetAttribute(deconv_gemm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
        g_smem_configured = true;
    }
    dim3 grid(tiles, 1, 4);
    deconv_gemm_kernel<<<grid, NTHREADS, smem, stream>>>(
        xc.data_ptr<float>(), wp.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
        y.data_ptr<float>(), N, H, W, Ho, 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",
    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 TF32 tensor-core implicit-GEMM
    CUDA kernel with fused (- bias + conv_bias, tanh) epilogue.
    """
    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,
        )
```