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

# ----------------------------------------------------------------------------------
# ConvTranspose2d (stride=1, K=3, C=OC=64) with a custom tensor-core CUDA kernel.
#
# Math: ConvTranspose2d(stride=1, padding=p, kernel=K=3, groups=1) is exactly a
# cross-correlation with spatially flipped, channel-swapped weights and effective
# padding p' = 2-p:
#     out[b,o,y,x] = sum_{c,u,v} x[b,c,y+u-p',x+v-p'] * wf[u,v,c,o]
#     wf[u,v,c,o]  = w[c,o,2-u,2-v]
#
# Implementation: single fused implicit-GEMM kernel using WMMA
# (fp16 inputs, fp32 accumulate):
#   * each block computes one output image row segment of TM=128 pixels x all
#     64 output channels for one (batch, y);
#   * the 3 required input rows (64 channels x 130 pixels) are loaded straight
#     from the fp32 NCHW input, converted to fp16 on the fly and staged in a
#     padded shared-memory NHWC tile (channel stride 72 halves kills bank
#     conflicts, every fragment row stays 16B aligned);
#   * WMMA: A fragments = staged pixels (row_major, ldm=72), B fragments =
#     fp16 weights read from global memory (L1 resident), accumulators are
#     written straight to the fp32 NCHW output with column-major stores;
#   * tiles overlap at row ends so no edge path is ever needed.
# This removes the fp32->fp16 preprocessing pass entirely: DRAM traffic is
# just input read + output write, while running at fp16 tensor-core rate.
# ----------------------------------------------------------------------------------

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

using namespace nvcuda;

#define TM 128
#define THREADS 128
#define LC 72                     // padded channel stride (multiple of 8)
#define PW (TM + 2)               // staged pixels per row
#define SMEM_BYTES (3 * PW * LC * 2)

// blockIdx.x -> pixel tile, blockIdx.y -> output row y, blockIdx.z -> batch b
__global__ void __launch_bounds__(THREADS) fused_tconv_kernel(
    const float* __restrict__ x,   // [B, 64, H, W] fp32
    const half*  __restrict__ wf,  // [9, 64, 64] fp16 (u, v, c, o)
    float* __restrict__ out,       // [B, 64, OH, OW] fp32
    int H, int W, int OH, int OW, int pprime, int tiles_per_row)
{
    extern __shared__ half xs[];   // [3][PW][LC]

    const int b   = blockIdx.z;
    const int y   = blockIdx.y;
    const int ti  = blockIdx.x;
    int x0 = ti * TM;
    if (x0 > OW - TM) x0 = OW - TM;   // overlap last tile; OW >= TM guaranteed

    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

    const long S = (long)H * (long)W;
    const float* xb = x + (long)b * 64 * S;

    // ---------------- stage 3 input rows into shared memory ----------------
#pragma unroll
    for (int u = 0; u < 3; ++u) {
        const int yy = y + u - pprime;
        const bool rok = ((unsigned)yy < (unsigned)H);
        half* xu = xs + u * PW * LC;
        // 4 c-chunks of 16 channels, 32 lanes -> pixels (coalesced reads)
#pragma unroll
        for (int cb = 0; cb < 4; ++cb) {
            const int c = cb * 16 + wid * 4 + (lane >> 3);
            const int lx = lane & 7;
            const float* xrow = xb + (long)c * S + (long)yy * W;
            half* srow = xu + c;
            if (rok) {
#pragma unroll
                for (int p0 = 0; p0 < 160; p0 += 32) {
                    const int p = p0 + lx * 4;
                    if (p < PW) {
                        const int xg = x0 - pprime + p;
                        float v0 = 0.f, v1 = 0.f, v2 = 0.f, v3 = 0.f;
                        if ((unsigned)xg < (unsigned)W) {
                            if (xg + 3 < W) {
                                const float4 f = *reinterpret_cast<const float4*>(xrow + xg);
                                v0 = f.x; v1 = f.y; v2 = f.z; v3 = f.w;
                            } else {
                                v0 = xrow[xg];
                                v1 = xrow[xg + 1];
                                v2 = xrow[xg + 2];
                            }
                        }
                        half2* d = reinterpret_cast<half2*>(srow + (long)p * LC);
                        const half2 z01 = __floats2half2_rn(v0, v1);
                        const half2 z23 = __floats2half2_rn(v2, v3);
                        d[0] = z01;
                        d[LC >> 1] = z23;
                    }
                }
            } else {
                const half2 zz = __floats2half2_rn(0.0f, 0.0f);
#pragma unroll
                for (int p0 = 0; p0 < 160; p0 += 32) {
                    const int p = p0 + lx * 4;
                    if (p < PW) {
                        half2* d = reinterpret_cast<half2*>(srow + (long)p * LC);
                        d[0] = zz;
                        d[LC >> 1] = zz;
                    }
                }
            }
        }
    }
    __syncthreads();

    // ------------------------------ implicit GEMM ---------------------------
    const int warp = wid;
    const int m0 = warp * 32;       // this warp's first pixel (of TM)

    wmma::fragment<wmma::accumulator, 16, 16, 16, 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);

