```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
from torch.utils.cpp_extension import load_inline


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

#define KH 3
#define KW 5
#define CIN 64
#define KCONV 960          // CIN * KH * KW
#define KTILES 30          // KCONV / BK
#define BK 32
#define BN 128
#define TX 16
#define TY 8
#define RR (TY + KH - 1)   // 10 input rows per output tile
#define CC (TX + KW - 1)   // 20 input cols per output tile
#define TPSTR 20           // temp tile pixel stride
#define TP_ELEMS (CIN * RR * TPSTR)         // 12800
#define PC 24              // padded pixel cols per halo row
#define NP (RR * PC)       // 240 halo pixels
#define S2 68              // InS2 pixel stride (floats), %32 == 4
#define INS2_ELEMS (NP * S2)                // 16320
#define B_STRIDE (BK + 4)
#define B_STAGE_ELEMS (BN * B_STRIDE)
#define STAGES 4
#define NTHREADS 256
#define SMEM_FLOATS (TP_ELEMS + INS2_ELEMS + STAGES * B_STAGE_ELEMS)

__device__ __forceinline__ float to_tf32(float v) {
    unsigned u;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(u) : "f"(v));
    return __uint_as_float(u);
}

__device__ __forceinline__ void cp_async16(void* smem_ptr, const void* gmem_ptr, bool pred) {
    unsigned s = (unsigned)__cvta_generic_to_shared(smem_ptr);
    int sz = pred ? 16 : 0;
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
                 :: "r"(s), "l"(gmem_ptr), "r"(sz));
}

__device__ __forceinline__ void cp_commit() {
    asm volatile("cp.async.commit_group;\n" ::);
}

template <int N>
__device__ __forceinline__ void cp_wait() {
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

__device__ __forceinline__ void ldmatrix_x4(const void* addr, unsigned& r0, unsigned& r1,
                                            unsigned& r2, unsigned& r3) {
    unsigned s = (unsigned)__cvta_generic_to_shared(addr);
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                 : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(s));
}

// Repack ConvTranspose2d weight (C_in, C_out, KH, KW) -> BgT[C_out][tap*64 + ic],
// with tap = ky*KW + kx, pre-rounded to TF32.
__global__ void repack_weight_kernel(const float* __restrict__ w,
                                     float* __restrict__ BgT,
                                     int C_out) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= KCONV) return;
    int ic = k & 63;
    int tap = k >> 6;
    const float* wp = w + (long)ic * C_out * (KH * KW) + tap;
    float* bp = BgT + k;
    for (int oc = 0; oc < C_out; ++oc) {
        *bp = to_tf32(*wp);
        wp += (KH * KW);
        bp += KCONV;
    }
}

// Direct transposed convolution: stride=1, pad=0, dilation=1, groups=1,
// C_in=64, KH=3, KW=5, C_out<=128, W_in%4==0, sm80+.
//   out[n, oc, y, x] = sum_{ic,ky,kx} w[ic,oc,ky,kx] * in[n,ic,y-ky,x-kx]
// Each block computes a TY x TX output tile (all C_out channels) of one image.
// GEMM view: M = TY*TX output pixels, N = C_out, K = tap*64 + ic.
// Input halo cached in smem as [pixel][channel] (ldmatrix-friendly, TF32
// pre-rounded); weights stream through a 4-stage cp.async pipeline;
// TF32 tensor-core mma with fp32 accumulation.
__global__ void __launch_bounds__(NTHREADS, 1)
tconv_direct_kernel(const float* __restrict__ input,
                    const float* __restrict__ BgT,
                    const float* __restrict__ bias,
                    float* __restrict__ output,
                    int N_batch, int H_in, int W_in,
                    int H_out, int W_out, int C_out,
                    int xtiles, int ytiles) {
    extern __shared__ float smem[];
    float* InT  = smem;                       // [64][200] staging
    float* InS2 = smem + TP_ELEMS;            // [240][68] tf32-rounded
    float* Bs   = smem + TP_ELEMS + INS2_ELEMS;  // 4 stages of [128][36]

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int tg   = lane & 3;
    const int g    = lane >> 2;

    const int bidx = blockIdx.x;
    const int xt = bidx % xtiles;
    const int yt = (bidx / xtiles) % ytiles;
    const int n  = bidx / (xtiles * ytiles);
    const int x0 = xt * TX;
    const int y0 = yt * TY;
    const int txl = min(TX, W_out - x0);
    const int tyl = min(TY, H_out - y0);

    const long HW  = (long)H_in * W_in;
    const long HWo = (long)H_out * W_out;
    const float* __restrict__ in_n = input + (long)n * CIN * HW;

    // ---- stage the input halo tile into InT (16B cp.async, zero-fill OOB) ----
    {
        #pragma unroll
        for (int it = 0; it < 10; ++it) {
            int idx = it * NTHREADS + tid;   // 64*10*5 = 3200 chunks
            if (idx < CIN * RR * (CC / 4)) {
                int ic = idx / (RR * (CC / 4));
                int rem = idx - ic * (RR * (CC / 4));
                int r = rem / (CC / 4);
                int c = rem - r * (CC / 4);
                int iy = y0 - (KH - 1) + r;
                int ix = x0 - (KW - 1) + 4 * c;
                bool ok = ((unsigned)iy < (unsigned)H_in) &&
                          ((unsigned)ix < (unsigned)W_in);
                const float* src = ok ? (in_n + (long)ic * HW + (long)iy * W_in + ix)
                                      : input;
                cp_async16(InT + ic * (RR * TPSTR) + r * TPSTR + 4 * c, src, ok);
            }
        }
    }
    cp_commit();

    // ---- weight (B) pipeline ----
    auto issue_B = [&](int k0, int stage) {
        float* Bsp = Bs + stage * B_STAGE_ELEMS;
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int idx = r * NTHREADS + tid;   // 1024 chunks: 128 rows x 8 x 16B
            int nl = idx >> 3;
            int c  = idx & 7;
            bool pred = nl < C_out;
            const float* src = pred ? (BgT + (long)nl * KCONV + k0 + 4 * c) : BgT;
            cp_async16(Bsp + nl * B_STRIDE + 4 * c, src, pred);
        }
        cp_commit();
    };

    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) issue_B(s * BK, s);

    cp_wait<STAGES - 2>();   // input tile + first weight stage resident
    __syncthreads();

    // ---- transpose + TF32-round: InT[ic][p] -> InS2[p'][ic] ----
    {
        #pragma unroll 1
        for (int ic4 = 0; ic4 < CIN / 4; ++ic4) {
            int p = tid;
            if (p < RR * TPSTR) {
                int r = p / TPSTR;
                int p2 = p + (PC - TPSTR) * r;   // pad each row 20 -> 24
                const float* src = InT + 4 * ic4 * (RR * TPSTR) + p;
                float v0 = to_tf32(src[0]);
                float v1 = to_tf32(src[RR * TPSTR]);
                float v2 = to_tf32(src[2 * RR * TPSTR]);
                float v3 = to_tf32(src[3 * RR * TPSTR]);
                float4* dst = (float4*)(InS2 + p2 * S2 + 4 * ic4);
                *dst = make_float4(v0, v1, v2, v3);
            }
        }
    }
    __syncthreads();

    float acc[4][4][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            #pragma unroll
            for (int r = 0; r < 4; ++r) acc[i][j][r] = 0.0f;

    const int wm  = (warp >> 2) * 64;    // warp's m offset: 0 or 64
    const int wn  = (warp & 3) * 32;     // warp's n offset: 0/32/64/96
    const int tyb = wm >> 4;             // 0 or 4
    const int mrow = lane & 15;          // ldmatrix A row (pixel)
    const int kq   = (lane >> 4) * 4;    // ldmatrix A col half (k+4?)
    const int br   = lane & 7;           // ldmatrix B row
    const int bh   = lane & 8;           // ldmatrix B k-half
    const int bq   = lane >> 4;          // ldmatrix B j half

    float afr[2][4][4];
    float bfr[2][4][2];

    // k-tile kt -> tap = kt>>1, ic0 = (kt&1)*32  (k = tap*64 + ic)
    auto load_frag = [&](int buf, int stage, int ks, int ktx) {
        const int tap = ktx >> 1;
        const int ic0 = (ktx & 1) << 5;
        const int ky = tap / KW;
        const int kx = tap - ky * KW;
        const int ry = (KH - 1) - ky;
        const int rx = (KW - 1) - kx;
        const float* Ap = InS2 + ((ry + tyb) * PC + rx + mrow) * S2 + ic0 + ks + kq;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            ldmatrix_x4(Ap + i * PC * S2,
                        afr[buf][i][0], afr[buf][i][1], afr[buf][i][2], afr[buf][i][3]);
        }
        const float* Bp = Bs + stage * B_STAGE_ELEMS +
                          (wn + bq * 8 + br) * B_STRIDE + ks + (bh ? 4 : 0);
        ldmatrix_x4(Bp, bfr[buf][0][0], bfr[buf][0][1], bfr[buf][1][0], bfr[buf][1][1]);
        ldmatrix_x4(Bp + 16 * B_STRIDE,
                    bfr[buf][2][0], bfr[buf][2][1], bfr[buf][3][0], bfr[buf][3][1]);
    };

    load_frag(0, 0, 0, 0);

    int stage = 0;
    #pragma unroll 1
    for (int kt = 0; kt < KTILES; ++kt) {
        #pragma unroll
        for (int kslice = 0; kslice < BK / 8; ++kslice) {
            const int cur = kslice & 1;
            if (kslice == BK / 8 - 1) {
                cp_wait<STAGES - 2>();
                __syncthreads();
            }
            const int ks_next = (kslice * 8 + 8) & (BK - 1);
            const int kt_next = (kslice == BK / 8 - 1)
                                    ? (kt + 1 < KTILES ? kt + 1 : KTILES - 1) : kt;
            const int stage_next = (kslice == BK / 8 - 1)
                                       ? (stage + 1 >= STAGES ? 0 : stage + 1) : stage;
            load_frag(cur ^ 1, stage_next, ks_next, kt_next);
            if (kslice == 0) {
                if (kt + STAGES - 1 < KTILES) {
                    int ws = stage > 0 ? stage - 1 : STAGES - 1;
                    issue_B((kt + STAGES - 1) * BK, ws);
                } else {
                    cp_commit();
                }
            }
            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    asm volatile(
                        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
                        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
                        : "+f"(acc[i][j][0]), "+f"(acc[i][j][1]),
                          "+f"(acc[i][j][2]), "+f"(acc[i][j][3])
                        : "r"(afr[cur][i][0]), "r"(afr[cur][i][1]),
                          "r"(afr[cur][i][2]), "r"(afr[cur][i][3]),
                          "r"(bfr[cur][j][0]), "r"(bfr[cur][j][1]));
                }
            }
        }
        stage = stage + 1 >= STAGES ? 0 : stage + 1;
    }

    // ---- epilogue: write straight from accumulators (32B-sector coalesced) ----
    float* __restrict__ out_n = output + (long)n * C_out * HWo;
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        const int ty = tyb + i;
        if (ty >= tyl) continue;
        const long rowbase = (long)(y0 + ty) * W_out + x0;
        const bool c0 = g < txl;
        const bool c1 = (g + 8) < txl;
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int oc = wn + j * 8 + 2 * tg;
            if (oc >= C_out) continue;
            const bool ok1 = (oc + 1) < C_out;
            const float b0 = bias ? bias[oc] : 0.0f;
            const long o0 = rowbase + (long)oc * HWo;
            if (c0) out_n[o0 + g] = acc[i][j][0] + b0;
            if (c1) out_n[o0 + g + 8] = acc[i][j][2] + b0;
            if (ok1) {
                const float b1 = bias ? bias[oc + 1] : 0.0f;
                if (c0) out_n[o0 + HWo + g] = acc[i][j][1] + b1;
                if (c1) out_n[o0 + HWo + g + 8] = acc[i][j][3] + b1;
            }
        }
    }
}

// Generic fallback: direct fp32 transposed convolution (any device, any shape).
__global__ void tconv_fallback_kernel(const float* __restrict__ input,
                                      const float* __restrict__ weight,
                                      const float* __restrict__ bias,
                                      float* __restrict__ output,
                                      int N_batch, int C_in, int H_in, int W_in,
                                      int H_out, int W_out, int C_out,
                                      int kh, int kw, long total) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int x = idx % W_out;
    int y = (idx / W_out) % H_out;
    int oc = (idx / ((long)W_out * H_out)) % C_out;
    int n = idx / ((long)W_out * H_out * C_out);

    float sum = bias ? bias[oc] : 0.0f;
    const int kwkh = kh * kw;
    for (int ic = 0; ic < C_in; ++ic) {
        const float* ip = input + ((long)n * C_in + ic) * H_in * W_in;
        const float* wp = weight + ((long)ic * C_out + oc) * kwkh;
        for (int ky = 0; ky < kh; ++ky) {
            int iy = y - ky;
            if ((unsigned)iy >= (unsigned)H_in) continue;
            for (int kx = 0; kx < kw; ++kx) {
                int ix = x - kx;
                if ((unsigned)ix >= (unsigned)W_in) continue;
                sum += ip[iy * W_in + ix] * wp[ky * kw + kx];
            }
        }
    }
    output[idx] = sum;
}

torch::Tensor tconv_forward(torch::Tensor input, torch::Tensor weight,
                            torch::Tensor bias, long kh, long kw) {
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "fp32 expected");
    auto x = input.contiguous();
    auto w = weight.contiguous();

    const int N_batch = 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 C_out = w.size(1);
    const int H_out = H_in + (int)kh - 1;
    const int W_out = W_in + (int)kw - 1;
    const long total = (long)N_batch * C_out * H_out * W_out;

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

    const float* bias_ptr = nullptr;
    if (bias.defined() && bias.numel() > 0) bias_ptr = bias.contiguous().data_ptr<float>();

    const cudaDeviceProp* prop = at::cuda::getCurrentDeviceProperties();

    const bool aligned =
        ((uintptr_t)x.data_ptr<float>() & 15) == 0 &&
        ((uintptr_t)out.data_ptr<float>() & 15) == 0;

    bool fast = prop->major >= 8 && C_in == CIN && kh == KH && kw == KW &&
                C_out <= BN && (W_in % 4) == 0 && aligned &&
                H_in <= 65536 && W_in <= 65536 && N_batch <= 65535 &&
                x.numel() < 2147483647L && total < 2147483647L;

    if (fast) {
        auto BgT = torch::empty({C_out, KCONV}, x.options());
        int threads = 256;
        int blocks = (KCONV + threads - 1) / threads;
        repack_weight_kernel<<<blocks, threads, 0, stream>>>(
            w.data_ptr<float>(), BgT.data_ptr<float>(), C_out);

        const int smem_bytes = SMEM_FLOATS * 4;
        static bool configured = false;
        if (!configured) {
            cudaFuncSetAttribute(tconv_direct_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 smem_bytes);
            configured = true;
        }
        const int xtiles = (W_out + TX - 1) / TX;
        const int ytiles = (H_out + TY - 1) / TY;
        long grid = (long)N_batch * ytiles * xtiles;
        tconv_direct_kernel<<<(unsigned)grid, NTHREADS, smem_bytes, stream>>>(
            x.data_ptr<float>(), BgT.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N_batch, H_in, W_in, H_out, W_out,
            C_out, xtiles, ytiles);
    } else {
        int threads = 256;
        long blocks = (total + threads - 1) / threads;
        tconv_fallback_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N_batch, C_in, H_in, W_in, H_out, W_out,
            C_out, (int)kh, (int)kw, total);
    }
    return out;
}
"""

cpp_source = "torch::Tensor tconv_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, long kh, long kw);"

tconv_ext = load_inline(
    name="tconv_direct_tf32_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["tconv_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    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.conv_transpose2d = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding, dilation=dilation,
            groups=groups, bias=bias)
        s = tuple(stride) if hasattr(stride, "__len__") else (stride, stride)
        p = tuple(padding) if hasattr(padding, "__len__") else (padding, padding)
        op = tuple(output_padding) if hasattr(output_padding, "__len__") else (output_padding, output_padding)
        d = tuple(dilation) if hasattr(dilation, "__len__") else (dilation, dilation)
        self._fast = (s == (1, 1) and p == (0, 0) and op == (0, 0)
                      and d == (1, 1) and groups == 1)
        k = kernel_size if hasattr(kernel_size, "__len__") else (kernel_size, kernel_size)
        self._kh, self._kw = int(k[0]), int(k[1])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast and x.is_cuda and x.dtype == torch.float32:
            w = self.conv_transpose2d.weight
            b = self.conv_transpose2d.bias
            if b is None:
                b = x.new_empty(0)
            return tconv_ext.tconv_forward(x, w, b, self._kh, self._kw)
        return self.conv_transpose2d(x)
```