import torch
import torch.nn as nn

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

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        output_padding (int, optional): Additional size added to one side of the output shape. Defaults to 0.
        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: int, stride: int = 1, padding: int = 0, output_padding: int = 0, 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, 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 = 8
in_channels = 32
out_channels = 32
kernel_size = 3
# large asymmetric input
height_in = 512
width_in = 1024

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>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

torch::Tensor deconv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride, int64_t padding, int64_t output_padding, int64_t groups);
'''

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

// Specialized path for the benchmark case: N x 32 x H x W, 32 output channels,
// 3x3 kernel, stride=1, padding=0, groups=1.  One warp computes the complete
// 32-channel output vector for one spatial location.  Input scalar loads are
// broadcast within the warp, while weight loads are naturally coalesced over OC.
__global__ __launch_bounds__(256, 2)
void deconv_k3s1p0_c32_warp_kernel(const float* __restrict__ x,
                                    const float* __restrict__ w,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    int N, int H, int W, int OH, int OW,
                                    bool has_bias) {
    constexpr int WARPS_PER_BLOCK = 8;
    int warp_in_block = threadIdx.x >> 5;
    int lane = threadIdx.x & 31;      // output channel
    long long warp_id = (long long)blockIdx.x * WARPS_PER_BLOCK + warp_in_block;
    long long npos = (long long)N * OH * OW;
    if (warp_id >= npos) return;

    int ow = (int)(warp_id % OW);
    long long t = warp_id / OW;
    int oh = (int)(t % OH);
    int n = (int)(t / OH);

    float acc = has_bias ? b[lane] : 0.0f;
    unsigned mask = 0xffffffffu;

    // Fast, branch-free interior.  For k=3,p=0, all 9 input taps are valid iff
    // oh in [2,H-1] and ow in [2,W-1].  This covers almost all output pixels.
    if (oh >= 2 && oh < H && ow >= 2 && ow < W) {
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            int ih = oh - kh;
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                int iw = ow - kw;
                #pragma unroll
                for (int ic = 0; ic < 32; ++ic) {
                    float xv = 0.0f;
                    if (lane == 0) xv = x[((n * 32 + ic) * H + ih) * W + iw];
                    xv = __shfl_sync(mask, xv, 0);
                    float wt = w[((ic * 32 + lane) * 3 + kh) * 3 + kw];
                    acc = fmaf(xv, wt, acc);
                }
            }
        }
    } else {
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            int ih = oh - kh;
            if ((unsigned)ih < (unsigned)H) {
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    int iw = ow - kw;
                    if ((unsigned)iw < (unsigned)W) {
                        #pragma unroll
                        for (int ic = 0; ic < 32; ++ic) {
                            float xv = 0.0f;
                            if (lane == 0) xv = x[((n * 32 + ic) * H + ih) * W + iw];
                            xv = __shfl_sync(mask, xv, 0);
                            float wt = w[((ic * 32 + lane) * 3 + kh) * 3 + kw];
                            acc = fmaf(xv, wt, acc);
                        }
                    }
                }
            }
        }
    }

    y[((n * 32 + lane) * OH + oh) * OW + ow] = acc;
}

// General correctness kernel for less common constructor arguments.
__global__
void deconv_generic_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           long long total,
                           int N, int inC, int H, int W,
                           int outC, int OH, int OW,
                           int kH, int kW,
                           int stride, int padding, int groups,
                           bool has_bias) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = (int)(idx % OW);
    long long q = idx / OW;
    int oh = (int)(q % OH);
    q /= OH;
    int oc = (int)(q % outC);
    int n = (int)(q / outC);

    int out_per_g = outC / groups;
    int in_per_g = inC / groups;
    int g = oc / out_per_g;
    int ocg = oc - g * out_per_g;
    int ic0 = g * in_per_g;
    int ic1 = ic0 + in_per_g;

    float acc = has_bias ? b[oc] : 0.0f;
    for (int ic = ic0; ic < ic1; ++ic) {
        for (int kh = 0; kh < kH; ++kh) {
            int th = oh + padding - kh;
            if (th % stride != 0) continue;
            int ih = th / stride;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < kW; ++kw) {
                int tw = ow + padding - kw;
                if (tw % stride != 0) continue;
                int iw = tw / stride;
                if ((unsigned)iw >= (unsigned)W) continue;
                float xv = x[((n * inC + ic) * H + ih) * W + iw];
                float wt = w[(((ic * out_per_g + ocg) * kH + kh) * kW + kw)];
                acc = fmaf(xv, wt, acc);
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor deconv_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride64, int64_t padding64, int64_t output_padding64, int64_t groups64) {
    CHECK_CUDA(x); CHECK_CUDA(weight);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight);
    CHECK_FLOAT(x); CHECK_FLOAT(weight);
    if (bias.numel() > 0) { CHECK_CUDA(bias); CHECK_CONTIGUOUS(bias); CHECK_FLOAT(bias); }
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));

    int N = (int)x.size(0);
    int inC = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int kH = (int)weight.size(2);
    int kW = (int)weight.size(3);
    int stride = (int)stride64;
    int padding = (int)padding64;
    int output_padding = (int)output_padding64;
    int groups = (int)groups64;
    int outC = (int)weight.size(1) * groups;
    int OH = (H - 1) * stride - 2 * padding + kH + output_padding;
    int OW = (W - 1) * stride - 2 * padding + kW + output_padding;

    auto y = torch::empty({N, outC, OH, OW}, x.options());
    bool has_bias = bias.numel() > 0;
    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (inC == 32 && outC == 32 && kH == 3 && kW == 3 && stride == 1 && padding == 0 && output_padding == 0 && groups == 1) {
        constexpr int threads = 256;
        constexpr int warps_per_block = 8;
        long long nwarps = (long long)N * OH * OW;
        int blocks = (int)((nwarps + warps_per_block - 1) / warps_per_block);
        deconv_k3s1p0_c32_warp_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), N, H, W, OH, OW, has_bias);
    } else {
        const int threads = 256;
        long long total = (long long)N * outC * OH * OW;
        int blocks = (int)((total + threads - 1) / threads);
        deconv_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(), total,
            N, inC, H, W, outC, OH, OW, kH, kW, stride, padding, groups, has_bias);
    }
    return y;
}
'''

_ext = load_inline(
    name='conv_transpose2d_32_k3_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['deconv_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: int, stride: int = 1, padding: int = 0, output_padding: int = 0, 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, groups=groups, bias=bias)
        self.stride = stride if isinstance(stride, int) else stride[0]
        self.padding = padding if isinstance(padding, int) else padding[0]
        self.output_padding = output_padding if isinstance(output_padding, int) else output_padding[0]
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # The benchmark inputs and module parameters are contiguous.  Keeping this
        # explicit also preserves correctness for non-contiguous callers.
        xc = x.contiguous()
        w = self.conv_transpose2d.weight.contiguous()
        b = self.conv_transpose2d.bias
        if b is None:
            b = torch.empty(0, device=xc.device, dtype=xc.dtype)
        else:
            b = b.contiguous()
        return _ext.deconv_forward(xc, w, b, int(self.stride), int(self.padding), int(self.output_padding), int(self.groups))