import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 2D convolution operation with asymmetric input and kernel size.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Tuple of integers representing the kernel size (height, width).
        stride (tuple, optional): Tuple of integers representing the stride of the convolution. Defaults to (1, 1).
        padding (tuple, optional): Tuple of integers representing the padding applied to the input. Defaults to (0, 0).
        output_padding (tuple, optional): Tuple of integers representing the additional size added to one side of the output shape. Defaults to (0, 0).
        dilation (tuple, optional): Tuple of integers representing the spacing between kernel elements. Defaults to (1, 1).
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    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(Model, 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)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height_in, width_in).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv_transpose2d(x)

# Test code
batch_size = 64
in_channels = 64
out_channels = 128
kernel_size = (3, 5)
height_in = 128
width_in = 256

def get_inputs():
    x = torch.rand(batch_size, in_channels, height_in, width_in)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor convt_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 bool has_bias, int64_t stride_h, int64_t stride_w,
                                 int64_t pad_h, int64_t pad_w, int64_t outpad_h, int64_t outpad_w,
                                 int64_t dil_h, int64_t dil_w, int64_t groups);

torch::Tensor convt_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            bool has_bias, int64_t stride_h, int64_t stride_w,
                            int64_t pad_h, int64_t pad_w, int64_t outpad_h, int64_t outpad_w,
                            int64_t dil_h, int64_t dil_w, int64_t groups) {
    return convt_forward_cuda(x, weight, bias, has_bias, stride_h, stride_w, pad_h, pad_w,
                              outpad_h, outpad_w, dil_h, dil_w, groups);
}
"""

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

#define BM 16
#define BN 64
#define BK 32

// Specialized implicit-GEMM kernel for the benchmark shape:
// N=64, Cin=64, HxW=128x256, Cout=128, K=3x5, stride=1, padding=0, dilation=1.
// Computes a 16 x 64 tile of output per block; each thread accumulates 4 columns.
__global__ __launch_bounds__(256, 2)
void convt2d_64_128_3x5_kernel(const float* __restrict__ x,
                                const float* __restrict__ w,
                                const float* __restrict__ bias,
                                float* __restrict__ y,
                                int has_bias) {
    __shared__ float As[BM][BK + 1];
    __shared__ float Bs[BK][BN + 1];

    const int tx = threadIdx.x;       // 0..15
    const int ty = threadIdx.y;       // 0..15
    const int tid = ty * 16 + tx;

    const int n = blockIdx.z;
    const int oc = blockIdx.y * BM + ty;
    const int col_base = blockIdx.x * BN;
    const int col0 = col_base + tx;
    const int col1 = col0 + 16;
    const int col2 = col0 + 32;
    const int col3 = col0 + 48;

    float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f;

    #pragma unroll
    for (int k0 = 0; k0 < 960; k0 += BK) {
        // load A tile: weights laid out [IC, OC, KH, KW]
        for (int t = tid; t < BM * BK; t += 256) {
            const int r = t / BK;
            const int kk = t - r * BK;
            const int k = k0 + kk;
            const int ic = k / 15;
            const int rem = k - ic * 15;
            const int kh = rem / 5;
            const int kw = rem - kh * 5;
            const int ocr = blockIdx.y * BM + r;
            As[r][kk] = w[((ic * 128 + ocr) * 3 + kh) * 5 + kw];
        }

        // load B tile: implicit full-convolution input matrix
        for (int t = tid; t < BK * BN; t += 256) {
            const int kk = t / BN;
            const int c = t - kk * BN;
            const int k = k0 + kk;
            const int ic = k / 15;
            const int rem = k - ic * 15;
            const int kh = rem / 5;
            const int kw = rem - kh * 5;
            const int col = col_base + c;
            float v = 0.0f;
            if (col < 33800) {
                const int oh = col / 260;
                const int ow = col - oh * 260;
                const int ih = oh - kh;
                const int iw = ow - kw;
                if ((unsigned)ih < 128u && (unsigned)iw < 256u) {
                    v = x[((n * 64 + ic) * 128 + ih) * 256 + iw];
                }
            }
            Bs[kk][c] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            const float a = As[ty][kk];
            acc0 = fmaf(a, Bs[kk][tx],      acc0);
            acc1 = fmaf(a, Bs[kk][tx + 16], acc1);
            acc2 = fmaf(a, Bs[kk][tx + 32], acc2);
            acc3 = fmaf(a, Bs[kk][tx + 48], acc3);
        }
        __syncthreads();
    }

    if (oc < 128) {
        const float b = has_bias ? bias[oc] : 0.0f;
        if (col0 < 33800) {
            const int oh = col0 / 260; const int ow = col0 - oh * 260;
            y[((n * 128 + oc) * 130 + oh) * 260 + ow] = acc0 + b;
        }
        if (col1 < 33800) {
            const int oh = col1 / 260; const int ow = col1 - oh * 260;
            y[((n * 128 + oc) * 130 + oh) * 260 + ow] = acc1 + b;
        }
        if (col2 < 33800) {
            const int oh = col2 / 260; const int ow = col2 - oh * 260;
            y[((n * 128 + oc) * 130 + oh) * 260 + ow] = acc2 + b;
        }
        if (col3 < 33800) {
            const int oh = col3 / 260; const int ow = col3 - oh * 260;
            y[((n * 128 + oc) * 130 + oh) * 260 + ow] = acc3 + b;
        }
    }
}

// Simple correctness fallback for unusual constructor arguments (not optimized).
__global__ void convt2d_generic_kernel(const float* __restrict__ x,
                                       const float* __restrict__ w,
                                       const float* __restrict__ bias,
                                       float* __restrict__ y,
                                       int total, int B, int IC, int H, int Wd, int OC, int KH, int KW,
                                       int OH, int OW, int sh, int sw, int ph, int pw, int dh, int dw,
                                       int groups, int has_bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    int t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int oc = t % OC;
    int n = t / OC;
    int OCpg = OC / groups;
    int ICpg = IC / groups;
    int g = oc / OCpg;
    int ocg = oc - g * OCpg;
    float acc = has_bias ? bias[oc] : 0.0f;
    for (int icg = 0; icg < ICpg; ++icg) {
        int ic = g * ICpg + icg;
        for (int kh = 0; kh < KH; ++kh) {
            int th = oh + ph - kh * dh;
            if (th % sh != 0) continue;
            int ih = th / sh;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < KW; ++kw) {
                int tw = ow + pw - kw * dw;
                if (tw % sw != 0) continue;
                int iw = tw / sw;
                if ((unsigned)iw < (unsigned)Wd) {
                    float xv = x[((n * IC + ic) * H + ih) * Wd + iw];
                    float wv = w[((ic * OCpg + ocg) * KH + kh) * KW + kw];
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor convt_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 bool has_bias, int64_t stride_h, int64_t stride_w,
                                 int64_t pad_h, int64_t pad_w, int64_t outpad_h, int64_t outpad_w,
                                 int64_t dil_h, int64_t dil_w, int64_t groups) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat, "float32 only");
    const int B = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int H = (int)x.size(2);
    const int Wd = (int)x.size(3);
    const int KH = (int)weight.size(2);
    const int KW = (int)weight.size(3);
    const int OCpg = (int)weight.size(1);
    const int OC = OCpg * (int)groups;
    const int OH = (H - 1) * (int)stride_h - 2 * (int)pad_h + (int)dil_h * (KH - 1) + (int)outpad_h + 1;
    const int OW = (Wd - 1) * (int)stride_w - 2 * (int)pad_w + (int)dil_w * (KW - 1) + (int)outpad_w + 1;
    auto y = torch::empty({B, OC, OH, OW}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;

    if (B == 64 && IC == 64 && H == 128 && Wd == 256 && OC == 128 && KH == 3 && KW == 5 &&
        stride_h == 1 && stride_w == 1 && pad_h == 0 && pad_w == 0 && outpad_h == 0 && outpad_w == 0 &&
        dil_h == 1 && dil_w == 1 && groups == 1 && x.is_contiguous() && weight.is_contiguous()) {
        dim3 block(16, 16, 1);
        dim3 grid((33800 + BN - 1) / BN, 8, 64);
        convt2d_64_128_3x5_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), has_bias ? 1 : 0);
    } else {
        int total = B * OC * OH * OW;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        convt2d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), total, B, IC, H, Wd, OC, KH, KW,
            OH, OW, (int)stride_h, (int)stride_w, (int)pad_h, (int)pad_w, (int)dil_h, (int)dil_w,
            (int)groups, has_bias ? 1 : 0);
    }
    return y;
}
"""

_ext = load_inline(
    name="convt2d_asym_implicit_gemm_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["convt_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

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)
        self.stride = tuple(stride)
        self.padding = tuple(padding)
        self.output_padding = tuple(output_padding)
        self.dilation = tuple(dilation)
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv_transpose2d.weight
        b = self.conv_transpose2d.bias
        has_bias = b is not None
        if not has_bias:
            b = w.new_empty((0,))
        # Benchmark inputs are contiguous; make the extension robust for occasional non-contiguous input.
        if not x.is_contiguous():
            x = x.contiguous()
        return _ext.convt_forward(x, w, b, has_bias,
                                  int(self.stride[0]), int(self.stride[1]),
                                  int(self.padding[0]), int(self.padding[1]),
                                  int(self.output_padding[0]), int(self.output_padding[1]),
                                  int(self.dilation[0]), int(self.dilation[1]), int(self.groups))