#pragma unroll
    for (int u = 0; u < 3; ++u) {
        const half* xu = xs + u * PW * LC;
        const half* wu = wf + (long)u * 3 * 4096;
#pragma unroll
        for (int v = 0; v < 3; ++v) {
            const half* wuv = wu + (long)v * 4096;
            const half* arow = xu + (long)(m0 + v) * LC;
#pragma unroll
            for (int c0 = 0; c0 < 64; c0 += 16) {
                wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a[2];
                wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> bb[4];
                wmma::load_matrix_sync(a[0], arow + c0, LC);
                wmma::load_matrix_sync(a[1], arow + 16 * LC + c0, LC);
#pragma unroll
                for (int j = 0; j < 4; ++j)
                    wmma::load_matrix_sync(bb[j], wuv + c0 * 64 + j * 16, 64);
#pragma unroll
                for (int i = 0; i < 2; ++i)
#pragma unroll
                    for (int j = 0; j < 4; ++j)
                        wmma::mma_sync(acc[i][j], a[i], bb[j], acc[i][j]);
            }
        }
    }

    // --------------------------- stores (no edges) ---------------------------
    const long OHOW = (long)OH * (long)OW;
    float* obase = out + (long)b * 64 * OHOW + (long)y * OW + x0 + m0;
#pragma unroll
    for (int i = 0; i < 2; ++i) {
#pragma unroll
        for (int j = 0; j < 4; ++j) {
            wmma::store_matrix_sync(obase + (long)(j * 16) * OHOW + i * 16,
                                    acc[i][j], OHOW, wmma::mem_col_major);
        }
    }
}

// --------------------------------- host wrapper ------------------------------
torch::Tensor conv_transpose2d_fused(torch::Tensor x, torch::Tensor wf,
                                     int64_t pad) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32, "x must be CUDA fp32");
    TORCH_CHECK(wf.is_cuda() && wf.dtype() == torch::kHalf, "wf must be CUDA fp16");
    const int B = x.size(0);
    const int C = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);
    const int OC = wf.size(2);
    TORCH_CHECK(C == 64 && OC == 64, "fast kernel requires C == OC == 64");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const int pprime = 2 - (int)pad;               // K = 3
    const int OH = H + 2 * pprime - 2;
    const int OW = W + 2 * pprime - 2;
    TORCH_CHECK(OH > 0 && OW >= TM, "output too small for fast kernel");

    auto out = torch::empty({B, OC, OH, OW}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(fused_tconv_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             SMEM_BYTES);
        configured = true;
    }

    const int tiles_per_row = (OW + TM - 1) / TM;
    dim3 grid((unsigned)tiles_per_row, (unsigned)OH, (unsigned)B);
    fused_tconv_kernel<<<grid, THREADS, SMEM_BYTES, stream>>>(
        x.data_ptr<float>(),
        reinterpret_cast<const half*>(wf.data_ptr<at::Half>()),
        out.data_ptr<float>(),
        H, W, OH, OW, pprime, tiles_per_row);

    return out;
}
"""

cpp_src = "torch::Tensor conv_transpose2d_fused(torch::Tensor x, torch::Tensor wf, int64_t pad);"

_conv_mod = load_inline(
    name="conv_transpose2d_fused_v5",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["conv_transpose2d_fused"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose2d accelerated with a fused tensor-core implicit-GEMM CUDA
    kernel (fast path: stride=1, groups=1, kernel 3x3, padding in {0,1,2},
    output_padding=0, in_channels = out_channels = 64, no bias, CUDA fp32
    input, output width >= 128). All other configurations fall back to the
    cuDNN implementation.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1,
                 padding: int = 0, output_padding: int = 0, 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,
                                                   groups=groups, bias=bias)
        self.stride = stride
        self.padding = padding
        self.output_padding = output_padding
        self.groups = groups
        self.kernel_size = kernel_size

        self._fast = (stride == 1 and groups == 1 and kernel_size == 3
                      and output_padding == 0 and 0 <= padding <= 2
                      and in_channels == 64 and out_channels == 64)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv_transpose2d
        if (self._fast and w.bias is None and x.is_cuda
                and x.dtype == torch.float32 and x.dim() == 4 and x.size(1) == 64):
            ow = x.size(3) + 2 - 2 * self.padding
            if ow >= 128:
                # rearranged weights wf[u,v,c,o] = w[c,o,2-u,2-v], rebuilt every
                # call (tiny: 36k elements) so weight updates are always honored.
                wf = (w.weight.detach()
                      .permute(2, 3, 0, 1)          # [3,3,C,OC]
                      .flip(0, 1)                   # spatial flip
                      .contiguous()
                      .to(torch.float16))
                return _conv_mod.conv_transpose2d_fused(x.contiguous(), wf, self.padding)
        return F.conv_transpose2d(x, w.weight, w.bias, self.stride, self.padding,
                                  self.output_padding, self.groups)
```