```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) accelerated with a custom tensor-core CUDA kernel.
#
# Math: ConvTranspose2d(stride=1, padding=p, kernel=K, groups=1)
#     out[b,o,y,x] = sum_{c,u,v} x[b,c,y+u-p,x+v-p] * wf[o,c,u,v]
#     wf[o,c,u,v]  = w[c,o,K-1-u,K-1-v]
# i.e. a cross-correlation with spatially flipped / channel-swapped weights.
#
# Implementation: implicit GEMM evaluated with WMMA (fp16 inputs, fp32 accum).
#   * x is first converted NCHW(fp32) -> NHWC(fp16) with padded halo columns by a
#     custom tiled-transpose kernel. In NHWC the 16-channel GEMM-k slices are
#     contiguous and 16B-aligned for any spatial shift, so A fragments are read
#     straight from global memory with ldmatrix (no im2col, no gather, no smem
#     staging for the main loop).
#   * Rearranged weights wf are stored as [9][C][OC] fp16; B fragments are also
#     read straight from global memory.
#   * Boundary output tiles (where the full 128-wide m tile does not fit) are
#     staged through shared memory with predicated stores.
# ----------------------------------------------------------------------------------

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

// ------------------------- NCHW fp32 -> NHWC fp16 (padded) ----------------------
// x : [B, C, H*W]  ->  xn : [B, H*Wp, C] with halo columns zeroed.
// Per batch this is a [C, S] -> [S, C] matrix transpose with fp32->fp16 cast.
__global__ void nchw2nhwc_pad_kernel(const float* __restrict__ x,
                                     half* __restrict__ xn,
                                     int B, int C, int H, int W, int Wp) {
    __shared__ float tile[32][33];
    const long S = (long)H * (long)W;
    const int s0 = blockIdx.x * 32;
    const int c0 = blockIdx.y * 32;
    const int b  = blockIdx.z;
    const int tx = threadIdx.x;   // 0..31
    const int ty = threadIdx.y;   // 0..7

    const float* xb = x + (long)b * C * S;
    // read 32(s) x 32(c), coalesced along s
#pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int s = s0 + ty + i;
        const int c = c0 + tx;
        float v = 0.0f;
        if (s < S && c < C) v = xb[(long)c * S + s];
        tile[ty + i][tx] = v;
    }
    __syncthreads();

    half* xnb = xn + (long)b * H * Wp * C;
#pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int c = c0 + ty + i;
        const int s = s0 + tx;
        if (s >= S || c >= C) continue;
        const int y = s / W;
        const int xx = s - y * W;
        xnb[((long)y * Wp + xx) * C + c] = __float2half(tile[tx][ty + i]);
    }
}

// zero the halo columns xx in [W, Wp) and yy in [H, OH)
__global__ void zero_halo_kernel(half* __restrict__ xn,
                                 int B, int H, int W, int Wp, int C, int OH) {
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long total = (long)B * OH * (Wp - W) * C;
    if (idx >= total) return;
    const int c  = (int)(idx % C);
    const long t = idx / C;
    const int dx = (int)(t % (Wp - W));
    const long u = t / (Wp - W);
    const int yy = (int)(u % OH);
    const int b  = (int)(u / OH);
    xn[(((long)b * OH + yy) * Wp + (W + dx)) * C + c] = __float2half(0.0f);
}

// ----------------------------- WMMA conv kernel ---------------------------------
// block: 4 warps, computes out[b, 0:OC, y, x0:x0+TM]; OC must be 64.
__global__ void __launch_bounds__(THREADS) imtconv_wmma_kernel(
    const half* __restrict__ xn,   // [B, OH, Wp, C]
    const half* __restrict__ w2,   // [9, C, OC] fp16
    float* __restrict__ out,       // [B, OC, OH, OW]
    int B, int C, int W, int OC, int OH, int OW, int Wp,
    int pad, int tiles_per_row)
{
    extern __shared__ float Cs[];  // only used by boundary tiles: TM x OC

    const int tile_id = blockIdx.x;
    const int y   = tile_id % OH;
    const int tmp = tile_id / OH;
    const int ti  = tmp % tiles_per_row;
    const int b   = tmp / tiles_per_row;
    const int x0  = ti * TM;

    const int warp = threadIdx.x >> 5;

    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);

    const bool interior = (x0 + TM <= OW) && (x0 + TM <= W + pad);
    const long HWpC = (long)OH * Wp * C;
    const half* xb = xn + (long)b * HWpC;
    const long OHOW = (long)OH * (long)OW;

    const int m_base = warp * 32;

    for (int u = 0; u < 3; ++u) {
        const int yy = y + u - pad;
        if ((unsigned)yy >= (unsigned)OH) continue;   // halo row -> contributes 0
        const half* xrow = xb + (long)yy * Wp * C + m_base * C + x0 * C;
        const half* wu = w2 + (long)u * 3 * C * OC;
#pragma unroll
        for (int v = 0; v < 3; ++v) {
            const half* wuv = wu + (long)v * C * OC;
            const long aoff = (long)(v - pad) * C;
#pragma unroll
            for (int c0 = 0; c0 < 64; c0 += 16) {
                wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> bfrag[4];
#pragma unroll
                for (int nf = 0; nf < 4; ++nf)
                    wmma::load_matrix_sync(bfrag[nf], wuv + (long)c0 * OC + nf * 16, OC);
                const half* aptr = xrow + aoff + c0;
#pragma unroll
                for (int mf = 0; mf < 2; ++mf) {
                    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> afrag;
                    wmma::load_matrix_sync(afrag, aptr + mf * 16 * C, C);
#pragma unroll
                    for (int nf = 0; nf < 4; ++nf)
                        wmma::mma_sync(acc[mf][nf], afrag, bfrag[nf], acc[mf][nf]);
                }
            }
        }
    }

    if (interior) {
        float* outb = out + (long)b * OC * OHOW + (long)y * OW + x0 + m_base;
#pragma unroll
        for (int mf = 0; mf < 2; ++mf) {
#pragma unroll
            for (int nf = 0; nf < 4; ++nf) {
                wmma::store_matrix_sync(outb + (long)(nf * 16) * OHOW + mf * 16,
                                        acc[mf][nf], OHOW, wmma::mem_col_major);
            }
        }
    } else {
        // boundary tile: stage through shared memory, then predicated writes
        const int m_valid = OW - x0;
        float* Cw = Cs + warp * (32 * OC);
#pragma unroll
        for (int mf = 0; mf < 2; ++mf) {
#pragma unroll
            for (int nf = 0; nf < 4; ++nf) {
                wmma::store_matrix_sync(Cw + mf * 16 * OC + nf * 16,
                                        acc[mf][nf], OC, wmma::mem_row_major);
            }
        }
        __syncthreads();
        const long bbase = (long)b * OC * OHOW + (long)y * OW + x0;
        for (int idx = threadIdx.x; idx < m_valid * OC; idx += THREADS) {
            const int m = idx / OC;
            const int o = idx - m * OC;
            out[bbase + (long)o * OHOW + m] = Cs[m * OC + o];
        }
    }
}

// --------------------------------- host wrapper ----------------------------------
torch::Tensor conv_transpose2d_wmma(torch::Tensor x, torch::Tensor w2,
                                    int64_t pad) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32, "x must be CUDA fp32");
    TORCH_CHECK(w2.is_cuda() && w2.dtype() == torch::kHalf, "w2 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 = w2.size(2);
    TORCH_CHECK(C == 64 && OC == 64, "fast kernel requires C == OC == 64");

    const int OH = H + 2;          // K = 3
    const int OW = W + 2;
    const int Wp = W + TM + 2;     // padded columns (halo)

    auto opts_h = x.options().dtype(torch::kHalf);
    auto xn = torch::empty({(long)B * OH * Wp * C}, opts_h);

    auto stream = at::cuda::getCurrentCUDAStream();

    // halo zero fill
    {
        const long total = (long)B * OH * (Wp - W) * C;
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        zero_halo_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            reinterpret_cast<half*>(xn.data_ptr<at::Half>()),
            B, H, W, Wp, C, OH);
    }
    // transpose + cast
    {
        const long S = (long)H * W;
        dim3 block(32, 8);
        dim3 grid((unsigned)((S + 31) / 32), (unsigned)((C + 31) / 32), (unsigned)B);
        nchw2nhwc_pad_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(),
            reinterpret_cast<half*>(xn.data_ptr<at::Half>()),
            B, C, H, W, Wp);
    }

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

    const int tiles_per_row = (OW + TM - 1) / TM;
    const long blocks = (long)B * OH * tiles_per_row;

    const bool has_edge = (OW % TM) != 0;
    size_t sh = has_edge ? (size_t)TM * OC * sizeof(float) : 0;
    if (sh > 48 * 1024) {
        static size_t max_sh = 0;
        if (sh > max_sh) {
            cudaFuncSetAttribute(imtconv_wmma_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize, (int)sh);
            max_sh = sh;
        }
    }

    imtconv_wmma_kernel<<<(unsigned)blocks, THREADS, sh, stream>>>(
        reinterpret_cast<const half*>(xn.data_ptr<at::Half>()),
        reinterpret_cast<const half*>(w2.data_ptr<at::Half>()),
        out.data_ptr<float>(),
        B, C, W, OC, OH, OW, Wp, (int)pad, tiles_per_row);

    return out;
}
"""

cpp_src = "torch::Tensor conv_transpose2d_wmma(torch::Tensor x, torch::Tensor w2, int64_t pad);"

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


class ModelNew(nn.Module):
    """
    ConvTranspose2d accelerated with a custom tensor-core implicit-GEMM CUDA
    kernel (fast path: stride=1, groups=1, kernel 3x3, in_channels =
    out_channels = 64, no bias, CUDA fp32 input). All other configurations use
    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.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):
            # 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.
            w2 = (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_wmma(x.contiguous(), w2, self.padding)
        return F.conv_transpose2d(x, w.weight, w.bias, self.stride, self.padding,
                                  0, self.groups)
```