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

# ---------------------------------------------------------------------------
# Custom fused ConvTranspose2d (stride=1, dilation=1, groups=1) implemented as
# an implicit-GEMM using TF32 tensor cores (WMMA) with a generic FP32 fallback.
#
# Math: ConvTranspose2d with stride 1 / dilation 1 / groups 1 is equivalent to
#   Y[n,oc,oh,ow] = sum_{ic,kh,kw} X[n,ic,oh+kh-(KH-1-pad_h), ow+kw-(KW-1-pad_w)]
#                                 * W[ic,oc,KH-1-kh,KW-1-kw]
# i.e. an implicit GEMM over m=(n,oh,ow), n=oc, k=(ic,kh,kw).
# ---------------------------------------------------------------------------

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

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 32
#define WARP_M 64
#define WARP_N 32
#define WM_FRAGS 4
#define WN_FRAGS 2

__global__ void __launch_bounds__(256)
conv_transpose_wmma_kernel(const float* __restrict__ X,
                           const float* __restrict__ W,
                           const float* __restrict__ bias,
                           float* __restrict__ Y,
                           int batch, int IC, int OC,
                           int IH, int IW, int OH, int OW,
                           int KH, int KW, int padH, int padW) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    extern __shared__ float smem[];
    float* As = smem;                 // BK x BM, [k][m] col-major A
    float* Bs = smem + BK * BM;       // BK x BN, [k][n] row-major B

    const int block_m = blockIdx.x * BM;
    const int block_n = blockIdx.y * BN;

    const int M = batch * OH * OW;
    const int KWH = KH * KW;
    const int K = IC * KWH;

    const int tid = threadIdx.x;
    const int warp_id = tid >> 5;
    const int warp_m = warp_id & 1;
    const int warp_n = warp_id >> 1;

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

    for (int k0 = 0; k0 < K; k0 += BK) {
        // ---- load A tile (gather from input with implicit zero padding) ----
        #pragma unroll
        for (int it = 0; it < (BK * BM) / 256; ++it) {
            int idx = tid + it * 256;
            int k_local = idx / BM;
            int m_local = idx % BM;
            int k = k0 + k_local;
            int m = block_m + m_local;
            float v = 0.0f;
            if (k < K && m < M) {
                int ic = k / KWH;
                int r  = k - ic * KWH;
                int kh = r / KW;
                int kw = r - kh * KW;
                int ow = m % OW;
                int t  = m / OW;
                int oh = t % OH;
                int n  = t / OH;
                int ih = oh + kh - padH;
                int iw = ow + kw - padW;
                if (ih >= 0 && ih < IH && iw >= 0 && iw < IW)
                    v = X[((n * IC + ic) * IH + ih) * IW + iw];
            }
            As[idx] = wmma::__float_to_tf32(v);
        }
        // ---- load B tile (flipped weights, contiguous writes) ----
        #pragma unroll
        for (int it = 0; it < (BK * BN) / 256; ++it) {
            int idx = tid + it * 256;
            int k_local = idx / BN;
            int n_local = idx % BN;
            int k = k0 + k_local;
            int oc = block_n + n_local;
            float v = 0.0f;
            if (k < K && oc < OC) {
                int ic = k / KWH;
                int r  = k - ic * KWH;
                int kh = r / KW;
                int kw = r - kh * KW;
                v = W[((ic * OC + oc) * KH + (KH - 1 - kh)) * KW + (KW - 1 - kw)];
            }
            Bs[idx] = wmma::__float_to_tf32(v);
        }
        __syncthreads();

        // ---- tensor core compute on the staged tiles ----
        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::col_major> a_frag[WM_FRAGS];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag[WN_FRAGS];
            #pragma unroll
            for (int i = 0; i < WM_FRAGS; ++i)
                wmma::load_matrix_sync(a_frag[i], &As[kk * BM + warp_m * WARP_M + i * 16], BM);
            #pragma unroll
            for (int j = 0; j < WN_FRAGS; ++j)
                wmma::load_matrix_sync(b_frag[j], &Bs[kk * BN + warp_n * WARP_N + j * 16], BN);
            #pragma unroll
            for (int i = 0; i < WM_FRAGS; ++i)
                #pragma unroll
                for (int j = 0; j < WN_FRAGS; ++j)
                    wmma::mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]);
        }
        __syncthreads();
    }

    // ---- stage accumulators in shared, then write out coalesced ----
    float* Ys = smem;  // reuse: BM x BN
    #pragma unroll
    for (int i = 0; i < WM_FRAGS; ++i)
        #pragma unroll
        for (int j = 0; j < WN_FRAGS; ++j)
            wmma::store_matrix_sync(&Ys[(warp_m * WARP_M + i * 16) * BN + warp_n * WARP_N + j * 16],
                                    acc[i][j], BN, wmma::mem_row_major);
    __syncthreads();

    #pragma unroll
    for (int it = 0; it < (BM * BN) / 256; ++it) {
        int idx = tid + it * 256;
        int m_local = idx / BN;
        int n_local = idx % BN;
        int m = block_m + m_local;
        int oc = block_n + n_local;
        if (m < M && oc < OC) {
            float v = Ys[idx];
            if (bias != nullptr) v += bias[oc];
            int ow = m % OW;
            int t  = m / OW;
            int oh = t % OH;
            int n  = t / OH;
            Y[((n * OC + oc) * OH + oh) * OW + ow] = v;
        }
    }
