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


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

#define BM 64
#define BN 128
#define BK 32
#define NTHREADS 256
#define MAXK 1024

// TF32 rounding: destination register of cvt.rna.tf32.f32 must be .b32
__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);
}

// ---------------------------------------------------------------------------
// Repack ConvTranspose2d weight (C_in, C_out, kh, kw) into BgT[C_out][K] with
// k = ic*kh*kw + ky*kw + kx, pre-rounded to TF32.
// ---------------------------------------------------------------------------
__global__ void repack_weight_kernel(const float* __restrict__ w,
                                     float* __restrict__ BgT,
                                     int C_in, int C_out, int kh, int kw) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    int K = C_in * kh * kw;
    if (k >= K) return;
    int kwkh = kh * kw;
    int ic = k / kwkh;
    int ky = (k % kwkh) / kw;
    int kx = k % kw;
    const float* wp = w + ((long)ic * C_out) * kwkh + ky * kw + kx;
    float* bp = BgT + k;
    for (int oc = 0; oc < C_out; ++oc) {
        *bp = to_tf32(*wp);
        wp += kwkh;
        bp += K;
    }
}

// ---------------------------------------------------------------------------
// Implicit-GEMM transposed convolution (stride=1, pad=0, dilation=1, groups=1)
//   out[n, oc, y, x] = sum_{ic,ky,kx} w[ic,oc,ky,kx] * in[n,ic,y-ky,x-kx]
//   M = N*H_out*W_out, N = C_out, K = C_in*kh*kw
// TF32 tensor cores (mma.m16n8k8), fp32 accumulation, fp32 output.
// ---------------------------------------------------------------------------
__global__ __launch_bounds__(NTHREADS)
void tconv_implicit_gemm_kernel(const float* __restrict__ input,
                                const float* __restrict__ BgT,
                                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, int K, long M) {
    constexpr int SMEM_MAIN = BK * (BM + 4) + BN * (BK + 4);
    constexpr int SMEM_EPI  = BM * (BN + 4);
    constexpr int SMEM_F    = SMEM_MAIN > SMEM_EPI ? SMEM_MAIN : SMEM_EPI;
    __shared__ float smem[SMEM_F];
    __shared__ int s_koff[MAXK];
    __shared__ int s_kk[MAXK];
    __shared__ int s_base[BM];
    __shared__ int s_yx[BM];
    __shared__ long s_obase[BM];

    float* As = smem;                     // [BK][BM+4]
    float* Bs = smem + BK * (BM + 4);     // [BN][BK+4]

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int warp_m = warp >> 2;         // 0..1
    const int warp_n = warp & 3;          // 0..3
    const int cr = lane & 3;
    const int cg = lane >> 2;
    const long m0 = (long)blockIdx.x * BM;
    const int HW = H_in * W_in;
    const long HWo = (long)H_out * W_out;
    const int kwkh = kh * kw;

    // per-m lookup tables
    if (tid < BM) {
        long m = m0 + tid;
        if (m < M) {
            int spatial = (int)(m % HWo);
            int n = (int)(m / HWo);
            int y = spatial / W_out;
            int x = spatial % W_out;
            s_base[tid] = n * (C_in * HW) + y * W_in + x;
            s_yx[tid] = (y << 16) | x;
            s_obase[tid] = (long)n * C_out * HWo + spatial;
        } else {
            s_base[tid] = -1;
            s_yx[tid] = 0;
            s_obase[tid] = 0;
        }
    }
    // full-K lookup tables (div/mod done once per kernel)
    for (int i = tid; i < K; i += NTHREADS) {
        int ic = i / kwkh;
        int ky = (i % kwkh) / kw;
        int kx = i % kw;
        s_koff[i] = ic * HW - ky * W_in - kx;
        s_kk[i] = (ky << 8) | kx;
    }
    __syncthreads();

    float acc[2][4][4];
    #pragma unroll
    for (int i = 0; i < 2; ++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;

    for (int k0 = 0; k0 < K; k0 += BK) {
        // stage A tile (implicit im2col gather)
        #pragma unroll
        for (int r = 0; r < (BK * BM) / NTHREADS; ++r) {
            int idx = r * NTHREADS + tid;
            int kl = idx >> 6;            // / BM
            int ml = idx & (BM - 1);
            int k = k0 + kl;
            float v = 0.0f;
            int base = s_base[ml];
            if (k < K && base >= 0) {
                int kk = s_kk[k];
                int yx = s_yx[ml];
                int iy = (yx >> 16) - (kk >> 8);
                int ix = (yx & 0xffff) - (kk & 0xff);
                if ((unsigned)iy < (unsigned)H_in && (unsigned)ix < (unsigned)W_in)
                    v = input[base + s_koff[k]];
            }
            As[kl * (BM + 4) + ml] = to_tf32(v);
        }
        // stage B tile
        #pragma unroll
        for (int r = 0; r < (BK * BN) / NTHREADS; ++r) {
            int idx = r * NTHREADS + tid;
            int kl = idx & (BK - 1);
            int nl = idx >> 5;            // / BK
            int k = k0 + kl;
            float v = (k < K && nl < C_out) ? BgT[(long)nl * K + k] : 0.0f;
            Bs[nl * (BK + 4) + kl] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int ks = 0; ks < BK; ks += 8) {
            float a[2][4], b[4][2];
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const float* Ap = As + ks * (BM + 4) + warp_m * 32 + i * 16;
                a[i][0] = Ap[cr * (BM + 4) + cg];
                a[i][1] = Ap[cr * (BM + 4) + cg + 8];
                a[i][2] = Ap[(cr + 4) * (BM + 4) + cg];
                a[i][3] = Ap[(cr + 4) * (BM + 4) + cg + 8];
            }
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                const float* Bp = Bs + (warp_n * 32 + j * 8) * (BK + 4) + ks;
                b[j][0] = Bp[cg * (BK + 4) + cr];
                b[j][1] = Bp[cg * (BK + 4) + cr + 4];
            }
            #pragma unroll
            for (int i = 0; i < 2; ++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"(__float_as_uint(a[i][0])), "r"(__float_as_uint(a[i][1])),
                          "r"(__float_as_uint(a[i][2])), "r"(__float_as_uint(a[i][3])),
                          "r"(__float_as_uint(b[j][0])), "r"(__float_as_uint(b[j][1])));
                }
            }
        }
        __syncthreads();
    }

    // stage accumulators (+bias) into shared tile for coalesced NCHW writeout
    float* Cs = smem;                     // [BM][BN+4]
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            int nl = warp_n * 32 + j * 8 + 2 * cr;
            #pragma unroll
            for (int rr = 0; rr < 2; ++rr) {
                int ml = warp_m * 32 + i * 16 + cg + rr * 8;
                if (nl + 1 < C_out) {
                    float b0 = bias ? bias[nl] : 0.0f;
                    float b1 = bias ? bias[nl + 1] : 0.0f;
                    Cs[ml * (BN + 4) + nl]     = acc[i][j][2 * rr] + b0;
                    Cs[ml * (BN + 4) + nl + 1] = acc[i][j][2 * rr + 1] + b1;
                } else if (nl < C_out) {
                    float b0 = bias ? bias[nl] : 0.0f;
                    Cs[ml * (BN + 4) + nl] = acc[i][j][2 * rr] + b0;
                }
            }
        }
    }
    __syncthreads();

    #pragma unroll
    for (int r = 0; r < (BM * BN) / NTHREADS; ++r) {
        int idx = r * NTHREADS + tid;
        int oc = idx >> 6;                // / BM
        int ml = idx & (BM - 1);
        long m = m0 + ml;
        if (m < M && oc < C_out)
            output[s_obase[ml] + (long)oc * HWo] = Cs[ml * (BN + 4) + oc];
    }
}

// ---------------------------------------------------------------------------
// Reference-style fallback kernel (any GPU, any shape): fp32 direct tconv.
// ---------------------------------------------------------------------------
__global__ void tconv_direct_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;
    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 int K = C_in * (int)(kh * kw);
    const long M = (long)N_batch * 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();

    bool fast = prop->major >= 8 && (M % BM) == 0 && K > 0 && K <= MAXK &&
                C_in <= 65535 && kh <= 255 && kw <= 255 &&
                H_out <= 65535 && W_out <= 65535 &&
                x.numel() < 2147483647L;

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

        long grid = M / BM;
        tconv_implicit_gemm_kernel<<<(unsigned)grid, NTHREADS, 0, stream>>>(
            x.data_ptr<float>(), BgT.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, K, M);
    } else {
        long total = M * (long)C_out;
        int threads = 256;
        long blocks = (total + threads - 1) / threads;
        tconv_direct_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_implicit_gemm_v2",
    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)
        # fast path supported only for this configuration
        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)
```