```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) 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 custom 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 32B 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, correctly-aligned 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).
__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[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;
    // coalesced read along s
#pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int c = c0 + ty + i;
        const int s = s0 + tx;
        float v = 0.0f;
        if (c < C && s < S) v = xb[(long)c * S + s];
        tile[ty + i][tx] = v;
    }
    __syncthreads();

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

// zero everything in the halo buffer that is not the data region
__global__ void zero_halo_kernel(half* __restrict__ xn,
                                 int B, int C, int H, int W,
                                 int HB, int Wp, int XOFF) {
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    const long row_halo = (long)H * (Wp - W);        // halo cols of rows [0,H)
    const long tail_halo = (long)(HB - H) * Wp;      // full rows [H, HB)
    const long halo_pb = row_halo + tail_halo;
    const long total = (long)B * halo_pb * C;
    if (idx >= total) return;

    const int c = (int)(idx % C);
    const long t = idx / C;
    const int b = (int)(t / halo_pb);
    const long sh = t % halo_pb;

    int y, col;
    if (sh < row_halo) {
        y = (int)(sh / (Wp - W));
        const int hc = (int)(sh % (Wp - W));
        col = (hc < XOFF) ? hc : (XOFF + W + (hc - XOFF));
    } else {
        const long t2 = sh - row_halo;
        y = H + (int)(t2 / Wp);
        col = (int)(t2 % Wp);
    }
    xn[(((long)b * HB + y) * Wp + col) * C + c] = __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();

    const int m_valid = min(TM, OW - x0);
    const long bbase = (long)b * OC * OHOW + (long)y * OW + x0;
    for (int idx = threadIdx.x; idx < m_valid * OC; idx += THREADS) {
        const int o = idx / TM;
        const int m = idx - o * TM;
        out[bbase + (long)o * OHOW + m] = Cs[o * TM + 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;

    auto xn = torch::empty({(long)B * HB * Wp * C},
                           x.options().dtype(torch::kHalf));
    auto stream = at::cuda::getCurrentCUDAStream();

    // zero the halo region
    {
        const long total = (long)B * ((long)HB * Wp - (long)H * 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, C, H, W, HB, Wp, XOFF);
    }
    // transpose + cast into the halo buffer
    {
        const long S = (long)H * W;
        dim3 block(32, 8);
        dim3 grid((unsigned)((S + 31) / 32), (unsigned)((C + 31) / 32),
                  (unsigned)B);
        nchw2nhwc_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(),
            reinterpret_cast<half*>(xn.data_ptr<at::Half>()),
            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>>>(
        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, 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_v2",
    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},
    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)
```