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


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

using namespace nvcuda;

// ---------------------------------------------------------------------------
// Rearrange ConvTranspose2d weight W[c][o][kh][kw] into a GEMM friendly
// layout WT[tap][c][o] where tap = a*KW + b and
// WT[a*KW+b][c][o] = W[c][o][KH-1-a][KW-1-b].
// With this, ConvTranspose2d(stride=1) becomes the correlation:
//   out[n,o,h,w] = sum_{tap=(a,b)} sum_c
//        x[n, c, h + a - (KH-1) + pad_h, w + b - (KW-1) + pad_w] * WT[tap][c][o]
// ---------------------------------------------------------------------------
__global__ void rearrange_weight_kernel(const float* __restrict__ w,
                                        float* __restrict__ wt,
                                        int C_in, int C_out, int KH, int KW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = C_in * C_out * KH * KW;
    if (idx >= total) return;
    int kw = idx % KW;
    int kh = (idx / KW) % KH;
    int o  = (idx / (KW * KH)) % C_out;
    int c  = idx / (KW * KH * C_out);
    int a = KH - 1 - kh;
    int b = KW - 1 - kw;
    wt[((a * KW + b) * C_in + c) * C_out + o] = w[idx];
}

// ---------------------------------------------------------------------------
// Tensor core (TF32) implicit GEMM kernel.
// C[M, C_out] = sum_k A[M, k] * B[k, C_out]
//   M    = output pixels of one (h_tile x w_tile) block for a fixed batch n
//   k    = (tap, c) pairs, iterated as: for tap -> for c-chunk of 32
//   A[c][m] = x[n, c0+c, h(m)+a-OFF_H, w(m)+b-OFF_W]  (zero outside image)
//   B[c][o] = WT[tap][c0+c][o]
// Each block: 128 pixels (8 x 16 tile) x 128 output channels, 8 warps,
// warp tile 64(m) x 32(n), wmma m16n8k8.
// ---------------------------------------------------------------------------
#define WARPS_M 2
#define WARPS_N 4
#define FRAGS_M 4
#define FRAGS_N 4
#define KCHUNK 32
#define TILE_H 8
#define TILE_W 16
#define SMEM_STRIDE 132

