```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 BM 128
#define BN 128
#define BK 32
#define STAGES 4
#define NTHREADS 256
#define MAXK 1024

#define A_STAGE_ELEMS (BK * (BM + 8))
#define B_STAGE_ELEMS (BN * (BK + 4))
#define PIPE_ELEMS (STAGES * (A_STAGE_ELEMS + 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_async4(void* smem_ptr, const void* gmem_ptr, bool pred) {
    unsigned s = (unsigned)__cvta_generic_to_shared(smem_ptr);
    int sz = pred ? 4 : 0;
    asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %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));
}

// ---------------------------------------------------------------------------
// 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]
//   GEMM: M = N*H_out*W_out, N = C_out, K = C_in*kh*kw
// TF32 tensor cores (mma.m16n8k8), 4-stage cp.async pipeline, fp32 accum.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NTHREADS, 1)
tconv_igemm_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) {
    extern __shared__ float smem[];
    float* As = smem;                                  // 4 stages of [BK][BM+8]
    float* Bs = smem + STAGES * A_STAGE_ELEMS;         // 4 stages of [BN][BK+4]
    float* Cs = smem;                                  // epilogue reuse [BN][BM+8]
    int* s_koff = (int*)(smem + PIPE_ELEMS);
    int* s_ky   = s_koff + MAXK;
    int* s_kx   = s_ky + MAXK;
    int* s_base = s_kx + MAXK;
    int* s_y    = s_base + BM;
    int* s_x    = s_y + BM;
    long* s_ob  = (long*)(s_x + BM);

    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 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_y[tid] = y;
            s_x[tid] = x;
            s_ob[tid] = (long)n * C_out * HWo + spatial;
        } else {
            s_base[tid] = -1;
            s_y[tid] = 0;
            s_x[tid] = 0;
            s_ob[tid] = 0;
        }
    }
    // per-k lookup tables
    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_ky[i] = ky;
        s_kx[i] = kx;
    }
    __syncthreads();

    // A load indices: thread covers fixed ml, k = k0 + c + 2*r
    const int a_ml = tid & (BM - 1);
    const int a_c  = tid >> 7;              // 0..1
    const int a_base = s_base[a_ml];
    const int a_y = s_y[a_ml];
    const int a_x = s_x[a_ml];
    // B load indices: thread covers fixed kl, n = d + 8*r
    const int b_kl = tid & (BK - 1);
    const int b_d  = tid >> 5;              // 0..7

    auto issue_tile = [&](int k0, int stage) {
        float* Asp = As + stage * A_STAGE_ELEMS;
        float* Bsp = Bs + stage * B_STAGE_ELEMS;
        #pragma unroll
        for (int r = 0; r < 16; ++r) {
            int kl = a_c + 2 * r;
            int k = k0 + kl;
            bool pred = false;
            const float* src = input;
            if (k < K && a_base >= 0) {
                int iy = a_y - s_ky[k];
                int ix = a_x - s_kx[k];
                if ((unsigned)iy < (unsigned)H_in && (unsigned)ix < (unsigned)W_in) {
                    pred = true;
                    src = input + (long)a_base + s_koff[k];
                }
            }
            cp_async4(Asp + kl * (BM + 8) + a_ml, src, pred);
        }
        #pragma unroll
        for (int r = 0; r < 16; ++r) {
            int nl = b_d + 8 * r;
            int k = k0 + b_kl;
            bool pred = (k < K) && (nl < C_out);
            const float* src = BgT + (long)nl * K + k;
            cp_async4(Bsp + nl * (BK + 4) + b_kl, src, pred);
        }
        cp_commit();
    };

    // prologue: fill first STAGES-1 stages
    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        int k0 = s * BK;
        if (k0 < K) issue_tile(k0, s);
        else cp_commit();
    }
    cp_wait<STAGES - 2>();
    __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;

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

    const int wm = (warp >> 2) * 64;
    const int wn = (warp & 3) * 32;

    auto load_frag = [&](int buf, int stage, int ks) {
        const float* Ap = As + stage * A_STAGE_ELEMS;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const float* p = Ap + ks * (BM + 8) + wm + i * 16;
            afr[buf][i][0] = to_tf32(p[tg * (BM + 8) + g]);
            afr[buf][i][1] = to_tf32(p[tg * (BM + 8) + g + 8]);
            afr[buf][i][2] = to_tf32(p[(tg + 4) * (BM + 8) + g]);
            afr[buf][i][3] = to_tf32(p[(tg + 4) * (BM + 8) + g + 8]);
        }
        const float* Bp = Bs + stage * B_STAGE_ELEMS;
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const float* p = Bp + (wn + j * 8) * (BK + 4) + ks;
            bfr[buf][j][0] = p[g * (BK + 4) + tg];
            bfr[buf][j][1] = p[g * (BK + 4) + tg + 4];
        }
    };

    load_frag(0, 0, 0);

    const int ktiles = (K + BK - 1) / BK;
    int stage = 0;
    int k0 = 0;
    for (int kt = 0; kt < ktiles; ++kt, k0 += BK) {
        #pragma unroll
        for (int kslice = 0; kslice < BK / 8; ++kslice) {
            int ks = kslice * 8;
            int cur = kslice & 1;
            if (kslice == BK / 8 - 1) {
                cp_wait<STAGES - 2>();
                __syncthreads();
            }
            int ks_next = (ks + 8) % BK;
            int stage_next = (kslice == BK / 8 - 1) ? (stage + 1 >= STAGES ? 0 : stage + 1) : stage;
            load_frag(cur ^ 1, stage_next, ks_next);
            if (kslice == 0) {
                int k_load = k0 + (STAGES - 1) * BK;
                int write_stage = stage > 0 ? stage - 1 : STAGES - 1;
                if (k_load < K) issue_tile(k_load, write_stage);
                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"(__float_as_uint(afr[cur][i][0])), "r"(__float_as_uint(afr[cur][i][1])),
                          "r"(__float_as_uint(afr[cur][i][2])), "r"(__float_as_uint(afr[cur][i][3])),
                          "r"(__float_as_uint(bfr[cur][j][0])), "r"(__float_as_uint(bfr[cur][j][1])));
                }
            }
        }
        stage = stage + 1 >= STAGES ? 0 : stage + 1;
    }

    // epilogue: stage C tile (+bias) in shared as [oc][m] then coalesced write
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            int nl = wn + j * 8 + 2 * tg;
            #pragma unroll
            for (int rr = 0; rr < 2; ++rr) {
                int ml = wm + i * 16 + g + rr * 8;
                if (nl + 1 < C_out) {
                    float b0 = bias ? bias[nl] : 0.0f;
                    float b1 = bias ? bias[nl + 1] : 0.0f;
                    Cs[nl * (BM + 8) + ml]       = acc[i][j][2 * rr] + b0;
                    Cs[(nl + 1) * (BM + 8) + ml] = acc[i][j][2 * rr + 1] + b1;
                } else if (nl < C_out) {
                    float b0 = bias ? bias[nl] : 0.0f;
                    Cs[nl * (BM + 8) + ml] = 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 >> 7;
        int ml = idx & (BM - 1);
        if (oc < C_out && m0 + ml < M)
            output[s_ob[ml] + (long)oc * HWo] = Cs[oc * (BM + 8) + ml];
    }
}

// ---------------------------------------------------------------------------
// 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 && K > 0 && K <= MAXK &&
                H_in <= 32768 && W_in <= 32768 &&
                x.numel() < 2147483647L && M < 2147483647L &&
                (long)M * C_out < 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);

        const int smem_bytes = PIPE_ELEMS * 4 + 3 * MAXK * 4 + 3 * BM * 4 + BM * 8;
        static bool configured = false;
        if (!configured) {
            cudaFuncSetAttribute(tconv_igemm_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 smem_bytes);
            configured = true;
        }
        long grid = (M + BM - 1) / BM;
        tconv_igemm_kernel<<<(unsigned)grid, NTHREADS, smem_bytes, 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_igemm_pipe_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)
```