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

// ---------------------------------------------------------------------------
// Repack the ConvTranspose2d weight (C_in, C_out, kh, kw) into
// BgT[C_out][K] with K ordered as k = ic*kh*kw + ky*kw + kx  (TF32 rounded).
// ---------------------------------------------------------------------------
__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;   // 0 .. C_in*kh*kw-1
    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) {
        float v = *wp;
        asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=f"(v) : "f"(v));
        *bp = v;
        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_batch*H_out*W_out, N = C_out, K = C_in*kh*kw
// TF32 tensor cores (mma.m16n8k8), fp32 accumulation everywhere.
// ---------------------------------------------------------------------------
__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) {
    __shared__ float As[BK][BM + 4];
    __shared__ float Bs[BN][BK + 4];
    __shared__ int s_base[BM];
    __shared__ int s_yx[BM];
    __shared__ int s_koff[BK];
    __shared__ int s_kk[BK];

    const int HW = H_in * W_in;
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int warp_m = warp >> 2;      // 0..1  (2 warps along M)
    const int warp_n = warp & 3;       // 0..3  (4 warps along N)

    const long m0 = (long)blockIdx.x * BM;

    // per-m lookup: input base offset and packed (y, x)
    if (tid < BM) {
        long m = m0 + tid;
        if (m < M) {
            int spatial = (int)(m % (long)(H_out * W_out));
            int n = (int)(m / (long)(H_out * W_out));
            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;
        } else {
            s_base[tid] = -1;
            s_yx[tid] = 0;
        }
    }
    __syncthreads();

    // per-warp accumulators: warp tile 32x32 -> 2 x 4 mma tiles, 4 regs each
    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;

    const int kwkh = kh * kw;

    for (int k0 = 0; k0 < K; k0 += BK) {
        // per-k lookup tables for this chunk
        for (int i = tid; i < BK; i += NTHREADS) {
            int k = k0 + i;
            if (k < K) {
                int ic = k / kwkh;
                int ky = (k % kwkh) / kw;
                int kx = k % kw;
                s_koff[i] = ic * HW - ky * W_in - kx;
                s_kk[i] = (ky << 16) | kx;
            } else {
                s_koff[i] = 0;
                s_kk[i] = -1;
            }
        }
        __syncthreads();

        // 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);
            float v = 0.0f;
            int kk = s_kk[kl];
            int base = s_base[ml];
            if (kk >= 0 && base >= 0) {
                int yx = s_yx[ml];
                int iy = (yx >> 16) - (kk >> 16);
                int ix = (yx & 0xffff) - (kk & 0xffff);
                if ((unsigned)iy < (unsigned)H_in && (unsigned)ix < (unsigned)W_in)
                    v = input[base + s_koff[kl]];
            }
            asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=f"(v) : "f"(v));
            As[kl][ml] = v;
        }

        // stage B tile from the repacked weight
        #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][kl] = v;
        }
        __syncthreads();

        // tensor-core compute over this K chunk
        #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[0][warp_m * 32 + i * 16];
                int cr = lane & 3, cg = lane >> 2;
                a[i][0] = Ap[(ks + cr) * (BM + 4) + cg];
                a[i][1] = Ap[(ks + cr) * (BM + 4) + cg + 8];
                a[i][2] = Ap[(ks + cr + 4) * (BM + 4) + cg];
                a[i][3] = Ap[(ks + cr + 4) * (BM + 4) + cg + 8];
            }
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                const float* Bp = &Bs[warp_n * 32 + j * 8][0];
                b[j][0] = Bp[(lane >> 2) * (BK + 4) + ks + (lane & 3)];
                b[j][1] = Bp[(lane >> 2) * (BK + 4) + ks + (lane & 3) + 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();
    }

    // write results
    const int cr = lane & 3;
    const int cg = lane >> 2;
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            int mrow0 = warp_m * 32 + i * 16 + cg;
            int ncol0 = warp_n * 32 + j * 8 + 2 * cr;
            #pragma unroll
            for (int rr = 0; rr < 2; ++rr) {
                long m = m0 + mrow0 + rr * 8;
                if (m < M) {
                    float* op = output + m * C_out;
                    float bv0 = bias ? bias[ncol0] : 0.0f;
                    float bv1 = bias ? bias[ncol0 + 1] : 0.0f;
                    if (ncol0 + 1 < C_out) {
                        op[ncol0] = acc[i][j][2 * rr] + bv0;
                        op[ncol0 + 1] = acc[i][j][2 * rr + 1] + bv1;
                    } else if (ncol0 < C_out) {
                        op[ncol0] = acc[i][j][2 * rr] + bv0;
                    }
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Reference-style fallback kernel (any GPU): fp32 direct transposed conv.
// ---------------------------------------------------------------------------
__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>();

    int dev = x.get_device();
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, dev);

    if (prop.major >= 8 && (M % BM) == 0 && K > 0) {
        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",
    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)
```