__global__ void __launch_bounds__(256)
conv_tensor_kernel(const float* __restrict__ x,
                   const float* __restrict__ wt,
                   const float* __restrict__ bias,
                   float* __restrict__ out,
                   int C_in, int C_out,
                   int H_in, int W_in,
                   int H_out, int W_out,
                   int KH, int KW,
                   int pad_h, int pad_w,
                   int tiles_w) {
#if __CUDA_ARCH__ >= 800
    __shared__ float As[2][KCHUNK * SMEM_STRIDE];
    __shared__ float Bs[2][KCHUNK * SMEM_STRIDE];

    const int OFF_H = KH - 1 - pad_h;
    const int OFF_W = KW - 1 - pad_w;

    const int tid    = threadIdx.x;
    const int warp   = tid >> 5;
    const int warp_m = warp % WARPS_M;
    const int warp_n = warp / WARPS_M;

    const int tile_m = blockIdx.x;
    const int n      = blockIdx.y;
    const int h0 = (tile_m / tiles_w) * TILE_H;
    const int w0 = (tile_m % tiles_w) * TILE_W;

    // pixel (h,w) of each smem-A column this thread will store into epilogue
    // (also used for the final write back)
    const int e_m  = tid & 127;              // 0..127
    const int e_oc = (tid >> 7) * 64;        // 0 or 64
    const int e_h  = h0 + (e_m >> 4);
    const int e_w  = w0 + (e_m & 15);
    const bool e_valid = (e_h < H_out) && (e_w < W_out);

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

    const int mtile = tid >> 4;              // 0..15 : (c_pair, h_half)
    const int wfrag = tid & 15;              // 0..15 : w position
    const int ld_c  = (mtile & 1) * 16;      // 0 or 16
    const int ld_hh = (mtile >> 1) * 4;      // 0,4,8,12
    const int ld_h1 = h0 + ld_hh + 2;        // global h of second row of pair
    const int ld_w  = w0 + wfrag;            // global w

    const int taps = KH * KW;
    int stage = 0;

    for (int tap = 0; tap < taps; ++tap) {
        const int a = tap / KW;
        const int b = tap - a * KW;
        const int gx_h = ld_h1 + a - OFF_H;      // row of second load
        const int gx_w = ld_w  + b - OFF_W;
        const bool row_ok = (gx_h >= 1) && (gx_h <= H_in);   // both gx_h-1 and gx_h in [0,H_in)
        const bool col_ok = (gx_w >= 0) && (gx_w <  W_in);
        const bool ld_ok  = row_ok && col_ok;

        const float* wt_tap = wt + (long)tap * C_in * C_out;

        for (int c0 = 0; c0 < C_in; c0 += KCHUNK) {
            // ---- load A tile (16B per thread) and B tile (16B per thread) ----
            {
                float4 va = make_float4(0.f, 0.f, 0.f, 0.f);
                if (ld_ok) {
                    const float* src = x + (((long)(n * C_in + c0 + ld_c) * H_in + (gx_h - 1)) * W_in + gx_w);
                    va.x = src[0];            va.y = src[W_in];
                    va.z = src[16L * H_in * W_in];
                    va.w = src[16L * H_in * W_in + W_in];
                }
                *reinterpret_cast<float4*>(&As[stage][ld_c * SMEM_STRIDE + (ld_hh * 2) * 16 + wfrag]) = va;
            }
            {
                const float4 vb = *reinterpret_cast<const float4*>(
                    wt_tap + ((long)(c0 + (tid >> 3)) * C_out) + ((tid & 31) * 4));
                *reinterpret_cast<float4*>(&Bs[stage][(tid >> 3) * SMEM_STRIDE + (tid & 31) * 4]) = vb;
            }
            __syncthreads();

            // ---- compute ----
            const float* As_ = As[stage];
            const float* Bs_ = Bs[stage];
            #pragma unroll
            for (int kk = 0; kk < KCHUNK; kk += 8) {
                wmma::fragment<wmma::matrix_a, 16, 8, 8, wmma::precision::tf32, wmma::col_major> fa[FRAGS_M];
                wmma::fragment<wmma::matrix_b, 16, 8, 8, wmma::precision::tf32, wmma::row_major> fb[FRAGS_N];
                #pragma unroll
                for (int i = 0; i < FRAGS_M; ++i)
                    wmma::load_matrix_sync(fa[i], As_ + (kk * SMEM_STRIDE) + warp_m * 64 + i * 16, SMEM_STRIDE);
                #pragma unroll
                for (int j = 0; j < FRAGS_N; ++j)
                    wmma::load_matrix_sync(fb[j], Bs_ + (kk * SMEM_STRIDE) + warp_n * 32 + j * 8, SMEM_STRIDE);
                #pragma unroll
                for (int i = 0; i < FRAGS_M; ++i)
                    #pragma unroll
                    for (int j = 0; j < FRAGS_N; ++j)
                        wmma::mma_sync(acc[i][j], fa[i], fb[j], acc[i][j]);
            }
            __syncthreads();
            stage ^= 1;
        }
    }

    // ---- epilogue: store accumulators through shared memory ----
    float* Cs = As[0];
    #pragma unroll
    for (int i = 0; i < FRAGS_M; ++i)
        #pragma unroll
        for (int j = 0; j < FRAGS_N; ++j)
            wmma::store_matrix_sync(Cs + (warp_m * 64 + i * 16) * SMEM_STRIDE + warp_n * 32 + j * 8,
                                    acc[i][j], SMEM_STRIDE, wmma::mem_row_major);
    __syncthreads();

    if (e_valid) {
        const float* crow = Cs + e_m * SMEM_STRIDE + e_oc;
        if (bias != nullptr) {
            const float* bp = bias + n * 0;  // silence unused warning pattern
            #pragma unroll
            for (int oi = 0; oi < 64; ++oi) {
                out[(((long)(n * C_out + e_oc + oi) * H_out + e_h) * W_out + e_w)] =
                    crow[oi] + bias[e_oc + oi];
            }
            (void)bp;
        } else {
            #pragma unroll
            for (int oi = 0; oi < 64; ++oi) {
                out[(((long)(n * C_out + e_oc + oi) * H_out + e_h) * W_out + e_w)] = crow[oi];
            }
        }
    }
#endif
}

// ---------------------------------------------------------------------------
// FP32 fallback kernel (used on pre-Ampere GPUs or small channel counts).
// One thread per output element.
// ---------------------------------------------------------------------------
__global__ void conv_transpose_fallback_kernel(
        const float* __restrict__ x, const float* __restrict__ w,
        const float* __restrict__ bias, float* __restrict__ out,
        int N, int C_in, int C_out, int H_in, int W_in,
        int H_out, int W_out, int KH, int KW, int pad_h, int pad_w) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long total = (long)N * C_out * H_out * W_out;
    if (idx >= total) return;
    int w_o = idx % W_out;
    int h_o = (idx / W_out) % H_out;
    int o   = (idx / ((long)W_out * H_out)) % C_out;
    int n   = idx / ((long)W_out * H_out * C_out);
    float acc = bias ? bias[o] : 0.f;
    for (int c = 0; c < C_in; ++c) {
        const float* xp = x + ((long)(n * C_in + c) * H_in) * W_in;
        const float* wp = w + ((long)(c * C_out + o) * KH) * KW;
        for (int kh = 0; kh < KH; ++kh) {
            int hi = h_o + pad_h - kh;
            if (hi < 0 || hi >= H_in) continue;
            for (int kw = 0; kw < KW; ++kw) {
                int wi = w_o + pad_w - kw;
                if (wi < 0 || wi >= W_in) continue;
                acc += xp[(long)hi * W_in + wi] * wp[kh * KW + kw];
            }
        }
    }
    out[idx] = acc;
}