#endif
}

torch::Tensor conv_transpose2d_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                                    int64_t padH, int64_t padW) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only float32 is supported");
    const at::cuda::CUDAGuard guard(x.device());

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int batch = xc.size(0);
    const int IC = xc.size(1);
    const int IH = xc.size(2);
    const int IW = xc.size(3);
    const int OC = wc.size(1);
    const int KH = wc.size(2);
    const int KW = wc.size(3);
    const int OH = IH + KH - 1 - 2 * (int)padH;
    const int OW = IW + KW - 1 - 2 * (int)padW;

    auto Y = at::empty({batch, OC, OH, OW}, xc.options());

    const int M = batch * OH * OW;
    dim3 grid((M + BM - 1) / BM, (OC + BN - 1) / BN);
    const int shbytes = BM * BN * (int)sizeof(float);  // 64 KB (max of tiles/staging)

    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(conv_transpose_wmma_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, shbytes);
        attr_set = true;
    }

    const float* bias_ptr = (bias.defined() && bias.numel() > 0)
                                ? bias.contiguous().data_ptr<float>() : nullptr;

    auto stream = at::cuda::getCurrentCUDAStream();
    conv_transpose_wmma_kernel<<<grid, 256, shbytes, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr, Y.data_ptr<float>(),
        batch, IC, OC, IH, IW, OH, OW, KH, KW, (int)padH, (int)padW);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return Y;
}
"""

cpp_source = (
    "torch::Tensor conv_transpose2d_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "int64_t padH, int64_t padW);"
)

_custom_conv_transpose = load_inline(
    name="conv_transpose2d_wmma",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv_transpose2d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose2d accelerated with a custom implicit-GEMM TF32 tensor-core
    CUDA kernel (fast path for stride=1, dilation=1, groups=1), with a
    fallback to the standard ATen operator for all other configurations.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: tuple = (1, 1), padding: tuple = (0, 0),
                 output_padding: tuple = (0, 0), dilation: tuple = (1, 1),
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding, dilation=dilation,
            groups=groups, bias=bias)
        self.stride = self.conv_transpose2d.stride
        self.padding = self.conv_transpose2d.padding
        self.output_padding = self.conv_transpose2d.output_padding
        self.dilation = self.conv_transpose2d.dilation
        self.groups = self.conv_transpose2d.groups
        self.kernel_size = self.conv_transpose2d.kernel_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (x.is_cuda and x.dtype == torch.float32
                and self.groups == 1
                and self.stride[0] == 1 and self.stride[1] == 1
                and self.dilation[0] == 1 and self.dilation[1] == 1
                and self.output_padding[0] == 0 and self.output_padding[1] == 0
                and torch.cuda.get_device_capability(x.device)[0] >= 8):
            w = self.conv_transpose2d.weight
            b = self.conv_transpose2d.bias
            if b is None:
                b = torch.empty(0, device=x.device, dtype=x.dtype)
            kh, kw = self.kernel_size
            pad_h = kh - 1 - self.padding[0]
            pad_w = kw - 1 - self.padding[1]
            return _custom_conv_transpose.conv_transpose2d_cuda(x, w, b, pad_h, pad_w)
        return self.conv_transpose2d(x)
```