```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) accelerated with a custom tensor-core
# CUDA kernel.
#
# Math: ConvTranspose2d(stride=1, padding=p, kernel=K, groups=1) is exactly a
# cross-correlation with spatially flipped, channel-swapped weights and padding
# p' = K-1-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,K-1-u,K-1-v]
#
# Implementation: implicit GEMM with WMMA (fp16 inputs, fp32 accumulate).
#   * A tiled-transpose kernel converts x from NCHW(fp32) to NHWC(fp16) inside a
#     zero-padded halo buffer [B, HB, Wp, C] so every 16-channel GEMM-k slice is
#     contiguous and 128B aligned for any spatial shift. A fragments are read
#     straight from global memory (no im2col, no gather).
#   * Accumulators are staged through shared memory and written to the NCHW
#     output with plain coalesced 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 (halo buffer) ---------------------
// x : [B, C, H*W]  ->  xn : [B, HB, Wp, C], data rows [0,H), cols [XOFF, XOFF+W).
// block: (64,4) threads, 64x64 shared tile; requires C == 64 (guarded on host).
__global__ void nchw2nhwc_kernel(const float* __restrict__ x,
                                 half* __restrict__ xn,
                                 int B, int C, int H, int W,
                                 int HB, int Wp, int XOFF) {
    __shared__ float tile[64][65];
    const long S = (long)H * (long)W;
    const int s0 = blockIdx.x * 64;
    const int c0 = blockIdx.y * 64;
    const int b  = blockIdx.z;
    const int tx = threadIdx.x;   // 0..63
    const int ty = threadIdx.y;   // 0..3

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

    half* xnb = xn + (long)b * HB * Wp * C;
    const int c = c0 + tx;
    if (c >= C) return;

    // coalesced write along c; walk s in strides of 4 with cheap row tracking
    int s  = s0 + ty;
    int y  = s / W;
    int xx = s - y * W;
#pragma unroll
    for (int i = 0; i < 16; ++i) {
        if (s < S) {
            xnb[((long)y * Wp + XOFF + xx) * C + c] =
                __float2half(tile[tx][ty + 4 * i]);
        }
        s += 4;
        xx += 4;
        while (xx >= W) { xx -= W; ++y; }
    }
}

// zero everything in the halo buffer that is not the data region (C == 64).
__global__ void zero_halo_kernel(half* __restrict__ xn,
                                 int B, int H, int W,
                                 int HB, int Wp, int XOFF) {
    const int y = blockIdx.x;   // 0..HB-1
    const int b = blockIdx.y;
    half* row = xn + ((long)b * HB + y) * (long)Wp * 64;
    if (y < H) {
        const int n = (Wp - W) * 64;   // halo columns of a data row
        for (int idx = threadIdx.x; idx < n; idx += blockDim.x) {
            const int hc  = idx >> 6;
            const int c   = idx & 63;
            const int col = (hc < XOFF) ? hc : (XOFF + W + hc - XOFF);
            row[(long)col * 64 + c] = __float2half(0.0f);
        }
    } else {
        const int n = Wp * 64;         // full bottom rows
        for (int idx = threadIdx.x; idx < n; idx += blockDim.x) {
            row[idx] = __float2half(0.0f);
        }
    }
}

// ----------------------------- WMMA conv kernel ----------------------------------
// block: 4 warps -> out[b, 0:OC, y, x0:x0+TM]; requires C == OC == 64.
__global__ void __launch_bounds__(THREADS) imtconv_wmma_kernel(
    const half* __restrict__ xn,   // [B, HB, 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 HB, int Wp, int XOFF, int pprime, int tiles_per_row)
{
    extern __shared__ float Cs[];  // TM x OC staging (layout [o][m])

    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;
    const int m_base = warp * 32;

    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 half* xb = xn + (long)b * HB * Wp * C;
    const long OHOW = (long)OH * (long)OW;

    for (int u = 0; u < 3; ++u) {
        const int yy = y + u - pprime;
        if ((unsigned)yy >= (unsigned)HB) continue;   // outside halo buffer
        const half* xrow = xb + ((long)yy * Wp + XOFF + x0 + m_base) * 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 - pprime) * 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]);
                }
            }
        }
    }

    // stage accumulators through shared memory: Cs[o * TM + m]
#pragma unroll
    for (int mf = 0; mf < 2; ++mf)
#pragma unroll
        for (int nf = 0; nf < 4; ++nf)
            wmma::store_matrix_sync(Cs + (nf * 16) * TM + warp * 32 + mf * 16,
                                    acc[mf][nf], TM, wmma::mem_col_major);
    __syncthreads();

    // coalesced NCHW stores; correct for partial (edge) tiles as well
    const int m_valid = min(TM, OW - x0);
    const long bbase = (long)b * OC * OHOW + (long)y * OW + x0;
    for (int o = 0; o < OC; ++o) {
        const float* srow = Cs + o * TM;
        float* orow = out + bbase + (long)o * OHOW;
        for (int m = threadIdx.x; m < m_valid; m += THREADS)
            orow[m] = srow[m];
    }
}

// --------------------------------- 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 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 > 0, "invalid output size");

    const int HB   = H + 2;
    const int XOFF = 2;
    const int Wp   = W + TM + 2 * XOFF + 8;   // extra slack for edge-tile reads

    auto xn = torch::empty({(long)B * HB * Wp * C},
                           x.options().dtype(torch::kHalf));
    auto stream = at::cuda::getCurrentCUDAStream();
    half* xn_ptr = reinterpret_cast<half*>(xn.data_ptr<at::Half>());

    // zero the halo region
    {
        dim3 grid((unsigned)HB, (unsigned)B);
        zero_halo_kernel<<<grid, 256, 0, stream>>>(xn_ptr, B, H, W, HB, Wp, XOFF);
    }
    // transpose + cast into the halo buffer
    {
        const long S = (long)H * W;
        dim3 block(64, 4);
        dim3 grid((unsigned)((S + 63) / 64), (unsigned)((C + 63) / 64),
                  (unsigned)B);
        nchw2nhwc_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), xn_ptr, B, C, H, W, HB, Wp, XOFF);
    }

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

    const int tiles_per_row = (OW + TM - 1) / TM;
    const long blocks = (long)B * tiles_per_row * OH;
    const size_t sh = (size_t)TM * OC * sizeof(float);   // 32 KB

    imtconv_wmma_kernel<<<(unsigned)blocks, THREADS, sh, stream>>>(
        xn_ptr,
        reinterpret_cast<const half*>(w2.data_ptr<at::Half>()),
        out.data_ptr<float>(),
        B, C, W, OC, OH, OW, HB, Wp, XOFF, pprime, 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_v3",
    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, padding in {0,1,2},
    output_padding=0, in_channels = out_channels = 64, no bias, CUDA fp32
    input). 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):
            # 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,
                                  self.output_padding, self.groups)
```