// ---------------------------------------------------------------------------
torch::Tensor conv_transpose_forward(torch::Tensor x, torch::Tensor weight,
                                     bool has_bias, torch::Tensor bias,
                                     int64_t stride_h, int64_t stride_w,
                                     int64_t pad_h, int64_t pad_w,
                                     int64_t out_pad_h, int64_t out_pad_w,
                                     int64_t dil_h, int64_t dil_w,
                                     int64_t groups) {
    // Fast custom path: stride 1, dilation 1, groups 1, output_padding 0
    bool fast = (stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1 &&
                 groups == 1 && out_pad_h == 0 && out_pad_w == 0);
    if (!fast) {
        return at::conv_transpose2d(x, weight, has_bias ? bias : torch::Tensor(),
                                    {stride_h, stride_w}, {pad_h, pad_w},
                                    {out_pad_h, out_pad_w}, groups, {dil_h, dil_w});
    }

    const int N     = x.size(0);
    const int C_in  = x.size(1);
    const int H_in  = x.size(2);
    const int W_in  = x.size(3);
    const int KH    = weight.size(2);
    const int KW    = weight.size(3);
    const int C_out = weight.size(1);
    const int H_out = H_in - 2 * (int)pad_h + KH;
    const int W_out = W_in - 2 * (int)pad_w + KW;

    auto out = at::empty({N, C_out, H_out, W_out}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    int dev = 0; cudaGetDevice(&dev);
    cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev);
    bool tensor_ok = (prop.major >= 8) && (C_in % 32 == 0) && (C_out % 128 == 0);

    const float* bias_ptr = has_bias ? bias.data_ptr<float>() : nullptr;

    if (tensor_ok) {
        auto wt = at::empty({KH * KW * C_in * C_out}, x.options());
        int total = C_in * C_out * KH * KW;
        rearrange_weight_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
            weight.data_ptr<float>(), wt.data_ptr<float>(), C_in, C_out, KH, KW);

        int tiles_w = (W_out + TILE_W - 1) / TILE_W;
        int tiles_h = (H_out + TILE_H - 1) / TILE_H;
        dim3 grid(tiles_w * tiles_h, N, 1);
        conv_tensor_kernel<<<grid, 256, 0, stream>>>(
            x.data_ptr<float>(), wt.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), C_in, C_out, H_in, W_in, H_out, W_out,
            KH, KW, (int)pad_h, (int)pad_w, tiles_w);
    } else {
        long totalel = (long)N * C_out * H_out * W_out;
        int threads = 256;
        long blocks = (totalel + threads - 1) / threads;
        conv_transpose_fallback_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N, C_in, C_out, H_in, W_in, H_out, W_out,
            KH, KW, (int)pad_h, (int)pad_w);
    }
    return out;
}
"""

conv_transpose_cpp_source = (
    "torch::Tensor conv_transpose_forward(torch::Tensor x, torch::Tensor weight, "
    "bool has_bias, torch::Tensor bias, int64_t stride_h, int64_t stride_w, "
    "int64_t pad_h, int64_t pad_w, int64_t out_pad_h, int64_t out_pad_w, "
    "int64_t dil_h, int64_t dil_w, int64_t groups);"
)

conv_transpose_mod = load_inline(
    name="conv_transpose_tensor",
    cpp_sources=conv_transpose_cpp_source,
    cuda_sources=conv_transpose_source,
    functions=["conv_transpose_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose2d accelerated with a custom TF32 tensor-core implicit GEMM
    CUDA kernel (with automatic fallbacks for unsupported 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.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = tuple(kernel_size)
        self.stride = tuple(stride)
        self.padding = tuple(padding)
        self.output_padding = tuple(output_padding)
        self.dilation = tuple(dilation)
        self.groups = groups

        kh, kw = self.kernel_size
        self.weight = nn.Parameter(
            torch.empty(in_channels, out_channels // groups, kh, kw))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_channels))
        else:
            self.register_parameter("bias", None)
        self.reset_parameters()

        self._fast = (self.stride == (1, 1) and self.dilation == (1, 1) and
                      self.groups == 1 and self.output_padding == (0, 0))

    def reset_parameters(self):
        kh, kw = self.kernel_size
        fan_in = self.in_channels * kh * kw
        bound = 1.0 / (fan_in ** 0.5) if fan_in > 0 else 0.0
        nn.init.uniform_(self.weight, -bound, bound)
        if self.bias is not None:
            nn.init.uniform_(self.bias, -bound, bound)

    def _load_state_dict_from_ref(self, ref: nn.ConvTranspose2d):
        with torch.no_grad():
            self.weight.copy_(ref.weight)
            if self.bias is not None and ref.bias is not None:
                self.bias.copy_(ref.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast and x.is_cuda:
            return conv_transpose_mod.conv_transpose_forward(
                x, self.weight, self.bias is not None,
                self.bias if self.bias is not None else self.weight,
                self.stride[0], self.stride[1],
                self.padding[0], self.padding[1],
                self.output_padding[0], self.output_padding[1],
                self.dilation[0], self.dilation[1], self.groups)
        return F.conv_transpose2d(
            x, self.weight, self.bias, self.stride, self.padding,
            self.output_padding, self.groups, self.dilation)
```