```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>
#include <mma.h>

using namespace nvcuda;

// ---------------------------------------------------------------------------
// 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 TF32 tensor-core (WMMA) implicit-GEMM kernel with a
// two-stage shared-memory software pipeline, with (+conv_bias, -bias, tanh)
// fused in the epilogue.
// ---------------------------------------------------------------------------

#define TILE_M 128
#define TILE_N 64
#define SLICE_K 16
#define NTHREADS 256
#define LDA (TILE_M + 8)
#define LDB (TILE_N + 8)
#define LDC (TILE_M + 8)

__device__ __forceinline__ float to_tf32(float x) {
    // round-to-nearest (ties up) truncation to tf32 precision
    unsigned int u = __float_as_uint(x);
    u = (u + 0x1000u) & ~0x1FFFu;
    return __uint_as_float(u);
}

// wp[phase][k=(ci,a,b)][co]  with  k = ci*4 + a*2 + b  (tf32-rounded)
// also computes combined bias comb[co] = conv_bias[co] - sub_bias[co]
__global__ void rearrange_weight_kernel(const float* __restrict__ w,
                                        const float* __restrict__ conv_bias,
                                        const float* __restrict__ sub_bias,
                                        float* __restrict__ wp,
                                        float* __restrict__ comb,
                                        int has_cb) {
    if (blockIdx.x == 256) {
        int co = threadIdx.x;
        if (co < 64) {
            float cb = has_cb ? __ldg(conv_bias + co) : 0.0f;
            comb[co] = cb - __ldg(sub_bias + co);
        }
        return;
    }
    int idx = blockIdx.x * blockDim.x + threadIdx.x;   // 4*256*64 elements
    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] = to_tf32(w[(long)ci * 1024 + co * 16 + kh * 4 + kw]);
}

__global__ void __launch_bounds__(NTHREADS) deconv_fused_tc_kernel(
        const float* __restrict__ x,
        const float* __restrict__ wp,
        const float* __restrict__ comb,
        float* __restrict__ y,
        int N, int H, int W, int Ho, int Wo) {

    __shared__ union {
        struct {
            float As[2][SLICE_K][LDA];
            float Bs[2][SLICE_K][LDB];
        } ab;
        float Cs[TILE_N][LDC];   // transposed accumulator tile [co][m]
    } smem;
    __shared__ int soff[TILE_M][4];
    __shared__ long long sbase[TILE_M];
    __shared__ long long sorow[TILE_M];
    __shared__ float scomb[TILE_N];

    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;
        }
    }
    if (tid < TILE_N) scomb[tid] = __ldg(comb + tid);
    __syncthreads();

    const int warp = tid >> 5;
    const int wm = warp & 3;    // 4 warps along M (32 rows each)
    const int wn = warp >> 2;   // 2 warps along N (32 cols each)
    const int m_off = wm * 32;
    const int n_off = wn * 32;

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

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

    // stage k-slice `ks` (16 k-values) into buffer `buf`
    auto stage = [&](int ks, int buf) {
        int kbase = ks * SLICE_K;
        #pragma unroll
        for (int it = 0; it < (SLICE_K * TILE_M) / NTHREADS; ++it) {  // 8
            int idx = tid + it * NTHREADS;
            int m  = idx & (TILE_M - 1);
            int kk = idx >> 7;
            int kg = kbase + 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);
            smem.ab.As[buf][kk][m] = to_tf32(v);
        }
        #pragma unroll
        for (int it = 0; it < (SLICE_K * TILE_N) / NTHREADS; ++it) {  // 4
            int idx = tid + it * NTHREADS;
            int co = idx & 63;
            int kk = idx >> 6;
            smem.ab.Bs[buf][kk][co] = __ldg(wphase + (long)(kbase + kk) * 64 + co);
        }
    };

    stage(0, 0);
    __syncthreads();

    #pragma unroll
    for (int ks = 0; ks < 256 / SLICE_K; ++ks) {
        int buf = ks & 1;
        if (ks + 1 < 256 / SLICE_K) stage(ks + 1, buf ^ 1);
        #pragma unroll
        for (int kk = 0; kk < SLICE_K / 8; ++kk) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::col_major> afrag[2];
            wmma::load_matrix_sync(afrag[0], &smem.ab.As[buf][kk * 8][m_off], LDA);
            wmma::load_matrix_sync(afrag[1], &smem.ab.As[buf][kk * 8][m_off + 16], LDA);
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> bfrag;
                wmma::load_matrix_sync(bfrag, &smem.ab.Bs[buf][kk * 8][n_off + j * 16], LDB);
                wmma::mma_sync(cfrag[0][j], afrag[0], bfrag, cfrag[0][j]);
                wmma::mma_sync(cfrag[1][j], afrag[1], bfrag, cfrag[1][j]);
            }
        }
        __syncthreads();
    }

    // write accumulators to shared memory (transposed: [co][m])
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::store_matrix_sync(&smem.Cs[n_off + j * 16][m_off + i * 16],
                                    cfrag[i][j], LDC, wmma::mem_col_major);
    __syncthreads();

    // fused epilogue: + combined bias, tanh, write out (m-fastest for coalescing)
    const int me = tid & (TILE_M - 1);
    const int cg = tid >> 7;   // 0 or 1
    const long long orow = sorow[me];
    if (orow >= 0) {
        #pragma unroll
        for (int it = 0; it < (TILE_M * TILE_N) / NTHREADS; ++it) {  // 32
            int co = 2 * it + cg;
            float v = smem.Cs[co][me] + scomb[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 (tf32-rounded) + combined bias
    auto wp = torch::empty({4 * 256 * 64}, xc.options());
    auto comb = torch::empty({64}, xc.options());
    rearrange_weight_kernel<<<257, 256, 0, stream>>>(
        wc.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
        wp.data_ptr<float>(), comb.data_ptr<float>(), has_cb);

    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_tc_kernel<<<grid, NTHREADS, 0, stream>>>(
        xc.data_ptr<float>(), wp.data_ptr<float>(), comb.data_ptr<float>(),
        y.data_ptr<float>(), N, H, W, (int)Ho, (int)Wo);
    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_tf32_tc",
    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 TF32 tensor-core
    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,
        )
```